diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 000000000..15f7bdd79 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,11 @@ +version: 2 +updates: + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + cooldown: + default-days: 7 + groups: + actions: + patterns: ["*"] diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 000000000..724766281 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,63 @@ +name: CI + +on: + pull_request: + push: + branches: [main] + +permissions: + contents: read + +concurrency: + group: ci-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + ruff: + name: Ruff + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + cache-dependency-glob: uv.lock + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Install dependencies + run: uv sync --locked --extra dev + + - name: Run Ruff + run: uv run ruff check . + + - name: Check formatting + run: uv run ruff format --check . + + tests: + name: Tests + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + cache-dependency-glob: uv.lock + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Install dependencies + run: uv sync --locked --extra dev + + - name: Run tests + run: uv run pytest diff --git a/.github/workflows/claude-review.yml b/.github/workflows/claude-review.yml new file mode 100644 index 000000000..ecda2212a --- /dev/null +++ b/.github/workflows/claude-review.yml @@ -0,0 +1,78 @@ +name: Claude PR Review + +on: + pull_request_target: + types: [opened, synchronize, ready_for_review, reopened] + +permissions: + contents: read + pull-requests: write + issues: read + id-token: write + +concurrency: + group: claude-review-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + review: + if: github.event.pull_request.draft == false + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + # On pull_request_target, keep checkout on the trusted base-repo ref. + # The Claude action can review the PR via GitHub context/API without + # executing untrusted fork code with repository secrets. + persist-credentials: false + + - name: Compose review prompt + id: compose + run: | + { + printf 'prompt<> "$GITHUB_OUTPUT" + + - name: Prepare Claude Code bin directory + run: mkdir -p "$HOME/.local/bin" + + - uses: anthropics/claude-code-action@v1.0.137 + with: + anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }} + # Bypass the OIDC -> Claude GitHub App token exchange. That exchange + # rejects OIDC tokens minted for pull_request_target events with + # "401 Invalid OIDC token", which broke every review after the switch + # away from pull_request. Using the workflow's GITHUB_TOKEN works for + # both same-repo and fork PRs; comments post as github-actions[bot] + # instead of claude[bot], which is the documented trade-off. + github_token: ${{ secrets.GITHUB_TOKEN }} + track_progress: true + prompt: ${{ steps.compose.outputs.prompt }} diff --git a/.github/workflows/claude.yml b/.github/workflows/claude.yml new file mode 100644 index 000000000..74ea6068b --- /dev/null +++ b/.github/workflows/claude.yml @@ -0,0 +1,35 @@ +name: Claude on Mention + +on: + issue_comment: + types: [created] + pull_request_review_comment: + types: [created] + pull_request_review: + types: [submitted] + issues: + types: [opened, assigned] + +permissions: + contents: write + pull-requests: write + issues: write + id-token: write + +jobs: + claude: + if: | + (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) || + (github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) || + (github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) || + (github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude'))) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - uses: anthropics/claude-code-action@v1.0.137 + with: + anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }} + track_progress: true diff --git a/.gitignore b/.gitignore index d758b077f..c10ab3552 100644 --- a/.gitignore +++ b/.gitignore @@ -56,6 +56,7 @@ frontend/yarn-error.log* eval/ # Project-specific +scratch/ session_logs/ /logs hf-agent-leaderboard/ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..e7e76e879 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,57 @@ +# Agent Notes + +## Local Dev Servers + +- Frontend: from `frontend/`, run `npm ci` if dependencies are missing, then `npm run dev`. +- Backend: from `backend/`, run `uv run uvicorn main:app --host ::1 --port 7860`. +- Frontend URL: http://localhost:5173/ +- Backend health check: `curl -g http://[::1]:7860/api` +- Frontend proxy health check: `curl http://localhost:5173/api` + +Notes: + +- Vite proxies `/api` and `/auth` to `http://localhost:7860`. +- If `127.0.0.1:7860` is already owned by another local process, binding the backend to `::1` lets the Vite proxy resolve `localhost` cleanly. +- Prefer `npm ci` over `npm install` for setup, since `npm install` may rewrite `frontend/package-lock.json` metadata depending on npm version. +- Non-local LLM calls use `https://router.huggingface.co/v1` with the active Hugging Face user's token. Web sessions and the CLI default to GLM 5.2. For local development, set `HF_TOKEN` and optionally `ML_INTERN_DEFAULT_MODEL_ID`. +- When asked to start the local server, export the GitHub CLI token first with `export GITHUB_TOKEN="$(gh auth token)"`. +- When debugging a web app issue tied to a session ID, inspect the session data in `smolagents/ml-intern-sessions` for additional context. + +## Development Checks + +- Before every commit, run `uv run ruff check .` and `uv run ruff format --check .`. +- If formatting fails, run `uv run ruff format .`, then re-run the Ruff checks before committing. + +## Git Workflow + +- Before creating any new branch or worktree, switch to `main` and pull the latest changes. + +## GitHub CLI + +- Always use the `gh` CLI for GitHub operations such as opening, editing, inspecting, or commenting on PRs and issues. +- For multiline PR descriptions, prefer `gh pr edit --body-file ` over inline `--body` so shell quoting, `$` env-var names, backticks, and newlines are preserved correctly. +- If `gh` reports an invalid token or auth failure, retry the command with `GH_TOKEN` and `GITHUB_TOKEN` unset, for example `env -u GH_TOKEN -u GITHUB_TOKEN gh pr create ...`, so `gh` can use the stored login token instead of a stale environment token. +- In Codex, sandboxed `gh` auth checks can report a valid keyring login as invalid when GitHub network access is restricted. Before telling the user to re-authenticate, retry with both env tokens unset and GitHub network access enabled. + +## GitHub PRs + +- Open code changes as GitHub PRs first. Do not push code changes directly to the Hugging Face Space deployment branch or Space remote before the PR has been opened, reviewed, and merged, unless the user explicitly asks to bypass the PR flow. +- After implementing a plan, run the required checks, commit the changes, open a GitHub PR, then start the backend and frontend local dev servers for testing. + +## Hugging Face Space Deploys + +- The Space remote is `space` and points to `https://huggingface.co/spaces/smolagents/ml-intern`. +- Deploy GitHub `main` to the Space from the local `space-main` branch by merging `origin/main` into `space-main` with a single merge commit, then pushing `space-main:main` to the `space` remote. +- Keep the Space-only README frontmatter on `space-main`; `.gitattributes` should contain `README.md merge=ours` and the local repo config should include `merge.ours.driver=true`. +- Local dev commonly uses a personal `HF_TOKEN`, but the deployed Space uses HF OAuth tokens. When adding Hub features, make sure the Space README `hf_oauth_scopes` frontmatter and the backend OAuth request in `backend/routes/auth.py` include the scopes required by the Hub APIs being called. A feature can work locally with a broad PAT and still fail in production with 403s if OAuth scopes are missing; after changing scopes, users may need to log out and log in again to receive a fresh token. +- Recommended deploy flow: + +```bash +git pull --ff-only origin main +git switch space-main +git config merge.ours.driver true +git merge --no-ff origin/main -m "Deploy $(date +%Y-%m-%d)" \ + -m "Co-authored-by: OpenAI Codex " +git push space space-main:main +git switch main +``` diff --git a/LICENSE b/LICENSE new file mode 100644 index 000000000..261eeb9e9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md index 29fe439b8..e0bf7384a 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,14 @@ smolagents logo

+

+ License + Website +

+ # ML Intern -An ML intern that autonomously researches, writes, and ships good quality ML releated code using the Hugging Face ecosystem — with deep access to docs, papers, datasets, and cloud compute. +An ML intern that autonomously researches, writes, and ships good quality ML related code using the Hugging Face ecosystem — with deep access to docs, papers, datasets, and cloud compute. ## Quick Start @@ -26,21 +31,21 @@ ml-intern Create a `.env` file in the project root (or export these in your shell): ```bash -ANTHROPIC_API_KEY= # if using anthropic models -HF_TOKEN= -GITHUB_TOKEN= +HF_TOKEN= # HF Router inference + Hub actions +GITHUB_TOKEN= ``` -If no `HF_TOKEN` is set, the CLI will prompt you to paste one on first launch. To get a GITHUB_TOKEN follow the tutorial [here](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens#creating-a-fine-grained-personal-access-token). + +All API-based model calls go through Hugging Face [Inference Providers](https://huggingface.co/docs/inference-providers/en/index), so your `HF_TOKEN` must be allowed to make Inference Provider calls. If no `HF_TOKEN` is set, the CLI will prompt you to paste one on first launch unless you start on a local model. To get a `GITHUB_TOKEN` follow the tutorial [here](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens#creating-a-fine-grained-personal-access-token). See the [local models section below](#local-models) for instructions on using agents that run on your hardware. ### Usage -**Interactive mode** (start a chat session): +#### Interactive mode (start a chat session): ```bash ml-intern ``` -**Headless mode** (single prompt, auto-approve): +#### Headless mode (single prompt, auto-approve): ```bash ml-intern "fine-tune llama on my dataset" @@ -49,9 +54,160 @@ ml-intern "fine-tune llama on my dataset" **Options:** ```bash -ml-intern --model anthropic/claude-opus-4-6 "your prompt" +ml-intern --sandbox-tools "your prompt" # use HF Space sandbox tools ml-intern --max-iterations 100 "your prompt" ml-intern --no-stream "your prompt" +# Change model +ml-intern --model moonshotai/Kimi-K2.7-Code:novita "your prompt" +ml-intern --model openai/gpt-5.5:fal-ai "your prompt" +``` + +Run `ml-intern` then `/model` to see the full list of suggested model ids +(Claude, GPT, HF Router models like MiniMax, Kimi, GLM, DeepSeek, and local +model prefixes). + +Hosted inference is billed to the active Hugging Face user. See below on how to run `ml-intern` with local models. + +#### Local models + +Local model support uses OpenAI-compatible HTTP endpoints through LiteLLM. The +agent does not load model weights directly from disk; start your inference +server first, then select it with a provider-specific model prefix: + +```bash +ml-intern --model ollama/llama3.1:8b "your prompt" +ml-intern --model vllm/meta-llama/Llama-3.1-8B-Instruct "your prompt" +``` + +Inside interactive mode, switch with `/model`: + +```text +/model ollama/llama3.1:8b +/model lm_studio/google/gemma-3-4b +/model llamacpp/llama-3.1-8b-instruct +``` + +Supported local prefixes are `ollama/`, `vllm/`, `lm_studio/`, and +`llamacpp/`. + +```bash +LOCAL_LLM_BASE_URL=http://localhost:8000 +LOCAL_LLM_API_KEY= +``` + +Set `LOCAL_LLM_BASE_URL` and optional `LOCAL_LLM_API_KEY` to use one shared +local endpoint, or override a specific provider with its matching `*_BASE_URL` +/ `*_API_KEY` variable, such as `OLLAMA_BASE_URL` or `VLLM_API_KEY`. +Provider-specific variables take precedence over the shared local variables. +Base URLs may include or omit `/v1`. + +**CLI tool runtime:** + +By default, the CLI runs `bash`, `read`, `write`, and `edit` on your local +filesystem. To use HF Space sandbox tools instead, including `sandbox_create`, +opt in with `--sandbox-tools`: + +```bash +ml-intern --sandbox-tools "test this training script in a GPU sandbox" +ml-intern --model llamacpp/ggml-org/gemma-3-1b-it-GGUF --sandbox-tools +``` + +Sandbox tool runtime requires `HF_TOKEN`, even when the selected model is local, +because it creates private HF Spaces. You can also make sandbox tools your CLI +default in `~/.config/ml-intern/cli_agent_config.json`: + +```json +{ "tool_runtime": "sandbox" } +``` + +Use the default local runtime when you want tools to inspect or edit files in +your checkout. Use sandbox runtime when you want the agent to create or replace +an HF Space sandbox, test code remotely, or request GPU sandbox hardware before +launching larger HF Jobs. + +## Sharing Traces + +Every session is auto-uploaded to your **own private Hugging Face dataset** +in [Claude Code JSONL format](https://huggingface.co/changelog/agent-trace-viewer), +which the HF Agent Trace Viewer auto-detects so you can browse turns, tool +calls, and model responses directly on the Hub. + +By default the dataset is named `{your-hf-username}/ml-intern-sessions` and is +**created private**. You can flip it to public from inside the CLI: + +```bash +/share-traces # show current visibility + dataset URL +/share-traces public # publish (anyone can view) +/share-traces private # lock it back down +``` + +You can also flip visibility from the dataset page on huggingface.co — the +agent honours whatever you set there for subsequent uploads. + +To opt out entirely, set in your CLI config (e.g. `configs/cli_agent_config.json` +or `~/.config/ml-intern/cli_agent_config.json`): + +```json +{ "share_traces": false } +``` + +To override the destination repo, set: + +```json +{ "personal_trace_repo_template": "{hf_user}/my-custom-traces" } +``` + +The shared `smolagents/ml-intern-sessions` dataset is unrelated and only +receives anonymized telemetry rows used by the backend KPI scheduler. + +## Supported Gateways + +ML Intern currently supports one-way notification gateways from CLI sessions. +These gateways send out-of-band status updates; they do not accept inbound chat +messages. + +### Slack + +Slack notifications use the Slack Web API to post messages when the agent needs +approval, hits an error, or completes a turn. Create a Slack app with a bot token +that has `chat:write`, invite the bot to the target channel, then set: + +```bash +SLACK_BOT_TOKEN=xoxb-... +SLACK_CHANNEL_ID=C... +``` + +The CLI automatically creates a `slack.default` destination when both variables +are present. Optional environment variables for the env-only default: + +```bash +ML_INTERN_SLACK_NOTIFICATIONS=false +ML_INTERN_SLACK_DESTINATION=slack.ops +ML_INTERN_SLACK_AUTO_EVENTS=approval_required,error,turn_complete +ML_INTERN_SLACK_ALLOW_AGENT_TOOL=true +ML_INTERN_SLACK_ALLOW_AUTO_EVENTS=true +``` + +For a persistent user-level config, put overrides in +`~/.config/ml-intern/cli_agent_config.json` or point `ML_INTERN_CLI_CONFIG` at a +JSON file: + +```json +{ + "messaging": { + "enabled": true, + "auto_event_types": ["approval_required", "error", "turn_complete"], + "destinations": { + "slack.ops": { + "provider": "slack", + "token": "${SLACK_BOT_TOKEN}", + "channel": "${SLACK_CHANNEL_ID}", + "allow_agent_tool": true, + "allow_auto_events": true + } + } + } +} ``` ## Architecture @@ -185,6 +341,18 @@ The agent emits the following events via `event_queue`: ## Development +### Pre-commit Checks + +Run Ruff before every commit: + +```bash +uv run ruff check . +uv run ruff format --check . +``` + +If the format check fails, run `uv run ruff format .` and re-run the checks +before committing. + ### Adding Built-in Tools Edit `agent/core/tools.py`: @@ -210,11 +378,12 @@ def create_builtin_tools() -> list[ToolSpec]: ### Adding MCP Servers -Edit `configs/main_agent_config.json`: +Edit `configs/cli_agent_config.json` for CLI defaults, or +`configs/frontend_agent_config.json` for web-session defaults: ```json { - "model_name": "anthropic/claude-sonnet-4-5-20250929", + "model_name": "zai-org/GLM-5.2:novita", "mcpServers": { "your-server-name": { "transport": "http", @@ -228,3 +397,14 @@ Edit `configs/main_agent_config.json`: ``` Note: Environment variables like `${YOUR_TOKEN}` are auto-substituted from `.env`. + +## Cite ml-intern +If you use `ml-intern` in your work, please cite it by using the following BibTeX entry or similar. +```bibtex +@Misc{ml-intern, + title = {ml-intern: an agent that autonomously researches, writes, and ships good quality ML related code using the Hugging Face ecosystem}, + author = {Aksel Joonas Reedi, Henri Bonamy, Yoan Di Cosmo, Leandro von Werra, Lewis Tunstall}, + howpublished = {\url{https://github.com/huggingface/ml-intern}}, + year = {2026} +} +``` diff --git a/REVIEW.md b/REVIEW.md new file mode 100644 index 000000000..3f08c60a8 --- /dev/null +++ b/REVIEW.md @@ -0,0 +1,135 @@ +# Review instructions + +These rules override the default review guidance. Treat them as the highest-priority +instruction block for any review of this repo. If something here contradicts a more +generic review habit, follow these. + +## Severity levels + +Every finding carries one of three priority labels: + +- **P0** — blocks merge. +- **P1** — worth fixing, not blocking. +- **P2** — informational. + +Write labels as plain text (`P0`, `P1`, `P2`) in finding headers. Do not use +emoji or colored markers. Use judgment on what belongs at which level — this +repo does not enumerate P0 cases; read the code and decide. + +## Default bias: rigor + +Reviews gate merges. This is an open-source repo that takes PRs from anyone; the +maintainer team is small and relies on the review to catch what they don't have +time to verify themselves. **Default bias is rigor, not speed.** When in doubt +on a P0-class concern, investigate further before deciding whether to flag — a +false negative ships a bug to production, a false positive costs the contributor +one round trip. + +Rigor is not nitpicking. The P1 cap, "do not report" skip list, and verification +bar all still apply. Rigor means going deep on a small number of real concerns, +not surfacing a large number of shallow ones. Prefer one well-investigated P0 +over three speculative P1s. + +**Hold the line on P0.** If the author pushes back on a P0 finding without a fix +that actually addresses the root cause, re-state the concern with added +citations. Only accept the pushback if the author points to code or behavior you +missed. Do not soften a P0 because the contributor is polite or new to the repo. + +For P1 and P2: if the author defers or pushes back without fixing, accept it +silently — do not re-flag on subsequent commits. P1/P2 are informational; the +author may defer to a follow-up issue at their discretion. + +If Claude and the author repeatedly disagree on the same class of finding, the +signal is that REVIEW.md is missing a rule; note it once in the PR summary as +`suggest-rule: ` and stop. + +## Investigate before posting + +The depth of your analysis determines the strength of your finding. For any +P0-class concern, before writing it up: + +- Read the relevant callers and callees, not just the diff. Use Read and Grep + to open files the diff doesn't touch but the changed code interacts with. +- Trace the full chain end-to-end for routing, auth, and agent-loop findings. + Cite each hop by `file:line`, not just the suspicious line. +- Check whether the codebase already has an established pattern for this kind + of change (`grep` for similar call sites, similar tool definitions, similar + route guards). If the PR introduces a new approach where an established + pattern exists, flag that — divergence from the existing pattern is usually a + regression vector even when the new code "works." +- Confirm the specific behavior you're claiming. "This breaks X" must be + grounded in either the code handling X or a test exercising X, not in + inference from naming or structure. + +A finding you "spotted" by scanning the diff is more likely to be a false +positive than a finding you verified by reading the code around it. + +## P1 cap + +Report at most **3** P1 findings per review. If you found more, say "plus N +similar items" in the summary. If everything you found is P1 or below, open the +summary with "No blocking issues." + +## Re-review convergence + +If this PR has already received a Claude review (there is a prior review comment +by the `claude` bot), suppress new P1 findings and post only P0 ones. Do not +re-post P1s that were already flagged on earlier commits. If the author pushed a +fix for a previously flagged issue, acknowledge it in one line rather than +re-flagging. + +## Do not report + +Anything in these paths — skip entirely: + +- `frontend/node_modules/**`, `**/*.lock`, `uv.lock`, `package-lock.json` +- `hf_agent.egg-info/**`, `.ruff_cache/**`, `.pytest_cache/**`, `.venv/**` +- `session_logs/**`, `reports/**` +- Anything under a `gen/` or `generated/` path + +Anything speculative — do not post: + +- "This might be slow" without a concrete complexity claim tied to a specific + input size +- Hypothetical race conditions without a concrete interleaving + +## Dependency PRs + +For PRs whose diff is only a lockfile bump, a `pyproject.toml` change, or a +new dependency, the code rules above don't apply — risks shift to provenance +and framing. Every claim in the title or body (CVE IDs, version numbers, +behavior fixes) must match what the diff actually does, and any new +transitive dep needs justification. A PR that lies in its framing is P0 +regardless of whether the code change is safe in isolation. + +## Verification bar + +Every behavior claim in a finding must cite `file:line`. "This breaks X" is not +actionable without a line reference. If you cannot cite a line, do not post +the finding. + +## Summary shape + +Open the review body with a single-line tally and an explicit merge verdict, on +two lines: + +``` +2 P0, 3 P1 +Verdict: changes requested +``` + +Valid verdicts: + +- **Verdict: ready to merge** — no P0 findings, contributor can merge as-is + once any CI passes +- **Verdict: changes requested** — at least one P0 that must be addressed + before merging +- **Verdict: needs discussion** — a design-level concern the maintainer should + weigh in on before the contributor iterates (use sparingly) + +If it's a clean review, write `LGTM` followed by `Verdict: ready to merge`. + +Then a **What I checked** bullet list — one line per major area you examined, +regardless of whether you found anything. This gives the maintainer visible +coverage at a glance and lets them decide whether to spot-check areas you +didn't touch. diff --git a/agent/README.md b/agent/README.md index 5cd172de4..567ce033d 100644 --- a/agent/README.md +++ b/agent/README.md @@ -7,7 +7,7 @@ Async agent loop with LiteLLM. **Queue-based async system:** - Submissions in (user input) → Agent Loop → Events output for possible UI updates - Session maintains state (context + tools) for possible future Context Engineering -- Handlers operations like (USER_INPUT, INTERRUPT, COMPACT, UNDO, SHUTDOWN) for possible UI control +- Handlers operations like (USER_INPUT, COMPACT, UNDO, SHUTDOWN) for possible UI control ## Components diff --git a/agent/__init__.py b/agent/__init__.py index 2e301c8d7..0e4cf34cd 100644 --- a/agent/__init__.py +++ b/agent/__init__.py @@ -8,10 +8,8 @@ # backend entries share the same config. # drop_params: quietly drop unsupported params rather than raising # suppress_debug_info: hide the noisy "Give Feedback" banner on errors -# modify_params: let LiteLLM patch Anthropic's tool-call requirements -# (synthesize a dummy tool spec when we call completion on a history -# that contains tool_calls but aren't passing `tools=` — happens -# during summarization / session seeding). +# modify_params: let LiteLLM patch provider-specific schema requirements +# for router-compatible request bodies when possible. litellm.drop_params = True litellm.suppress_debug_info = True litellm.modify_params = True diff --git a/agent/config.py b/agent/config.py index b7e698ad7..c6784db92 100644 --- a/agent/config.py +++ b/agent/config.py @@ -2,21 +2,23 @@ import os import re from pathlib import Path -from typing import Any, Union +from typing import Any, Literal, Union from dotenv import load_dotenv - -# Project root: two levels up from this file (agent/config.py -> project root) -_PROJECT_ROOT = Path(__file__).resolve().parent.parent from fastmcp.mcp_config import ( RemoteMCPServer, StdioMCPServer, ) from pydantic import BaseModel +from agent.messaging.models import MessagingConfig + # These two are the canonical server config types for MCP servers. MCPServerConfig = Union[StdioMCPServer, RemoteMCPServer] +# Project root: two levels up from this file (agent/config.py -> project root) +_PROJECT_ROOT = Path(__file__).resolve().parent.parent + class Config(BaseModel): """Configuration manager""" @@ -24,24 +26,138 @@ class Config(BaseModel): model_name: str mcpServers: dict[str, MCPServerConfig] = {} save_sessions: bool = True - session_dataset_repo: str = "akseljoonas/hf-agent-sessions" - auto_save_interval: int = 3 # Save every N user turns (0 = disabled) + session_dataset_repo: str = "smolagents/ml-intern-sessions" + # Per-user private dataset that mirrors each session in Claude Code JSONL + # format so the HF Agent Trace Viewer auto-renders it + # (https://huggingface.co/changelog/agent-trace-viewer). Created private + # on first use; user flips it public via /share-traces. ``{hf_user}`` is + # substituted at upload time from the authenticated HF username. + share_traces: bool = True + personal_trace_repo_template: str = "{hf_user}/ml-intern-sessions" + auto_save_interval: int = 1 # Save every N user turns (0 = disabled) + # Mid-turn heartbeat: save + upload every N seconds while events are being + # emitted. Guards against losing trace data on long-running turns that + # crash before turn_complete (e.g. a multi-hour hf_jobs wait that OOMs). + # 0 = disabled. Consumed by agent.core.telemetry.HeartbeatSaver. + heartbeat_interval_s: int = 60 yolo_mode: bool = False # Auto-approve all tool calls without confirmation max_iterations: int = 300 # Max LLM calls per agent turn (-1 = unlimited) # Permission control parameters confirm_cpu_jobs: bool = True auto_file_upload: bool = False + tool_runtime: Literal["local", "sandbox"] = "local" # Reasoning effort *preference* — the ceiling the user wants. The probe # on `/model` walks a cascade down from here (``max`` → ``xhigh`` → ``high`` # → …) and caches per-model what the provider actually accepted in - # ``Session.model_effective_effort``. Default ``max`` because we'd rather - # burn tokens thinking than ship a wrong ML recipe; the cascade lands on - # whichever level the model supports (``high`` for GPT-5 / HF router, - # ``xhigh`` or ``max`` for Anthropic 4.6 / 4.7). ``None`` = thinking off. + # ``Session.model_effective_effort``. Default ``high`` because HF Router + # accepts low/medium/high generically and provider-specific higher levels + # should be discovered through explicit probes. ``None`` = thinking off. # Valid values: None | "minimal" | "low" | "medium" | "high" | "xhigh" | "max" - reasoning_effort: str | None = "max" + reasoning_effort: str | None = "high" + messaging: MessagingConfig = MessagingConfig() + + +USER_CONFIG_ENV_VAR = "ML_INTERN_CLI_CONFIG" +DEFAULT_USER_CONFIG_PATH = ( + Path.home() / ".config" / "ml-intern" / "cli_agent_config.json" +) +SLACK_DEFAULT_DESTINATION = "slack.default" +SLACK_DEFAULT_AUTO_EVENT_TYPES = ["approval_required", "error", "turn_complete"] + + +def _deep_merge_config( + base: dict[str, Any], override: dict[str, Any] +) -> dict[str, Any]: + merged = dict(base) + for key, value in override.items(): + current = merged.get(key) + if isinstance(current, dict) and isinstance(value, dict): + merged[key] = _deep_merge_config(current, value) + else: + merged[key] = value + return merged + + +def _load_json_config(path: Path) -> dict[str, Any]: + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, dict): + raise ValueError(f"Config file {path} must contain a JSON object") + return data + + +def _load_user_config() -> dict[str, Any]: + raw_path = os.environ.get(USER_CONFIG_ENV_VAR) + if raw_path: + path = Path(raw_path).expanduser() + if not path.exists(): + raise FileNotFoundError( + f"{USER_CONFIG_ENV_VAR} points to missing config file: {path}" + ) + return _load_json_config(path) + + if DEFAULT_USER_CONFIG_PATH.exists(): + return _load_json_config(DEFAULT_USER_CONFIG_PATH) + return {} + + +def _env_bool(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None: + return default + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + return default + + +def _env_list(name: str) -> list[str] | None: + value = os.environ.get(name) + if value is None: + return None + return [item.strip() for item in value.split(",") if item.strip()] + + +def apply_slack_user_defaults(raw_config: dict[str, Any]) -> dict[str, Any]: + """Enable a default Slack destination from user env vars, when present.""" + if not _env_bool("ML_INTERN_SLACK_NOTIFICATIONS", True): + return raw_config + + token = os.environ.get("SLACK_BOT_TOKEN") + channel = os.environ.get("SLACK_CHANNEL_ID") or os.environ.get("SLACK_CHANNEL") + if not token or not channel: + return raw_config + + config = dict(raw_config) + messaging = dict(config.get("messaging") or {}) + destinations = dict(messaging.get("destinations") or {}) + destination_name = ( + os.environ.get("ML_INTERN_SLACK_DESTINATION") or SLACK_DEFAULT_DESTINATION + ).strip() + + if destination_name not in destinations: + destinations[destination_name] = { + "provider": "slack", + "token": token, + "channel": channel, + "allow_agent_tool": _env_bool("ML_INTERN_SLACK_ALLOW_AGENT_TOOL", True), + "allow_auto_events": _env_bool("ML_INTERN_SLACK_ALLOW_AUTO_EVENTS", True), + } + + auto_events = _env_list("ML_INTERN_SLACK_AUTO_EVENTS") + if auto_events is not None: + messaging["auto_event_types"] = auto_events + elif "auto_event_types" not in messaging: + messaging["auto_event_types"] = SLACK_DEFAULT_AUTO_EVENT_TYPES + + messaging["enabled"] = True + messaging["destinations"] = destinations + config["messaging"] = messaging + return config def substitute_env_vars(obj: Any) -> Any: @@ -81,7 +197,10 @@ def replacer(match): return obj -def load_config(config_path: str = "config.json") -> Config: +def load_config( + config_path: str = "config.json", + include_user_defaults: bool = False, +) -> Config: """ Load configuration with environment variable substitution. @@ -93,8 +212,10 @@ def load_config(config_path: str = "config.json") -> Config: load_dotenv(_PROJECT_ROOT / ".env") load_dotenv(override=False) - with open(config_path, "r") as f: - raw_config = json.load(f) + raw_config = _load_json_config(Path(config_path)) + if include_user_defaults: + raw_config = _deep_merge_config(raw_config, _load_user_config()) + raw_config = apply_slack_user_defaults(raw_config) config_with_env = substitute_env_vars(raw_config) return Config.model_validate(config_with_env) diff --git a/agent/context_manager/manager.py b/agent/context_manager/manager.py index 373fb11ec..1858e926d 100644 --- a/agent/context_manager/manager.py +++ b/agent/context_manager/manager.py @@ -3,7 +3,7 @@ """ import logging -import os +import time import zoneinfo from datetime import datetime from pathlib import Path @@ -13,7 +13,11 @@ from jinja2 import Template from litellm import Message, acompletion -from agent.core.prompt_caching import with_prompt_caching +from agent.core.prompt_caching import ( + router_session_id_for, + with_prompt_cache_params, + with_prompt_caching, +) logger = logging.getLogger(__name__) @@ -78,6 +82,24 @@ def _get_hf_username(hf_token: str | None = None) -> str: "will be have to be filled in." ) +# Per-message ceiling. If a single message in the "untouched" tail is larger +# than this, compaction can't recover even after summarizing the middle — +# producing the infinite compaction loop seen 2026-05-03 in pod logs (200k +# context shrinks to 200k+ because one tool output is 80k tokens). We replace +# such messages with a placeholder before compaction runs. +_MAX_TOKENS_PER_MESSAGE = 50_000 + + +class CompactionFailedError(Exception): + """Raised when compaction can't reduce context below the threshold. + + Typically means an individual preserved message (system, first user, or + untouched tail) exceeds what truncation can fix in one pass. The caller + must terminate the session; retrying produces an infinite loop that burns + hosted inference budget. + """ + + # Used when seeding a brand-new session from prior browser-cached messages. # Here we're writing a note to *ourselves* — so preserve the tool-call trail, # files produced, and planned next steps in first person. Optimized for @@ -102,6 +124,8 @@ async def summarize_messages( max_tokens: int = 2000, tool_specs: list[dict] | None = None, prompt: str = _COMPACT_PROMPT, + session: Any = None, + kind: str = "compaction", ) -> tuple[str, int]: """Run a summarization prompt against a list of messages. @@ -110,21 +134,57 @@ async def summarize_messages( instead — it preserves the tool-call trail so the agent can answer follow-up questions about what it did. + ``session`` is optional; when provided, the call is recorded via + ``telemetry.record_llm_call`` so its cost lands in the session's + ``total_cost_usd``. Without it, the call still happens but is + invisible in telemetry, which used to hide a significant share of hosted + inference spend. + Returns ``(summary_text, completion_tokens)``. """ from agent.core.llm_params import _resolve_llm_params prompt_messages = list(messages) + [Message(role="user", content=prompt)] - llm_params = _resolve_llm_params(model_name, hf_token, reasoning_effort="high") + llm_params = _resolve_llm_params( + model_name, + hf_token, + reasoning_effort="high", + ) + llm_params = with_prompt_cache_params( + llm_params, + session_id=router_session_id_for(session), + ) + llm_params = {**llm_params, "max_completion_tokens": max_tokens} prompt_messages, tool_specs = with_prompt_caching( - prompt_messages, tool_specs, llm_params.get("model") + prompt_messages, tool_specs, llm_params ) + _t0 = time.monotonic() response = await acompletion( messages=prompt_messages, - max_completion_tokens=max_tokens, tools=tool_specs, **llm_params, ) + if session is not None: + from agent.core import telemetry + from agent.core.yolo_budget import maybe_pause_yolo_after_spend + + usage = await telemetry.record_llm_call( + session, + model=model_name, + response=response, + latency_ms=int((time.monotonic() - _t0) * 1000), + finish_reason=response.choices[0].finish_reason + if response.choices + else None, + kind=kind, + ) + await maybe_pause_yolo_after_spend( + session, + spend_kind=kind, + observed_cost_usd=usage.get("cost_usd") + if isinstance(usage, dict) + else None, + ) summary = response.choices[0].message.content or "" completion_tokens = response.usage.completion_tokens if response.usage else 0 return summary, completion_tokens @@ -141,13 +201,23 @@ def __init__( tool_specs: list[dict[str, Any]] | None = None, prompt_file_suffix: str = "system_prompt_v3.yaml", hf_token: str | None = None, + hf_username: str | None = None, local_mode: bool = False, + autonomous_mode: bool = False, ): + self.prompt_file_suffix = prompt_file_suffix + self.tool_specs = tool_specs or [] + self.hf_token = hf_token + self.hf_username = hf_username + self.local_mode = local_mode + self.autonomous_mode = autonomous_mode self.system_prompt = self._load_system_prompt( - tool_specs or [], - prompt_file_suffix="system_prompt_v3.yaml", + self.tool_specs, + prompt_file_suffix=self.prompt_file_suffix, hf_token=hf_token, + hf_username=hf_username, local_mode=local_mode, + autonomous_mode=autonomous_mode, ) # The model's real input-token ceiling (from litellm.get_model_info). # Compaction triggers at _COMPACT_THRESHOLD_RATIO below it — see @@ -160,13 +230,48 @@ def __init__( self.running_context_usage = 0 self.untouched_messages = untouched_messages self.items: list[Message] = [Message(role="system", content=self.system_prompt)] + self.on_message_added = None + + def refresh_system_prompt( + self, + *, + tool_specs: list[dict[str, Any]] | None = None, + hf_token: str | None = None, + hf_username: str | None = None, + local_mode: bool | None = None, + autonomous_mode: bool | None = None, + ) -> Message: + """Re-render the system prompt and return it as a system message.""" + if tool_specs is not None: + self.tool_specs = tool_specs + if hf_token is not None: + self.hf_token = hf_token + if hf_username is not None: + self.hf_username = hf_username + if local_mode is not None: + self.local_mode = local_mode + if autonomous_mode is not None: + self.autonomous_mode = autonomous_mode + self.system_prompt = self._load_system_prompt( + self.tool_specs, + prompt_file_suffix=getattr( + self, "prompt_file_suffix", "system_prompt_v3.yaml" + ), + hf_token=getattr(self, "hf_token", None), + hf_username=getattr(self, "hf_username", None), + local_mode=getattr(self, "local_mode", False), + autonomous_mode=getattr(self, "autonomous_mode", False), + ) + return Message(role="system", content=self.system_prompt) def _load_system_prompt( self, tool_specs: list[dict[str, Any]], prompt_file_suffix: str = "system_prompt.yaml", hf_token: str | None = None, + hf_username: str | None = None, local_mode: bool = False, + autonomous_mode: bool = False, ): """Load and render the system prompt from YAML file with Jinja2""" prompt_file = Path(__file__).parent.parent / "prompts" / f"{prompt_file_suffix}" @@ -182,18 +287,21 @@ def _load_system_prompt( current_time = now.strftime("%H:%M:%S.%f")[:-3] current_timezone = f"{now.strftime('%Z')} (UTC{now.strftime('%z')[:3]}:{now.strftime('%z')[3:]})" - # Get HF user info from OAuth token - hf_user_info = _get_hf_username(hf_token) + # Prefer the username already resolved by the caller; fall back to a + # token lookup for contexts that construct ContextManager directly. + hf_user_info = hf_username or _get_hf_username(hf_token) template = Template(template_str) static_prompt = template.render( tools=tool_specs, num_tools=len(tool_specs), + autonomous_mode=autonomous_mode, ) # CLI-specific context for local mode if local_mode: import os + cwd = os.getcwd() local_context = ( f"\n\n# CLI / Local mode\n\n" @@ -211,7 +319,7 @@ def _load_system_prompt( f"{static_prompt}\n\n" f"[Session context: Date={current_date}, Time={current_time}, " f"Timezone={current_timezone}, User={hf_user_info}, " - f"Tools={len(tool_specs)}]" + f"Tools={len(tool_specs)}, Autonomous={str(autonomous_mode).lower()}]" ) def add_message(self, message: Message, token_count: int = None) -> None: @@ -219,6 +327,8 @@ def add_message(self, message: Message, token_count: int = None) -> None: if token_count: self.running_context_usage = token_count self.items.append(message) + if self.on_message_added: + self.on_message_added(message) def get_messages(self) -> list[Message]: """Get all messages for sending to LLM. @@ -253,45 +363,53 @@ def _normalize_tool_calls(msg: Message) -> None: def _patch_dangling_tool_calls(self) -> None: """Add stub tool results for any tool_calls that lack a matching result. - Scans backwards to find the last assistant message with tool_calls, - which may not be items[-1] if some tool results were already added. + Ensures each assistant message's tool_calls are followed immediately + by matching tool-result messages. This has to work across the whole + history, not just the most recent turn, because a cancelled tool use + in an earlier turn can still poison the next provider request. """ if not self.items: return - # Find the last assistant message with tool_calls - assistant_msg = None - for i in range(len(self.items) - 1, -1, -1): + i = 0 + while i < len(self.items): msg = self.items[i] - if getattr(msg, "role", None) == "assistant" and getattr( + if getattr(msg, "role", None) != "assistant" or not getattr( msg, "tool_calls", None ): - assistant_msg = msg - break - # Stop scanning once we hit a user message — anything before - # that belongs to a previous (complete) turn. - if getattr(msg, "role", None) == "user": - break + i += 1 + continue + + self._normalize_tool_calls(msg) + + # Consume the contiguous tool-result block that immediately follows + # this assistant message. Any missing tool ids must be inserted + # before the next non-tool message to satisfy provider ordering. + j = i + 1 + immediate_ids: set[str | None] = set() + while ( + j < len(self.items) and getattr(self.items[j], "role", None) == "tool" + ): + immediate_ids.add(getattr(self.items[j], "tool_call_id", None)) + j += 1 + + missing: list[Message] = [] + for tc in msg.tool_calls: + if tc.id not in immediate_ids: + missing.append( + Message( + role="tool", + content="Tool was not executed (interrupted or error).", + tool_call_id=tc.id, + name=tc.function.name, + ) + ) - if not assistant_msg: - return + if missing: + self.items[j:j] = missing + j += len(missing) - self._normalize_tool_calls(assistant_msg) - answered_ids = { - getattr(m, "tool_call_id", None) - for m in self.items - if getattr(m, "role", None) == "tool" - } - for tc in assistant_msg.tool_calls: - if tc.id not in answered_ids: - self.items.append( - Message( - role="tool", - content="Tool was not executed (interrupted or error).", - tool_call_id=tc.id, - name=tc.function.name, - ) - ) + i = j def undo_last_turn(self) -> bool: """Remove the last complete turn (user msg + all assistant/tool msgs that follow). @@ -341,15 +459,106 @@ def compaction_threshold(self) -> int: @property def needs_compaction(self) -> bool: - return self.running_context_usage > self.compaction_threshold and bool(self.items) + return self.running_context_usage > self.compaction_threshold and bool( + self.items + ) + + def _truncate_oversized( + self, messages: list[Message], model_name: str + ) -> list[Message]: + """Replace any message > _MAX_TOKENS_PER_MESSAGE with a placeholder. + + These are typically tool outputs (CSV dumps, file contents) sitting in + the untouched tail or first-user position that compaction can't shrink + — they pass through verbatim, keeping context above threshold and + triggering an infinite compaction retry loop. + """ + from litellm import token_counter + + out: list[Message] = [] + for msg in messages: + # System messages are sacred — they're the agent's instructions. + # In edge cases (items < untouched_messages), the slice math in + # compact() can let items[0] (the system message) leak into the + # recent_messages list. Defense-in-depth: never truncate it. + if msg.role == "system": + out.append(msg) + continue + try: + n = token_counter(model=model_name, messages=[msg.model_dump()]) + except Exception: + # token_counter occasionally fails on edge-case content; + # don't drop the message, just keep it as-is. + out.append(msg) + continue + if n <= _MAX_TOKENS_PER_MESSAGE: + out.append(msg) + continue + placeholder = ( + f"[truncated for compaction — original was {n} tokens, " + f"removed to keep context under {self.compaction_threshold} tokens]" + ) + logger.warning( + "Truncating %s message: %d -> %d tokens for compaction", + msg.role, + n, + len(placeholder) // 4, + ) + # Preserve all known assistant-side fields (tool_calls, thinking_blocks, + # reasoning_content, provider_specific_fields) even when content is + # replaced. Historical traces may still contain provider reasoning + # metadata, and truncation should not silently discard it. + kept = { + k: getattr(msg, k, None) + for k in ( + "tool_call_id", + "tool_calls", + "name", + "thinking_blocks", + "reasoning_content", + "provider_specific_fields", + ) + if getattr(msg, k, None) is not None + } + out.append(Message(role=msg.role, content=placeholder, **kept)) + return out + + def _recompute_usage(self, model_name: str) -> None: + """Refresh ``running_context_usage`` from current items via real tokenizer.""" + from litellm import token_counter + + try: + self.running_context_usage = token_counter( + model=model_name, + messages=[m.model_dump() for m in self.items], + ) + except Exception as e: + logger.warning("token_counter failed (%s); rough estimate", e) + # Rough fallback: 4 chars per token. + self.running_context_usage = ( + sum(len(getattr(m, "content", "") or "") for m in self.items) // 4 + ) async def compact( self, model_name: str, tool_specs: list[dict] | None = None, hf_token: str | None = None, + session: Any = None, ) -> None: - """Remove old messages to keep history under target size""" + """Remove old messages to keep history under target size. + + ``session`` is optional — if passed, the underlying summarization + LLM call is recorded via ``telemetry.record_llm_call(kind= + "compaction")`` so its cost shows up in ``total_cost_usd``. + + Raises ``CompactionFailedError`` if the post-compact context is still + over the threshold. This happens when a preserved message (typically + a giant tool output stuck in the untouched tail) is too large for + truncation to fix. The caller must terminate the session — retrying + is what caused the 2026-05-03 infinite-compaction-loop pattern that + burned hosted inference budget invisibly. + """ if not self.needs_compaction: return @@ -372,12 +581,45 @@ async def compact( idx = len(self.items) - self.untouched_messages while idx > 1 and self.items[idx].role != "user": idx -= 1 + # The real invariant is "idx must be strictly after first_user_idx, + # otherwise recent_messages overlaps with the messages we put in + # head". The walk-back's `idx > 1` guard is necessary (no system in + # recent) but insufficient (first_user is also in head and would be + # duplicated). Chat providers can reject two consecutive user messages + # with a 400 — bot review on PR #213 caught this on the second clamp + # iteration. + if idx <= first_user_idx: + idx = first_user_idx + 1 recent_messages = self.items[idx:] - messages_to_summarize = self.items[first_user_idx + 1:idx] - - # improbable, messages would have to very long + messages_to_summarize = self.items[first_user_idx + 1 : idx] + + # Truncate any message that's larger than _MAX_TOKENS_PER_MESSAGE in + # the parts we PRESERVE through compaction (first_user + recent_tail). + # These are the only places where individual messages can defeat + # compaction by being intrinsically too large. Messages in + # ``messages_to_summarize`` are folded into the summary, so their size + # doesn't matter on its own. + if first_user_msg is not None: + truncated = self._truncate_oversized([first_user_msg], model_name) + first_user_msg = truncated[0] + recent_messages = self._truncate_oversized(recent_messages, model_name) + + # If there's nothing to summarize but the preserved messages are now + # truncated and small, just rebuild and recompute. This is rare but + # avoids returning silently with the old (over-threshold) state. if not messages_to_summarize: + head = [system_msg] if system_msg else [] + if first_user_msg: + head.append(first_user_msg) + self.items = head + recent_messages + self._recompute_usage(model_name) + if self.running_context_usage > self.compaction_threshold: + raise CompactionFailedError( + f"Nothing to summarize but context ({self.running_context_usage}) " + f"still over threshold ({self.compaction_threshold}) after truncation. " + f"System prompt or first user message likely exceeds the budget." + ) return summary, completion_tokens = await summarize_messages( @@ -387,8 +629,13 @@ async def compact( max_tokens=self.compact_size, tool_specs=tool_specs, prompt=_COMPACT_PROMPT, + session=session, + kind="compaction", + ) + summarized_message = Message( + role="assistant", + content=summary, ) - summarized_message = Message(role="assistant", content=summary) # Reconstruct: system + first user msg + summary + recent messages head = [system_msg] if system_msg else [] @@ -396,16 +643,19 @@ async def compact( head.append(first_user_msg) self.items = head + [summarized_message] + recent_messages - # Count the actual post-compact context — system prompt + first user - # turn + summary + the preserved tail all contribute, not just the - # summary. litellm.token_counter uses the model's real tokenizer. - from litellm import token_counter - - try: - self.running_context_usage = token_counter( - model=model_name, - messages=[m.model_dump() for m in self.items], + self._recompute_usage(model_name) + + # Hard verify: if compaction didn't bring us below the threshold even + # after truncating oversized preserved messages, retrying just burns + # hosted inference budget on the same useless compaction call. Raise so the + # caller can terminate the session cleanly. Pre-2026-05-04, the + # caller looped indefinitely (~$3/Opus retry) until the pod was + # killed — invisible to the dataset because the session never + # finished cleanly. + if self.running_context_usage > self.compaction_threshold: + raise CompactionFailedError( + f"Compaction ineffective: {self.running_context_usage} tokens " + f"still over threshold {self.compaction_threshold} after summarize " + f"and truncation. Likely the system prompt + first user + summary " + f"+ truncated tail still exceeds budget." ) - except Exception as e: - logger.warning("token_counter failed post-compact (%s); falling back to rough estimate", e) - self.running_context_usage = len(self.system_prompt) // 4 + completion_tokens diff --git a/agent/core/agent_loop.py b/agent/core/agent_loop.py index c3fd88bc8..2ba2eab2b 100644 --- a/agent/core/agent_loop.py +++ b/agent/core/agent_loop.py @@ -5,24 +5,200 @@ import asyncio import json import logging -import os -from dataclasses import dataclass - -from litellm import ChatCompletionMessageToolCall, Message, acompletion +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from litellm import ( + ChatCompletionMessageToolCall, + Message, + acompletion, +) from litellm.exceptions import ContextWindowExceededError from agent.config import Config +from agent.core.approval_policy import ( + is_scheduled_operation, + normalize_tool_operation, +) +from agent.core.cost_estimation import CostEstimate, estimate_tool_cost +from agent.messaging.gateway import NotificationGateway +from agent.core import telemetry from agent.core.doom_loop import check_for_doom_loop +from agent.core.hf_access import ( + HF_BILLING_URL, + HF_PRO_SUBSCRIBE_URL, + is_inference_billing_error, +) from agent.core.llm_params import _resolve_llm_params -from agent.core.prompt_caching import with_prompt_caching -from agent.core.session import Event, OpType, Session +from agent.core.prompt_caching import ( + router_session_id_for, + with_prompt_cache_params, + with_prompt_caching, +) +from agent.core.session import DEFAULT_SESSION_LOG_DIR, Event, OpType, Session from agent.core.tools import ToolRouter +from agent.core.usage_thresholds import ( + USAGE_THRESHOLD_TOOL_NAME, + is_usage_threshold_pending, + next_usage_warning_threshold, +) +from agent.core.yolo_budget import ( + BudgetDecision, + check_session_budget, + is_yolo_budget_pending, + maybe_pause_yolo_after_spend, + release_budget_reservation, + reserve_session_budget, + yolo_budget_can_resume, + yolo_budget_pending_to_tool, +) from agent.tools.jobs_tool import CPU_FLAVORS +from agent.tools.sandbox_tool import ( + DEFAULT_CPU_SANDBOX_HARDWARE, + start_cpu_sandbox_preload, + teardown_session_sandbox, +) logger = logging.getLogger(__name__) ToolCall = ChatCompletionMessageToolCall +_MALFORMED_TOOL_PREFIX = "ERROR: Tool call to '" +_MALFORMED_TOOL_SUFFIX = "' had malformed JSON arguments" +_NO_TOOL_INCOMPLETE_PLAN_RETRY_LIMIT = 2 + + +def _unfinished_plan_items(session: Session) -> list[dict[str, str]]: + plan = getattr(session, "current_plan", None) or [] + unfinished: list[dict[str, str]] = [] + for item in plan: + if not isinstance(item, dict): + continue + status = item.get("status") + if status in {"pending", "in_progress"}: + unfinished.append(item) + return unfinished + + +def _format_plan_items_for_guard(items: list[dict[str, str]], limit: int = 4) -> str: + formatted = [] + for item in items[:limit]: + item_id = item.get("id") or "?" + content = item.get("content") or "(unnamed task)" + status = item.get("status") or "unknown" + formatted.append(f"{item_id}. {content} [{status}]") + if len(items) > limit: + formatted.append(f"... and {len(items) - limit} more") + return "; ".join(formatted) + + +def _no_tool_incomplete_plan_prompt(items: list[dict[str, str]]) -> str: + summary = _format_plan_items_for_guard(items) + return ( + "[SYSTEM: CONTINUATION GUARD] Your previous response ended without any " + "tool calls, but the task is not complete. The current plan still has " + f"unfinished items: {summary}. Do not return control to the user yet. " + "Continue from the next unfinished item and make at least one tool call " + "now. If you genuinely cannot continue, first use tools to inspect the " + "state or verify the blocker." + ) + + +def _malformed_tool_name(message: Message) -> str | None: + """Return the tool name for malformed-json tool-result messages.""" + if getattr(message, "role", None) != "tool": + return None + content = getattr(message, "content", None) + if not isinstance(content, str): + return None + if not content.startswith(_MALFORMED_TOOL_PREFIX): + return None + end = content.find(_MALFORMED_TOOL_SUFFIX, len(_MALFORMED_TOOL_PREFIX)) + if end == -1: + return None + return content[len(_MALFORMED_TOOL_PREFIX) : end] + + +def _detect_repeated_malformed( + items: list[Message], + threshold: int = 2, +) -> str | None: + """Return the repeated malformed tool name if the tail contains a streak. + + Walk backward over the current conversation tail. A streak counts only + consecutive malformed tool-result messages for the same tool; any other + tool result breaks it. + """ + if threshold <= 0: + return None + + streak_tool: str | None = None + streak = 0 + + for item in reversed(items): + if getattr(item, "role", None) != "tool": + continue + + malformed_tool = _malformed_tool_name(item) + if malformed_tool is None: + break + + if streak_tool is None: + streak_tool = malformed_tool + streak = 1 + elif malformed_tool == streak_tool: + streak += 1 + else: + break + + if streak >= threshold: + return streak_tool + + return None + + +def _coerce_float(value: Any) -> float: + if isinstance(value, bool) or value is None: + return 0.0 + try: + return float(value) + except (TypeError, ValueError): + return 0.0 + + +def _usage_output_message(pending: dict[str, Any]) -> str: + current = _coerce_float(pending.get("current_spend_usd")) + next_threshold = _coerce_float(pending.get("next_threshold_usd")) + return ( + f"Current-session usage warning acknowledged at ${current:.2f}. " + f"The next warning is at ${next_threshold:.2f}." + ) + + +async def _maybe_pause_for_usage_threshold( + session: Session, + *, + continuation: str, + final_response: str | None = None, +) -> bool: + checker = getattr(session, "usage_threshold_checker", None) + if checker is None or session.pending_approval: + return False + payload: dict[str, Any] = { + "continuation": continuation, + "force_check": continuation == "complete_turn", + "history_size": len(session.context_manager.items), + } + if final_response is not None: + payload["final_response"] = final_response + try: + return bool(await checker(payload)) + except Exception as e: + logger.debug("Usage threshold check failed: %s", e) + return False + def _validate_tool_args(tool_args: dict) -> tuple[bool, str | None]: """ @@ -46,13 +222,42 @@ def _validate_tool_args(tool_args: dict) -> tuple[bool, str | None]: return True, None -def _needs_approval( +_IMMEDIATE_HF_JOB_RUNS = {"run", "uv"} + + +@dataclass(frozen=True) +class ApprovalDecision: + requires_approval: bool + auto_approved: bool = False + auto_approval_blocked: bool = False + block_reason: str | None = None + estimated_cost_usd: float | None = None + remaining_cap_usd: float | None = None + billable: bool = False + + +def _operation(tool_args: dict) -> str: + return normalize_tool_operation(tool_args.get("operation")) + + +def _is_immediate_hf_job_run(tool_name: str, tool_args: dict) -> bool: + return tool_name == "hf_jobs" and _operation(tool_args) in _IMMEDIATE_HF_JOB_RUNS + + +def _is_scheduled_hf_job_run(tool_name: str, tool_args: dict) -> bool: + return tool_name == "hf_jobs" and is_scheduled_operation(_operation(tool_args)) + + +def _is_budgeted_auto_approval_target(tool_name: str, tool_args: dict) -> bool: + return tool_name == "sandbox_create" or _is_immediate_hf_job_run( + tool_name, tool_args + ) + + +def _base_needs_approval( tool_name: str, tool_args: dict, config: Config | None = None ) -> bool: - """Check if a tool call requires user approval before execution.""" - # Yolo mode: skip all approvals - if config and config.yolo_mode: - return False + """Check if a tool call requires approval before YOLO policy is applied.""" # If args are malformed, skip approval (validation error will be shown later) args_valid, _ = _validate_tool_args(tool_args) @@ -60,11 +265,14 @@ def _needs_approval( return False if tool_name == "sandbox_create": - return True + hardware = tool_args.get("hardware") or DEFAULT_CPU_SANDBOX_HARDWARE + return hardware != DEFAULT_CPU_SANDBOX_HARDWARE if tool_name == "hf_jobs": - operation = tool_args.get("operation", "") - if operation not in ["run", "uv", "scheduled run", "scheduled uv"]: + operation = _operation(tool_args) + if is_scheduled_operation(operation): + return True + if operation not in _IMMEDIATE_HF_JOB_RUNS: return False # Check if this is a CPU-only job @@ -116,25 +324,234 @@ def _needs_approval( return False +def _session_auto_approval_enabled(session: Session | None) -> bool: + return bool(session and getattr(session, "auto_approval_enabled", False)) + + +def _effective_yolo_enabled(session: Session | None, config: Config | None) -> bool: + return bool( + (config and config.yolo_mode) or _session_auto_approval_enabled(session) + ) + + +async def _approval_decision( + tool_name: str, + tool_args: dict, + session: Session, + *, + reserved_spend_usd: float = 0.0, +) -> ApprovalDecision: + """Return the approval decision for one parsed tool call.""" + config = session.config + base_requires_approval = _base_needs_approval(tool_name, tool_args, config) + + # Scheduled jobs are recurring/unbounded enough that YOLO never bypasses + # the human confirmation, including legacy config.yolo_mode. + if _is_scheduled_hf_job_run(tool_name, tool_args): + reason = "Scheduled HF jobs always require manual approval." + if _session_auto_approval_enabled(session): + reason = "Scheduled HF jobs require disabling YOLO because their recurring cost is unbounded." + return ApprovalDecision( + requires_approval=True, + auto_approval_blocked=_effective_yolo_enabled(session, config), + block_reason=reason, + ) + + yolo_enabled = _effective_yolo_enabled(session, config) + budgeted_target = _is_budgeted_auto_approval_target(tool_name, tool_args) + + # Cost caps are a session-scoped web policy. Legacy config.yolo_mode + # remains uncapped for CLI/headless, except for scheduled jobs above. + session_yolo_enabled = _session_auto_approval_enabled(session) + if yolo_enabled and budgeted_target and session_yolo_enabled: + estimate = await estimate_tool_cost(tool_name, tool_args, session=session) + budget = check_session_budget( + session, + estimate, + reserved_spend_usd=reserved_spend_usd, + ) + if not budget.allowed: + return ApprovalDecision( + requires_approval=True, + auto_approval_blocked=True, + block_reason=budget.block_reason, + estimated_cost_usd=budget.estimated_cost_usd, + remaining_cap_usd=budget.remaining_cap_usd, + billable=estimate.billable, + ) + if base_requires_approval: + return ApprovalDecision( + requires_approval=False, + auto_approved=True, + estimated_cost_usd=budget.estimated_cost_usd, + remaining_cap_usd=budget.remaining_cap_usd, + billable=estimate.billable, + ) + return ApprovalDecision( + requires_approval=False, + estimated_cost_usd=budget.estimated_cost_usd, + remaining_cap_usd=budget.remaining_cap_usd, + billable=estimate.billable, + ) + + if base_requires_approval and yolo_enabled: + return ApprovalDecision(requires_approval=False, auto_approved=True) + + return ApprovalDecision(requires_approval=base_requires_approval) + + +def _record_estimated_spend( + session: Session, + decision: ApprovalDecision, + *, + reservation_id: str | None = None, +) -> BudgetDecision: + if not decision.billable or decision.estimated_cost_usd is None: + return BudgetDecision(allowed=True, billable=False) + return reserve_session_budget( + session, + CostEstimate( + estimated_cost_usd=decision.estimated_cost_usd, + billable=True, + ), + spend_kind="tool", + reservation_id=reservation_id, + ) + + +async def _record_manual_approved_spend_if_needed( + session: Session, + tool_name: str, + tool_args: dict, + *, + tool_call_id: str | None = None, +) -> BudgetDecision: + if not _session_auto_approval_enabled(session): + return BudgetDecision(allowed=True) + if _is_scheduled_hf_job_run(tool_name, tool_args): + return BudgetDecision( + allowed=False, + billable=True, + block_reason=( + "Scheduled HF jobs require disabling YOLO because their recurring " + "cost is unbounded." + ), + ) + if not _is_budgeted_auto_approval_target(tool_name, tool_args): + return BudgetDecision(allowed=True) + estimate = await estimate_tool_cost(tool_name, tool_args, session=session) + return reserve_session_budget( + session, + estimate, + spend_kind=tool_name, + reservation_id=tool_call_id, + ) + + +async def _check_manual_approved_budget( + session: Session, + tool_name: str, + tool_args: dict, + *, + reserved_spend_usd: float = 0.0, +) -> BudgetDecision: + if not _session_auto_approval_enabled(session): + return BudgetDecision(allowed=True) + if _is_scheduled_hf_job_run(tool_name, tool_args): + return BudgetDecision( + allowed=False, + billable=True, + block_reason=( + "Scheduled HF jobs require disabling YOLO because their recurring " + "cost is unbounded." + ), + ) + if not _is_budgeted_auto_approval_target(tool_name, tool_args): + return BudgetDecision(allowed=True) + estimate = await estimate_tool_cost(tool_name, tool_args, session=session) + return check_session_budget( + session, + estimate, + reserved_spend_usd=reserved_spend_usd, + ) + + # -- LLM retry constants -------------------------------------------------- _MAX_LLM_RETRIES = 3 _LLM_RETRY_DELAYS = [5, 15, 30] # seconds between retries +_LLM_RATE_LIMIT_RETRY_DELAYS = [30, 60] + + +def _is_rate_limit_error(error: Exception) -> bool: + """Return True for rate-limit / quota-bucket style provider errors.""" + err_str = str(error).lower() + rate_limit_patterns = [ + "429", + "rate limit", + "rate_limit", + "too many requests", + "too many tokens", + "request limit", + "throttl", + ] + return any(pattern in err_str for pattern in rate_limit_patterns) + + +def _is_context_overflow_error(error: Exception) -> bool: + """Return True when the prompt exceeded the model's context window.""" + if isinstance(error, ContextWindowExceededError): + return True + + err_str = str(error).lower() + overflow_patterns = [ + "context window exceeded", + "maximum context length", + "max context length", + "prompt is too long", + "context length exceeded", + "too many input tokens", + "input is too long", + ] + return any(pattern in err_str for pattern in overflow_patterns) + + +def _retry_delay_for(error: Exception, attempt_index: int) -> int | None: + """Return the delay for this retry attempt, or None if it should not retry.""" + if _is_rate_limit_error(error): + schedule = _LLM_RATE_LIMIT_RETRY_DELAYS + elif _is_transient_error(error): + schedule = _LLM_RETRY_DELAYS + else: + return None + + if attempt_index >= len(schedule): + return None + return schedule[attempt_index] def _is_transient_error(error: Exception) -> bool: """Return True for errors that are likely transient and worth retrying.""" err_str = str(error).lower() transient_patterns = [ - "timeout", "timed out", - "429", "rate limit", "rate_limit", - "503", "service unavailable", - "502", "bad gateway", - "500", "internal server error", - "overloaded", "capacity", - "connection reset", "connection refused", "connection error", - "eof", "broken pipe", + "timeout", + "timed out", + "503", + "service unavailable", + "502", + "bad gateway", + "500", + "internal server error", + "overloaded", + "capacity", + "connection reset", + "connection refused", + "connection error", + "eof", + "broken pipe", ] - return any(pattern in err_str for pattern in transient_patterns) + return _is_rate_limit_error(error) or any( + pattern in err_str for pattern in transient_patterns + ) def _is_effort_config_error(error: Exception) -> bool: @@ -146,11 +563,14 @@ def _is_effort_config_error(error: Exception) -> bool: doesn't work for the current model. We heal the cache and retry once. """ from agent.core.effort_probe import _is_invalid_effort, _is_thinking_unsupported + return _is_thinking_unsupported(error) or _is_invalid_effort(error) async def _heal_effort_and_rebuild_params( - session: Session, error: Exception, llm_params: dict, + session: Session, + error: Exception, + llm_params: dict, ) -> dict: """Update the session's effort cache based on ``error`` and return new llm_params. Called only when ``_is_effort_config_error(error)`` is True. @@ -161,7 +581,11 @@ async def _heal_effort_and_rebuild_params( • invalid-effort → re-run the full cascade probe; the result lands in the cache """ - from agent.core.effort_probe import ProbeInconclusive, _is_thinking_unsupported, probe_effort + from agent.core.effort_probe import ( + ProbeInconclusive, + _is_thinking_unsupported, + probe_effort, + ) model = session.config.model_name if _is_thinking_unsupported(error): @@ -170,11 +594,16 @@ async def _heal_effort_and_rebuild_params( else: try: outcome = await probe_effort( - model, session.config.reasoning_effort, session.hf_token, + model, + session.config.reasoning_effort, + session.hf_token, + session=session, ) session.model_effective_effort[model] = outcome.effective_effort logger.info( - "healed: %s effort cascade → %s", model, outcome.effective_effort, + "healed: %s effort cascade → %s", + model, + outcome.effective_effort, ) except ProbeInconclusive: # Transient during healing — strip thinking for safety, next @@ -189,26 +618,50 @@ async def _heal_effort_and_rebuild_params( ) -def _friendly_error_message(error: Exception) -> str | None: +def _inference_credit_error_message(user_plan: str | None = None) -> str: + plan = (user_plan or "unknown").lower() + if plan == "pro": + return ( + "Hugging Face Inference Providers credits are exhausted for this " + "account.\n\n" + f"Add credits to continue: {HF_BILLING_URL}" + ) + if plan == "free": + return ( + "Your monthly Hugging Face Inference Providers credits are exhausted.\n\n" + f"Subscribe to HF PRO for more monthly usage: {HF_PRO_SUBSCRIBE_URL}\n" + f"Or add pay-as-you-go credits: {HF_BILLING_URL}" + ) + return ( + "Hugging Face Inference Providers credits appear to be exhausted for " + "this account.\n\n" + f"Add pay-as-you-go credits: {HF_BILLING_URL}\n" + f"If this is a free account, HF PRO adds more monthly usage: {HF_PRO_SUBSCRIBE_URL}" + ) + + +def _friendly_error_message( + error: Exception, + *, + user_plan: str | None = None, +) -> str | None: """Return a user-friendly message for known error types, or None to fall back to traceback.""" err_str = str(error).lower() - if "authentication" in err_str or "unauthorized" in err_str or "invalid x-api-key" in err_str: + if ( + "authentication" in err_str + or "unauthorized" in err_str + or "invalid x-api-key" in err_str + ): return ( - "Authentication failed — your API key is missing or invalid.\n\n" - "To fix this, set the API key for your model provider:\n" - " • Anthropic: export ANTHROPIC_API_KEY=sk-...\n" - " • OpenAI: export OPENAI_API_KEY=sk-...\n" - " • HF Router: export HF_TOKEN=hf_...\n\n" + "Authentication failed - your Hugging Face token is missing or invalid.\n\n" + "To fix this, set HF_TOKEN=hf_... or run `hf auth login`.\n\n" "You can also add it to a .env file in the project root.\n" "To switch models, use the /model command." ) - if "insufficient" in err_str and "credit" in err_str: - return ( - "Insufficient API credits. Please check your account balance " - "at your model provider's dashboard." - ) + if is_inference_billing_error(error): + return _inference_credit_error_message(user_plan) if "not supported by provider" in err_str or "no provider supports" in err_str: return ( @@ -219,12 +672,11 @@ def _friendly_error_message(error: Exception) -> str | None: ) if "model_not_found" in err_str or ( - "model" in err_str - and ("not found" in err_str or "does not exist" in err_str) + "model" in err_str and ("not found" in err_str or "does not exist" in err_str) ): return ( "Model not found. Use '/model' to list suggestions, or paste an " - "HF model id like 'MiniMaxAI/MiniMax-M2.7'. Availability is shown " + "HF model id like 'MiniMaxAI/MiniMax-M3:novita'. Availability is shown " "when you switch." ) @@ -232,23 +684,69 @@ def _friendly_error_message(error: Exception) -> str | None: async def _compact_and_notify(session: Session) -> None: - """Run compaction and send event if context was reduced.""" + """Run compaction and send event if context was reduced. + + Catches ``CompactionFailedError`` and ends the session cleanly instead + of letting the caller retry. Pre-2026-05-04 the caller looped on + ContextWindowExceededError → compact → re-trigger, burning hosted + inference budget while the session never reached the upload path. + """ + from agent.context_manager.manager import CompactionFailedError + cm = session.context_manager old_usage = cm.running_context_usage logger.debug( "Compaction check: usage=%d, max=%d, threshold=%d, needs_compact=%s", - old_usage, cm.model_max_tokens, cm.compaction_threshold, cm.needs_compaction, - ) - await cm.compact( - model_name=session.config.model_name, - tool_specs=session.tool_router.get_tool_specs_for_llm(), - hf_token=session.hf_token, + old_usage, + cm.model_max_tokens, + cm.compaction_threshold, + cm.needs_compaction, ) + try: + await cm.compact( + model_name=session.config.model_name, + tool_specs=session.tool_router.get_tool_specs_for_llm(), + hf_token=session.hf_token, + session=session, + ) + except CompactionFailedError as e: + logger.error( + "Compaction failed for session %s: %s — terminating session", + session.session_id, + e, + ) + # Persist the failure event so the dataset has a record of WHY this + # session ended (and the cost it incurred up to that point) even if + # save_and_upload_detached has issues downstream. + await session.send_event( + Event( + event_type="session_terminated", + data={ + "reason": "compaction_failed", + "context_usage": cm.running_context_usage, + "context_threshold": cm.compaction_threshold, + "error": str(e)[:300], + "user_message": ( + "Your conversation has grown too large to continue. " + "The work you've done is saved — start a new session to keep going." + ), + }, + ) + ) + # Stop the agent loop; the finally in _run_session will fire + # cleanup_sandbox + save_trajectory so the dataset captures + # everything that did happen. + session.is_running = False + return + new_usage = cm.running_context_usage if new_usage != old_usage: logger.warning( "Context compacted: %d -> %d tokens (max=%d, %d messages)", - old_usage, new_usage, cm.model_max_tokens, len(cm.items), + old_usage, + new_usage, + cm.model_max_tokens, + len(cm.items), ) await session.send_event( Event( @@ -287,145 +785,419 @@ async def _cleanup_on_cancel(session: Session) -> None: @dataclass class LLMResult: """Result from an LLM call (streaming or non-streaming).""" + content: str | None tool_calls_acc: dict[int, dict] token_count: int finish_reason: str | None + usage: dict = field(default_factory=dict) + + +def _session_cancelled(session: Any) -> bool: + return bool(getattr(session, "is_cancelled", False)) + + +async def _sleep_for_retry_or_cancel(session: Session, delay: float) -> bool: + """Sleep for a retry delay, waking early if the session is interrupted.""" + if _session_cancelled(session): + return True + + cancel_event = getattr(session, "_cancelled", None) + if cancel_event is None or not hasattr(cancel_event, "wait"): + await asyncio.sleep(delay) + return _session_cancelled(session) + + sleep_task = asyncio.create_task(asyncio.sleep(delay)) + cancel_task = asyncio.create_task(cancel_event.wait()) + done, pending = await asyncio.wait( + {sleep_task, cancel_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + return cancel_task in done or _session_cancelled(session) + + +def _is_invalid_thinking_signature_error(exc: Exception) -> bool: + """Return True when a provider rejected replayed thinking metadata.""" + text = str(exc) + return ( + "Invalid `signature` in `thinking` block" in text + or "Invalid signature in thinking block" in text + ) + + +def _strip_thinking_state_from_messages(messages: list[Any]) -> int: + """Remove replayed thinking metadata from assistant history messages.""" + stripped = 0 + + for message in messages: + role = ( + message.get("role") + if isinstance(message, dict) + else getattr(message, "role", None) + ) + if role != "assistant": + continue + + if isinstance(message, dict): + if message.pop("thinking_blocks", None) is not None: + stripped += 1 + if message.pop("reasoning_content", None) is not None: + stripped += 1 + provider_fields = message.get("provider_specific_fields") + content = message.get("content") + else: + if getattr(message, "thinking_blocks", None) is not None: + message.thinking_blocks = None + stripped += 1 + if getattr(message, "reasoning_content", None) is not None: + message.reasoning_content = None + stripped += 1 + provider_fields = getattr(message, "provider_specific_fields", None) + content = getattr(message, "content", None) + + if isinstance(provider_fields, dict): + cleaned_fields = dict(provider_fields) + if cleaned_fields.pop("thinking_blocks", None) is not None: + stripped += 1 + if cleaned_fields.pop("reasoning_content", None) is not None: + stripped += 1 + if cleaned_fields != provider_fields: + if isinstance(message, dict): + message["provider_specific_fields"] = cleaned_fields + else: + message.provider_specific_fields = cleaned_fields + + if isinstance(content, list): + cleaned_content = [ + block + for block in content + if not ( + isinstance(block, dict) + and block.get("type") in {"thinking", "redacted_thinking"} + ) + ] + if len(cleaned_content) != len(content): + stripped += len(content) - len(cleaned_content) + if isinstance(message, dict): + message["content"] = cleaned_content + else: + message.content = cleaned_content + + return stripped + +async def _maybe_heal_invalid_thinking_signature( + session: Session, + messages: list[Any], + exc: Exception, + *, + already_healed: bool, +) -> bool: + if already_healed or not _is_invalid_thinking_signature_error(exc): + return False + + stripped = _strip_thinking_state_from_messages(messages) + if not stripped: + return False -async def _call_llm_streaming(session: Session, messages, tools, llm_params) -> LLMResult: + await session.send_event( + Event( + event_type="tool_log", + data={ + "tool": "system", + "log": ( + "The inference provider rejected stale thinking signatures; retrying " + "without replayed thinking metadata." + ), + }, + ) + ) + return True + + +def _assistant_message_from_result( + llm_result: LLMResult, + *, + tool_calls: list[ToolCall] | None = None, +) -> Message: + """Build an assistant history message for HF Router-compatible replay.""" + kwargs: dict[str, Any] = { + "role": "assistant", + "content": llm_result.content, + } + if tool_calls is not None: + kwargs["tool_calls"] = tool_calls + return Message(**kwargs) + + +async def _call_llm_streaming( + session: Session, messages, tools, llm_params +) -> LLMResult: """Call the LLM with streaming, emitting assistant_chunk events.""" - response = None _healed_effort = False # one-shot safety net per call - messages, tools = with_prompt_caching(messages, tools, llm_params.get("model")) + _healed_thinking_signature = False + t_start = time.monotonic() for _llm_attempt in range(_MAX_LLM_RETRIES): + if _session_cancelled(session): + return LLMResult( + content=None, + tool_calls_acc={}, + token_count=0, + finish_reason=None, + ) + full_content = "" + tool_calls_acc: dict[int, dict] = {} + token_count = 0 + finish_reason = None + final_usage_chunk = None try: + request_llm_params = with_prompt_cache_params( + llm_params, + session_id=router_session_id_for(session), + ) + cached_messages, cached_tools = with_prompt_caching( + messages, tools, request_llm_params + ) response = await acompletion( - messages=messages, - tools=tools, + messages=cached_messages, + tools=cached_tools, tool_choice="auto", stream=True, stream_options={"include_usage": True}, timeout=600, - **llm_params, + **request_llm_params, + ) + + async for chunk in response: + if session.is_cancelled: + tool_calls_acc.clear() + break + + choice = chunk.choices[0] if chunk.choices else None + if not choice: + if hasattr(chunk, "usage") and chunk.usage: + token_count = chunk.usage.total_tokens + final_usage_chunk = chunk + continue + + delta = choice.delta + if choice.finish_reason: + finish_reason = choice.finish_reason + + if delta.content: + full_content += delta.content + await session.send_event( + Event( + event_type="assistant_chunk", + data={"content": delta.content}, + ) + ) + + if delta.tool_calls: + for tc_delta in delta.tool_calls: + idx = tc_delta.index + if idx not in tool_calls_acc: + tool_calls_acc[idx] = { + "id": "", + "type": "function", + "function": {"name": "", "arguments": ""}, + } + if tc_delta.id: + tool_calls_acc[idx]["id"] = tc_delta.id + if tc_delta.function: + if tc_delta.function.name: + tool_calls_acc[idx]["function"]["name"] += ( + tc_delta.function.name + ) + if tc_delta.function.arguments: + tool_calls_acc[idx]["function"]["arguments"] += ( + tc_delta.function.arguments + ) + + if hasattr(chunk, "usage") and chunk.usage: + token_count = chunk.usage.total_tokens + final_usage_chunk = chunk + + usage = await telemetry.record_llm_call( + session, + model=llm_params.get("model", session.config.model_name), + response=final_usage_chunk, + latency_ms=int((time.monotonic() - t_start) * 1000), + finish_reason=finish_reason, + ) + return LLMResult( + content=full_content or None, + tool_calls_acc=tool_calls_acc, + token_count=token_count, + finish_reason=finish_reason, + usage=usage, ) - break except ContextWindowExceededError: raise except Exception as e: + stream_received_output = bool(full_content or tool_calls_acc) + if full_content: + await session.send_event( + Event(event_type="assistant_stream_end", data={}) + ) + if stream_received_output: + logger.warning( + "Streaming LLM error after partial response; not retrying " + "to avoid duplicating assistant output/tool calls: %s", + e, + ) + await telemetry.record_llm_call( + session, + model=llm_params.get("model", session.config.model_name), + response=final_usage_chunk, + latency_ms=int((time.monotonic() - t_start) * 1000), + finish_reason=finish_reason or "error", + ) + raise + if _is_context_overflow_error(e): + raise ContextWindowExceededError(str(e)) from e if not _healed_effort and _is_effort_config_error(e): _healed_effort = True - llm_params = await _heal_effort_and_rebuild_params(session, e, llm_params) - await session.send_event(Event( - event_type="tool_log", - data={"tool": "system", "log": "Reasoning effort not supported for this model — adjusting and retrying."}, - )) + llm_params = await _heal_effort_and_rebuild_params( + session, e, llm_params + ) + await session.send_event( + Event( + event_type="tool_log", + data={ + "tool": "system", + "log": "Reasoning effort not supported for this model — adjusting and retrying.", + }, + ) + ) + continue + if await _maybe_heal_invalid_thinking_signature( + session, + messages, + e, + already_healed=_healed_thinking_signature, + ): + _healed_thinking_signature = True continue - if _llm_attempt < _MAX_LLM_RETRIES - 1 and _is_transient_error(e): - _delay = _LLM_RETRY_DELAYS[_llm_attempt] + _delay = _retry_delay_for(e, _llm_attempt) + if _llm_attempt < _MAX_LLM_RETRIES - 1 and _delay is not None: logger.warning( "Transient LLM error (attempt %d/%d): %s — retrying in %ds", - _llm_attempt + 1, _MAX_LLM_RETRIES, e, _delay, + _llm_attempt + 1, + _MAX_LLM_RETRIES, + e, + _delay, ) - await session.send_event(Event( - event_type="tool_log", - data={"tool": "system", "log": f"LLM connection error, retrying in {_delay}s..."}, - )) - await asyncio.sleep(_delay) + await session.send_event( + Event( + event_type="tool_log", + data={ + "tool": "system", + "log": f"LLM connection error, retrying in {_delay}s...", + }, + ) + ) + if await _sleep_for_retry_or_cancel(session, _delay): + return LLMResult( + content=None, + tool_calls_acc={}, + token_count=0, + finish_reason=None, + ) continue raise - full_content = "" - tool_calls_acc: dict[int, dict] = {} - token_count = 0 - finish_reason = None - - async for chunk in response: - if session.is_cancelled: - tool_calls_acc.clear() - break - - choice = chunk.choices[0] if chunk.choices else None - if not choice: - if hasattr(chunk, "usage") and chunk.usage: - token_count = chunk.usage.total_tokens - continue - - delta = choice.delta - if choice.finish_reason: - finish_reason = choice.finish_reason - if delta.content: - full_content += delta.content - await session.send_event( - Event(event_type="assistant_chunk", data={"content": delta.content}) - ) - - if delta.tool_calls: - for tc_delta in delta.tool_calls: - idx = tc_delta.index - if idx not in tool_calls_acc: - tool_calls_acc[idx] = { - "id": "", "type": "function", - "function": {"name": "", "arguments": ""}, - } - if tc_delta.id: - tool_calls_acc[idx]["id"] = tc_delta.id - if tc_delta.function: - if tc_delta.function.name: - tool_calls_acc[idx]["function"]["name"] += tc_delta.function.name - if tc_delta.function.arguments: - tool_calls_acc[idx]["function"]["arguments"] += tc_delta.function.arguments - - if hasattr(chunk, "usage") and chunk.usage: - token_count = chunk.usage.total_tokens - - return LLMResult( - content=full_content or None, - tool_calls_acc=tool_calls_acc, - token_count=token_count, - finish_reason=finish_reason, - ) - - -async def _call_llm_non_streaming(session: Session, messages, tools, llm_params) -> LLMResult: +async def _call_llm_non_streaming( + session: Session, messages, tools, llm_params +) -> LLMResult: """Call the LLM without streaming, emit assistant_message at the end.""" response = None _healed_effort = False - messages, tools = with_prompt_caching(messages, tools, llm_params.get("model")) + _healed_thinking_signature = False + t_start = time.monotonic() for _llm_attempt in range(_MAX_LLM_RETRIES): + if _session_cancelled(session): + return LLMResult( + content=None, + tool_calls_acc={}, + token_count=0, + finish_reason=None, + ) try: + request_llm_params = with_prompt_cache_params( + llm_params, + session_id=router_session_id_for(session), + ) + cached_messages, cached_tools = with_prompt_caching( + messages, tools, request_llm_params + ) response = await acompletion( - messages=messages, - tools=tools, + messages=cached_messages, + tools=cached_tools, tool_choice="auto", stream=False, timeout=600, - **llm_params, + **request_llm_params, ) break except ContextWindowExceededError: raise except Exception as e: + if _is_context_overflow_error(e): + raise ContextWindowExceededError(str(e)) from e if not _healed_effort and _is_effort_config_error(e): _healed_effort = True - llm_params = await _heal_effort_and_rebuild_params(session, e, llm_params) - await session.send_event(Event( - event_type="tool_log", - data={"tool": "system", "log": "Reasoning effort not supported for this model — adjusting and retrying."}, - )) + llm_params = await _heal_effort_and_rebuild_params( + session, e, llm_params + ) + await session.send_event( + Event( + event_type="tool_log", + data={ + "tool": "system", + "log": "Reasoning effort not supported for this model — adjusting and retrying.", + }, + ) + ) continue - if _llm_attempt < _MAX_LLM_RETRIES - 1 and _is_transient_error(e): - _delay = _LLM_RETRY_DELAYS[_llm_attempt] + if await _maybe_heal_invalid_thinking_signature( + session, + messages, + e, + already_healed=_healed_thinking_signature, + ): + _healed_thinking_signature = True + continue + _delay = _retry_delay_for(e, _llm_attempt) + if _llm_attempt < _MAX_LLM_RETRIES - 1 and _delay is not None: logger.warning( "Transient LLM error (attempt %d/%d): %s — retrying in %ds", - _llm_attempt + 1, _MAX_LLM_RETRIES, e, _delay, + _llm_attempt + 1, + _MAX_LLM_RETRIES, + e, + _delay, + ) + await session.send_event( + Event( + event_type="tool_log", + data={ + "tool": "system", + "log": f"LLM connection error, retrying in {_delay}s...", + }, + ) ) - await session.send_event(Event( - event_type="tool_log", - data={"tool": "system", "log": f"LLM connection error, retrying in {_delay}s..."}, - )) - await asyncio.sleep(_delay) + if await _sleep_for_retry_or_cancel(session, _delay): + return LLMResult( + content=None, + tool_calls_acc={}, + token_count=0, + finish_reason=None, + ) continue raise @@ -454,11 +1226,20 @@ async def _call_llm_non_streaming(session: Session, messages, tools, llm_params) Event(event_type="assistant_message", data={"content": content}) ) + usage = await telemetry.record_llm_call( + session, + model=llm_params.get("model", session.config.model_name), + response=response, + latency_ms=int((time.monotonic() - t_start) * 1000), + finish_reason=finish_reason, + ) + return LLMResult( content=content, tool_calls_acc=tool_calls_acc, token_count=token_count, finish_reason=finish_reason, + usage=usage, ) @@ -473,6 +1254,44 @@ async def _abandon_pending_approval(session: Session) -> None: history stays valid) and notifies the frontend that those tools were abandoned. """ + if is_usage_threshold_pending( + session.pending_approval + ) or is_yolo_budget_pending(session.pending_approval): + pending = session.pending_approval + tool_call_id = str(pending.get("tool_call_id") or "") + tool_name = str(pending.get("kind") or USAGE_THRESHOLD_TOOL_NAME) + session.pending_approval = None + if tool_call_id: + await session.send_event( + Event( + event_type="tool_state_change", + data={ + "tool_call_id": tool_call_id, + "tool": tool_name, + "state": "abandoned", + }, + ) + ) + if pending.get("continuation") == "complete_turn": + final_response = pending.get("final_response") + await session.send_event( + Event( + event_type="turn_complete", + data={ + "history_size": int( + pending.get("history_size") + or len(session.context_manager.items) + ), + "final_response": final_response + if isinstance(final_response, str) + else None, + }, + ) + ) + session.increment_turn() + await session.auto_save_if_needed() + return + tool_calls = session.pending_approval.get("tool_calls", []) for tc in tool_calls: tool_name = tc.function.name @@ -505,7 +1324,8 @@ async def _abandon_pending_approval(session: Session) -> None: @staticmethod async def run_agent( - session: Session, text: str, + session: Session, + text: str, ) -> str | None: """ Handle user input (like user_input_or_turn in codex.rs:1291) @@ -534,14 +1354,32 @@ async def run_agent( final_response = None errored = False max_iterations = session.config.max_iterations + no_tool_incomplete_plan_retries = 0 while max_iterations == -1 or iteration < max_iterations: # ── Cancellation check: before LLM call ── if session.is_cancelled: break - - # Compact before calling the LLM if context is near the limit + if session.pending_approval: + return final_response + + # Compact before calling the LLM if context is near the limit. + # When _compact_and_notify catches CompactionFailedError it sets + # session.is_running = False; we MUST exit the loop here, otherwise + # the LLM call below fires with an over-threshold context, hits + # ContextWindowExceededError, and we end up looping again on the + # except path — exactly the bug this PR is supposed to fix. await _compact_and_notify(session) + if not session.is_running: + break + if session.pending_approval: + return final_response + + if await _maybe_pause_for_usage_threshold( + session, + continuation="continue_agent", + ): + return final_response # Doom-loop detection: break out of repeated tool call patterns doom_prompt = check_for_doom_loop(session.context_manager.items) @@ -549,12 +1387,28 @@ async def run_agent( session.context_manager.add_message( Message(role="user", content=doom_prompt) ) + + malformed_tool = _detect_repeated_malformed(session.context_manager.items) + if malformed_tool: + recovery_prompt = ( + "[SYSTEM: Repeated malformed tool arguments detected for " + f"'{malformed_tool}'. Stop retrying the same tool call shape. " + "Use a different strategy that produces smaller, valid JSON. " + "For large file writes, prefer bash with a heredoc or split the " + "edit into multiple smaller tool calls.]" + ) + session.context_manager.add_message( + Message(role="user", content=recovery_prompt) + ) await session.send_event( Event( event_type="tool_log", data={ "tool": "system", - "log": "Doom loop detected — injecting corrective prompt", + "log": ( + "Repeated malformed tool arguments detected — " + f"forcing a different strategy for {malformed_tool}" + ), }, ) ) @@ -569,12 +1423,19 @@ async def run_agent( llm_params = _resolve_llm_params( session.config.model_name, session.hf_token, - reasoning_effort=session.effective_effort_for(session.config.model_name), + reasoning_effort=session.effective_effort_for( + session.config.model_name + ), ) if session.stream: - llm_result = await _call_llm_streaming(session, messages, tools, llm_params) + llm_result = await _call_llm_streaming( + session, messages, tools, llm_params + ) else: - llm_result = await _call_llm_non_streaming(session, messages, tools, llm_params) + llm_result = await _call_llm_non_streaming( + session, messages, tools, llm_params + ) + llm_observed_cost_usd = llm_result.usage.get("cost_usd") content = llm_result.content tool_calls_acc = llm_result.tool_calls_acc @@ -606,7 +1467,7 @@ async def run_agent( " • For other tools: reduce the size of your arguments or use bash." ) if content: - assistant_msg = Message(role="assistant", content=content) + assistant_msg = _assistant_message_from_result(llm_result) session.context_manager.add_message(assistant_msg, token_count) session.context_manager.add_message( Message(role="user", content=f"[SYSTEM: {truncation_hint}]") @@ -618,7 +1479,10 @@ async def run_agent( await session.send_event( Event( event_type="tool_log", - data={"tool": "system", "log": f"Output truncated — retrying with smaller content ({dropped_names})"}, + data={ + "tool": "system", + "log": f"Output truncated — retrying with smaller content ({dropped_names})", + }, ) ) iteration += 1 @@ -647,6 +1511,54 @@ async def run_agent( # If no tool calls, add assistant message and we're done if not tool_calls: + unfinished_plan = _unfinished_plan_items(session) + if ( + unfinished_plan + and no_tool_incomplete_plan_retries + < _NO_TOOL_INCOMPLETE_PLAN_RETRY_LIMIT + ): + if await maybe_pause_yolo_after_spend( + session, + spend_kind="llm_call", + observed_cost_usd=llm_observed_cost_usd, + ): + return final_response + logger.info( + "No tool calls with unfinished plan; retrying agent turn " + "(attempt %d/%d)", + no_tool_incomplete_plan_retries + 1, + _NO_TOOL_INCOMPLETE_PLAN_RETRY_LIMIT, + ) + if content: + assistant_msg = _assistant_message_from_result(llm_result) + session.context_manager.add_message( + assistant_msg, token_count + ) + session.context_manager.add_message( + Message( + role="user", + content=_no_tool_incomplete_plan_prompt( + unfinished_plan + ), + ) + ) + no_tool_incomplete_plan_retries += 1 + await session.send_event( + Event( + event_type="tool_log", + data={ + "tool": "system", + "log": ( + "Plan still has unfinished items after a " + "text-only response — retrying instead of " + "returning to the prompt." + ), + }, + ) + ) + iteration += 1 + continue + logger.debug( "Agent loop ending: no tool calls. " "finish_reason=%s, token_count=%d, " @@ -662,11 +1574,30 @@ async def run_agent( (content or "")[:500], ) if content: - assistant_msg = Message(role="assistant", content=content) + assistant_msg = _assistant_message_from_result(llm_result) session.context_manager.add_message(assistant_msg, token_count) final_response = content + if await maybe_pause_yolo_after_spend( + session, + spend_kind="llm_call", + observed_cost_usd=llm_observed_cost_usd, + continuation="complete_turn", + final_response=final_response + if isinstance(final_response, str) + else None, + ): + return final_response break + no_tool_incomplete_plan_retries = 0 + + if await maybe_pause_yolo_after_spend( + session, + spend_kind="llm_call", + observed_cost_usd=llm_observed_cost_usd, + ): + return final_response + # Validate tool call args (one json.loads per call, once) # and split into good vs bad good_tools: list[tuple[ToolCall, str, dict]] = [] @@ -678,15 +1609,15 @@ async def run_agent( except (json.JSONDecodeError, TypeError, ValueError): logger.warning( "Malformed arguments for tool_call %s (%s) — skipping", - tc.id, tc.function.name, + tc.id, + tc.function.name, ) tc.function.arguments = "{}" bad_tools.append(tc) # Add assistant message with all tool calls to context - assistant_msg = Message( - role="assistant", - content=content, + assistant_msg = _assistant_message_from_result( + llm_result, tool_calls=tool_calls, ) session.context_manager.add_message(assistant_msg, token_count) @@ -699,48 +1630,92 @@ async def run_agent( f"arguments and was NOT executed. Retry with smaller content — " f"for 'write', split into multiple smaller writes using 'edit'." ) - session.context_manager.add_message(Message( - role="tool", - content=error_msg, - tool_call_id=tc.id, - name=tc.function.name, - )) - await session.send_event(Event( - event_type="tool_call", - data={"tool": tc.function.name, "arguments": {}, "tool_call_id": tc.id}, - )) - await session.send_event(Event( - event_type="tool_output", - data={"tool": tc.function.name, "tool_call_id": tc.id, "output": error_msg, "success": False}, - )) + session.context_manager.add_message( + Message( + role="tool", + content=error_msg, + tool_call_id=tc.id, + name=tc.function.name, + ) + ) + await session.send_event( + Event( + event_type="tool_call", + data={ + "tool": tc.function.name, + "arguments": {}, + "tool_call_id": tc.id, + }, + ) + ) + await session.send_event( + Event( + event_type="tool_output", + data={ + "tool": tc.function.name, + "tool_call_id": tc.id, + "output": error_msg, + "success": False, + }, + ) + ) # ── Cancellation check: before tool execution ── if session.is_cancelled: break - # Separate good tools into approval-required vs auto-execute - approval_required_tools: list[tuple[ToolCall, str, dict]] = [] - non_approval_tools: list[tuple[ToolCall, str, dict]] = [] + # Separate good tools into approval-required vs auto-execute. + # Track reserved spend while classifying a batch so two + # auto-approved jobs in one model response cannot jointly + # exceed the remaining session cap. + approval_required_tools: list[ + tuple[ToolCall, str, dict, ApprovalDecision] + ] = [] + non_approval_tools: list[ + tuple[ToolCall, str, dict, ApprovalDecision] + ] = [] + reserved_auto_spend_usd = 0.0 for tc, tool_name, tool_args in good_tools: - if _needs_approval(tool_name, tool_args, session.config): - approval_required_tools.append((tc, tool_name, tool_args)) + decision = await _approval_decision( + tool_name, + tool_args, + session, + reserved_spend_usd=reserved_auto_spend_usd, + ) + if decision.requires_approval: + approval_required_tools.append( + (tc, tool_name, tool_args, decision) + ) else: - non_approval_tools.append((tc, tool_name, tool_args)) + non_approval_tools.append((tc, tool_name, tool_args, decision)) + if ( + decision.auto_approved + and decision.billable + and decision.estimated_cost_usd is not None + ): + reserved_auto_spend_usd += decision.estimated_cost_usd # Execute non-approval tools (in parallel when possible) if non_approval_tools: # 1. Validate args upfront parsed_tools: list[ - tuple[ToolCall, str, dict, bool, str] + tuple[ToolCall, str, dict, ApprovalDecision, bool, str] ] = [] - for tc, tool_name, tool_args in non_approval_tools: + for tc, tool_name, tool_args, decision in non_approval_tools: args_valid, error_msg = _validate_tool_args(tool_args) parsed_tools.append( - (tc, tool_name, tool_args, args_valid, error_msg) + (tc, tool_name, tool_args, decision, args_valid, error_msg) ) # 2. Send all tool_call events upfront (so frontend shows them all) - for tc, tool_name, tool_args, args_valid, _ in parsed_tools: + for ( + tc, + tool_name, + tool_args, + _decision, + args_valid, + _, + ) in parsed_tools: if args_valid: await session.send_event( Event( @@ -758,22 +1733,42 @@ async def _exec_tool( tc: ToolCall, name: str, args: dict, + decision: ApprovalDecision, valid: bool, err: str, ) -> tuple[ToolCall, str, dict, str, bool]: if not valid: return (tc, name, args, err, False) + if decision.billable: + budget = _record_estimated_spend( + session, + decision, + reservation_id=tc.id, + ) + if not budget.allowed: + return ( + tc, + name, + args, + budget.block_reason + or "YOLO budget blocked this tool call.", + False, + ) out, ok = await session.tool_router.call_tool( name, args, session=session, tool_call_id=tc.id ) + if not ok and decision.billable: + release_budget_reservation(session, tc.id) return (tc, name, args, out, ok) - gather_task = asyncio.ensure_future(asyncio.gather( - *[ - _exec_tool(tc, name, args, valid, err) - for tc, name, args, valid, err in parsed_tools - ] - )) + gather_task = asyncio.ensure_future( + asyncio.gather( + *[ + _exec_tool(tc, name, args, decision, valid, err) + for tc, name, args, decision, valid, err in parsed_tools + ] + ) + ) cancel_task = asyncio.ensure_future(session._cancelled.wait()) done, _ = await asyncio.wait( @@ -788,12 +1783,18 @@ async def _exec_tool( except asyncio.CancelledError: pass # Notify frontend that in-flight tools were cancelled - for tc, name, _args, valid, _ in parsed_tools: + for tc, name, _args, _decision, valid, _ in parsed_tools: if valid: - await session.send_event(Event( - event_type="tool_state_change", - data={"tool_call_id": tc.id, "tool": name, "state": "cancelled"}, - )) + await session.send_event( + Event( + event_type="tool_state_change", + data={ + "tool_call_id": tc.id, + "tool": name, + "state": "cancelled", + }, + ) + ) await _cleanup_on_cancel(session) break @@ -826,30 +1827,60 @@ async def _exec_tool( if approval_required_tools: # Prepare batch approval data tools_data = [] - for tc, tool_name, tool_args in approval_required_tools: + blocked_payloads = [] + for tc, tool_name, tool_args, decision in approval_required_tools: # Resolve sandbox file paths for hf_jobs scripts so the # frontend can display & edit the actual file content. - if tool_name == "hf_jobs" and isinstance(tool_args.get("script"), str): + if tool_name == "hf_jobs" and isinstance( + tool_args.get("script"), str + ): from agent.tools.sandbox_tool import resolve_sandbox_script + sandbox = getattr(session, "sandbox", None) - resolved, _ = await resolve_sandbox_script(sandbox, tool_args["script"]) + resolved, _ = await resolve_sandbox_script( + sandbox, tool_args["script"] + ) if resolved: tool_args = {**tool_args, "script": resolved} - tools_data.append({ + tool_payload = { "tool": tool_name, "arguments": tool_args, "tool_call_id": tc.id, - }) - - await session.send_event(Event( - event_type="approval_required", - data={"tools": tools_data, "count": len(tools_data)}, - )) + } + if decision.auto_approval_blocked: + tool_payload.update( + { + "auto_approval_blocked": True, + "block_reason": decision.block_reason, + "estimated_cost_usd": decision.estimated_cost_usd, + "remaining_cap_usd": decision.remaining_cap_usd, + } + ) + blocked_payloads.append(tool_payload) + tools_data.append(tool_payload) + + event_data = {"tools": tools_data, "count": len(tools_data)} + if blocked_payloads: + first = blocked_payloads[0] + event_data.update( + { + "auto_approval_blocked": True, + "block_reason": first.get("block_reason"), + "estimated_cost_usd": first.get("estimated_cost_usd"), + "remaining_cap_usd": first.get("remaining_cap_usd"), + } + ) + await session.send_event( + Event( + event_type="approval_required", + data=event_data, + ) + ) # Store all approval-requiring tools (ToolCall objects for execution) session.pending_approval = { - "tool_calls": [tc for tc, _, _ in approval_required_tools], + "tool_calls": [tc for tc, _, _, _ in approval_required_tools], } # Return early - wait for EXEC_APPROVAL operation @@ -858,21 +1889,33 @@ async def _exec_tool( iteration += 1 except ContextWindowExceededError: - # Force compact and retry this iteration + # Force compact and retry this iteration. cm = session.context_manager logger.warning( "ContextWindowExceededError at iteration %d — forcing compaction " "(usage=%d, model_max_tokens=%d, messages=%d)", - iteration, cm.running_context_usage, cm.model_max_tokens, len(cm.items), + iteration, + cm.running_context_usage, + cm.model_max_tokens, + len(cm.items), ) cm.running_context_usage = cm.model_max_tokens + 1 await _compact_and_notify(session) + # Same guard as the top of the loop: if compaction couldn't + # bring us under threshold, _compact_and_notify has already + # emitted session_terminated and set is_running=False. Continue + # would just re-call the LLM with the same too-big context. + if not session.is_running: + break continue except Exception as e: import traceback - error_msg = _friendly_error_message(e) + error_msg = _friendly_error_message( + e, + user_plan=getattr(session, "user_plan", None), + ) if error_msg is None: error_msg = str(e) + "\n" + traceback.format_exc() @@ -889,10 +1932,23 @@ async def _exec_tool( await _cleanup_on_cancel(session) await session.send_event(Event(event_type="interrupted")) elif not errored: + if await _maybe_pause_for_usage_threshold( + session, + continuation="complete_turn", + final_response=final_response + if isinstance(final_response, str) + else None, + ): + return final_response await session.send_event( Event( event_type="turn_complete", - data={"history_size": len(session.context_manager.items)}, + data={ + "history_size": len(session.context_manager.items), + "final_response": final_response + if isinstance(final_response, str) + else None, + }, ) ) @@ -910,6 +1966,271 @@ async def undo(session: Session) -> None: logger.warning("Undo: no user message found to remove") await session.send_event(Event(event_type="undo_complete")) + @staticmethod + async def new_conversation(session: Session, *, clear_screen: bool = False) -> None: + """Start a fresh conversation inside the active runtime.""" + try: + result = session.start_new_conversation() + except Exception as e: + await session.send_event( + Event(event_type="error", data={"error": f"New chat failed: {e}"}) + ) + return + result["clear_screen"] = clear_screen + await session.send_event(Event(event_type="new_complete", data=result)) + + @staticmethod + async def resume(session: Session, path: str) -> None: + """Reload context from a saved session log into the active session.""" + from agent.core.session_resume import restore_session_from_log + + try: + result = restore_session_from_log(session, Path(path)) + except Exception as e: + await session.send_event( + Event(event_type="error", data={"error": f"Resume failed: {e}"}) + ) + return + await session.send_event(Event(event_type="resume_complete", data=result)) + + @staticmethod + async def _exec_usage_threshold_approval( + session: Session, approvals: list[dict] + ) -> None: + pending = ( + session.pending_approval + if isinstance(session.pending_approval, dict) + else {} + ) + tool_call_id = str(pending.get("tool_call_id") or "") + approval = next( + (item for item in approvals if item.get("tool_call_id") == tool_call_id), + {"approved": False}, + ) + approved = bool(approval.get("approved")) + + session.pending_approval = None + if not tool_call_id: + await session.send_event( + Event( + event_type="error", + data={"error": "Usage approval is missing its approval id"}, + ) + ) + return + + if not approved: + feedback = str(approval.get("feedback") or "Stopped by user").strip() + await session.send_event( + Event( + event_type="tool_state_change", + data={ + "tool_call_id": tool_call_id, + "tool": USAGE_THRESHOLD_TOOL_NAME, + "state": "rejected", + }, + ) + ) + await session.send_event( + Event( + event_type="tool_output", + data={ + "tool": USAGE_THRESHOLD_TOOL_NAME, + "tool_call_id": tool_call_id, + "output": feedback, + "success": False, + }, + ) + ) + await session.send_event(Event(event_type="interrupted")) + session.increment_turn() + await session.auto_save_if_needed() + return + + current_spend = _coerce_float(pending.get("current_spend_usd")) + acknowledged_threshold = _coerce_float(pending.get("threshold_usd")) + next_threshold = next_usage_warning_threshold( + current_spend, + acknowledged_threshold, + ) + session.usage_warning_next_threshold_usd = next_threshold + pending["next_threshold_usd"] = next_threshold + + await session.send_event( + Event( + event_type="tool_state_change", + data={ + "tool_call_id": tool_call_id, + "tool": USAGE_THRESHOLD_TOOL_NAME, + "state": "approved", + }, + ) + ) + await session.send_event( + Event( + event_type="tool_output", + data={ + "tool": USAGE_THRESHOLD_TOOL_NAME, + "tool_call_id": tool_call_id, + "output": _usage_output_message(pending), + "success": True, + }, + ) + ) + + if pending.get("continuation") == "complete_turn": + final_response = pending.get("final_response") + await session.send_event( + Event( + event_type="turn_complete", + data={ + "history_size": int( + pending.get("history_size") + or len(session.context_manager.items) + ), + "final_response": final_response + if isinstance(final_response, str) + else None, + }, + ) + ) + session.increment_turn() + await session.auto_save_if_needed() + return + + await Handlers.run_agent(session, "") + + @staticmethod + async def _exec_yolo_budget_approval( + session: Session, approvals: list[dict] + ) -> None: + pending = ( + session.pending_approval + if isinstance(session.pending_approval, dict) + else {} + ) + tool_call_id = str(pending.get("tool_call_id") or "") + approval = next( + (item for item in approvals if item.get("tool_call_id") == tool_call_id), + {"approved": False}, + ) + approved = bool(approval.get("approved")) + + if not tool_call_id: + session.pending_approval = None + await session.send_event( + Event( + event_type="error", + data={"error": "YOLO budget approval is missing its approval id"}, + ) + ) + return + + if not approved: + session.pending_approval = None + feedback = str(approval.get("feedback") or "Stopped by user").strip() + await session.send_event( + Event( + event_type="tool_state_change", + data={ + "tool_call_id": tool_call_id, + "tool": "yolo_budget", + "state": "rejected", + }, + ) + ) + await session.send_event( + Event( + event_type="tool_output", + data={ + "tool": "yolo_budget", + "tool_call_id": tool_call_id, + "output": feedback, + "success": False, + }, + ) + ) + await session.send_event(Event(event_type="interrupted")) + session.increment_turn() + await session.auto_save_if_needed() + return + + can_resume, reason = yolo_budget_can_resume(session, pending) + if not can_resume: + pending["reason"] = reason + pending["current_spend_usd"] = round( + float( + getattr(session, "auto_approval_estimated_spend_usd", 0.0) or 0.0 + ), + 6, + ) + pending["remaining_cap_usd"] = ( + None + if getattr(session, "auto_approval_cost_cap_usd", None) is None + else session.auto_approval_remaining_usd + ) + tool = yolo_budget_pending_to_tool(pending) + await session.send_event( + Event( + event_type="approval_required", + data={ + "tools": [tool], + "count": 1, + "yolo_budget": True, + "auto_approval_blocked": True, + "block_reason": reason, + "estimated_cost_usd": pending.get("estimated_next_usd"), + "remaining_cap_usd": pending.get("remaining_cap_usd"), + }, + ) + ) + return + + session.pending_approval = None + await session.send_event( + Event( + event_type="tool_state_change", + data={ + "tool_call_id": tool_call_id, + "tool": "yolo_budget", + "state": "approved", + }, + ) + ) + await session.send_event( + Event( + event_type="tool_output", + data={ + "tool": "yolo_budget", + "tool_call_id": tool_call_id, + "output": "YOLO budget check acknowledged.", + "success": True, + }, + ) + ) + + if pending.get("continuation") == "complete_turn": + final_response = pending.get("final_response") + await session.send_event( + Event( + event_type="turn_complete", + data={ + "history_size": int( + pending.get("history_size") + or len(session.context_manager.items) + ), + "final_response": final_response + if isinstance(final_response, str) + else None, + }, + ) + ) + session.increment_turn() + await session.auto_save_if_needed() + return + + await Handlers.run_agent(session, "") + @staticmethod async def exec_approval(session: Session, approvals: list[dict]) -> None: """Handle batch job execution approval""" @@ -922,6 +2243,13 @@ async def exec_approval(session: Session, approvals: list[dict]) -> None: ) return + if is_usage_threshold_pending(session.pending_approval): + await Handlers._exec_usage_threshold_approval(session, approvals) + return + if is_yolo_budget_pending(session.pending_approval): + await Handlers._exec_yolo_budget_approval(session, approvals) + return + tool_calls = session.pending_approval.get("tool_calls", []) if not tool_calls: await session.send_event( @@ -980,10 +2308,66 @@ async def exec_approval(session: Session, approvals: list[dict]) -> None: tool_args["script"] = edited_script was_edited = True logger.info(f"Using user-edited script for {tool_name} ({tc.id})") + selected_namespace = approval_decision.get("namespace") + if selected_namespace and tool_name == "hf_jobs": + tool_args["namespace"] = selected_namespace approved_tasks.append((tc, tool_name, tool_args, was_edited)) else: rejected_tasks.append((tc, tool_name, approval_decision)) + reserved_manual_spend_usd = 0.0 + blocked_manual_budget: tuple[ToolCall, str, BudgetDecision] | None = None + for tc, tool_name, tool_args, _was_edited in approved_tasks: + budget = await _check_manual_approved_budget( + session, + tool_name, + tool_args, + reserved_spend_usd=reserved_manual_spend_usd, + ) + if not budget.allowed: + blocked_manual_budget = (tc, tool_name, budget) + break + if budget.billable and budget.estimated_cost_usd is not None: + reserved_manual_spend_usd += budget.estimated_cost_usd + + if blocked_manual_budget is not None: + blocked_tc, _blocked_tool, blocked_budget = blocked_manual_budget + tools_data = [] + for tc in tool_calls: + try: + args = json.loads(tc.function.arguments) + except (json.JSONDecodeError, AttributeError, TypeError): + args = {} + payload = { + "tool": getattr(tc.function, "name", None), + "arguments": args, + "tool_call_id": tc.id, + } + if tc.id == blocked_tc.id: + payload.update( + { + "auto_approval_blocked": True, + "block_reason": blocked_budget.block_reason, + "estimated_cost_usd": blocked_budget.estimated_cost_usd, + "remaining_cap_usd": blocked_budget.remaining_cap_usd, + } + ) + tools_data.append(payload) + await session.send_event( + Event( + event_type="approval_required", + data={ + "tools": tools_data, + "count": len(tools_data), + "auto_approval_blocked": True, + "block_reason": blocked_budget.block_reason, + "estimated_cost_usd": blocked_budget.estimated_cost_usd, + "remaining_cap_usd": blocked_budget.remaining_cap_usd, + }, + ) + ) + return + # Clear pending approval immediately so a page refresh during # execution won't re-show the approval dialog. session.pending_approval = None @@ -1031,21 +2415,40 @@ async def execute_tool(tc, tool_name, tool_args, was_edited): ) ) + budget = await _record_manual_approved_spend_if_needed( + session, + tool_name, + tool_args, + tool_call_id=tc.id, + ) + if not budget.allowed: + return ( + tc, + tool_name, + budget.block_reason or "YOLO budget blocked this tool call.", + False, + was_edited, + ) + output, success = await session.tool_router.call_tool( tool_name, tool_args, session=session, tool_call_id=tc.id ) + if not success and budget.reservation: + release_budget_reservation(session, budget.reservation.reservation_id) return (tc, tool_name, output, success, was_edited) # Execute all approved tools concurrently (cancellable) if approved_tasks: - gather_task = asyncio.ensure_future(asyncio.gather( - *[ - execute_tool(tc, tool_name, tool_args, was_edited) - for tc, tool_name, tool_args, was_edited in approved_tasks - ], - return_exceptions=True, - )) + gather_task = asyncio.ensure_future( + asyncio.gather( + *[ + execute_tool(tc, tool_name, tool_args, was_edited) + for tc, tool_name, tool_args, was_edited in approved_tasks + ], + return_exceptions=True, + ) + ) cancel_task = asyncio.ensure_future(session._cancelled.wait()) done, _ = await asyncio.wait( @@ -1061,10 +2464,16 @@ async def execute_tool(tc, tool_name, tool_args, was_edited): pass # Notify frontend that approved tools were cancelled for tc, tool_name, _args, _was_edited in approved_tasks: - await session.send_event(Event( - event_type="tool_state_change", - data={"tool_call_id": tc.id, "tool": tool_name, "state": "cancelled"}, - )) + await session.send_event( + Event( + event_type="tool_state_change", + data={ + "tool_call_id": tc.id, + "tool": tool_name, + "state": "cancelled", + }, + ) + ) await _cleanup_on_cancel(session) await session.send_event(Event(event_type="interrupted")) session.increment_turn() @@ -1156,6 +2565,8 @@ async def shutdown(session: Session) -> bool: _ = session.save_and_upload_detached(repo_id) session.is_running = False + if not getattr(session, "local_mode", False): + await teardown_session_sandbox(session) await session.send_event(Event(event_type="shutdown")) return True @@ -1183,6 +2594,21 @@ async def process_submission(session: Session, submission) -> bool: await Handlers.undo(session) return True + if op.op_type == OpType.NEW: + clear_screen = bool((op.data or {}).get("clear_screen")) + await Handlers.new_conversation(session, clear_screen=clear_screen) + return True + + if op.op_type == OpType.RESUME: + path = op.data.get("path") if op.data else None + if path: + await Handlers.resume(session, path) + else: + await session.send_event( + Event(event_type="error", data={"error": "Resume requires a path"}) + ) + return True + if op.op_type == OpType.EXEC_APPROVAL: approvals = op.data.get("approvals", []) if op.data else [] await Handlers.exec_approval(session, approvals) @@ -1198,12 +2624,19 @@ async def process_submission(session: Session, submission) -> bool: async def submission_loop( submission_queue: asyncio.Queue, event_queue: asyncio.Queue, - config: Config | None = None, + config: Config, tool_router: ToolRouter | None = None, session_holder: list | None = None, hf_token: str | None = None, + user_id: str | None = None, + hf_username: str | None = None, local_mode: bool = False, + autonomous_mode: bool = False, stream: bool = True, + notification_gateway: NotificationGateway | None = None, + notification_destinations: list[str] | None = None, + defer_turn_complete_notification: bool = False, + user_plan: str | None = None, ) -> None: """ Main agent loop - processes submissions and dispatches to handlers. @@ -1212,17 +2645,34 @@ async def submission_loop( # Create session with tool router session = Session( - event_queue, config=config, tool_router=tool_router, hf_token=hf_token, - local_mode=local_mode, stream=stream, + event_queue, + config=config, + tool_router=tool_router, + hf_token=hf_token, + user_id=user_id, + hf_username=hf_username, + user_plan=user_plan, + local_mode=local_mode, + autonomous_mode=autonomous_mode, + stream=stream, + notification_gateway=notification_gateway, + notification_destinations=notification_destinations, + defer_turn_complete_notification=defer_turn_complete_notification, ) if session_holder is not None: session_holder[0] = session + if not local_mode: + start_cpu_sandbox_preload(session) logger.info("Agent loop started") - # Retry any failed uploads from previous sessions (fire-and-forget) + # Retry any failed uploads from previous sessions (fire-and-forget). + # Includes the personal trace repo when enabled so a session that failed + # to publish to the user's HF dataset gets a fresh attempt on next run. if config and config.save_sessions: Session.retry_failed_uploads_detached( - directory="session_logs", repo_id=config.session_dataset_repo + directory=str(DEFAULT_SESSION_LOG_DIR), + repo_id=config.session_dataset_repo, + personal_repo_id=session._personal_trace_repo_id(), ) try: @@ -1230,10 +2680,13 @@ async def submission_loop( async with tool_router: # Emit ready event after initialization await session.send_event( - Event(event_type="ready", data={ - "message": "Agent initialized", - "tool_count": len(tool_router.tools), - }) + Event( + event_type="ready", + data={ + "message": "Agent initialized", + "tool_count": len(tool_router.tools), + }, + ) ) while session.is_running: diff --git a/agent/core/approval_policy.py b/agent/core/approval_policy.py new file mode 100644 index 000000000..73098ca61 --- /dev/null +++ b/agent/core/approval_policy.py @@ -0,0 +1,11 @@ +"""Shared predicates for approval-gated tool operations.""" + +from typing import Any + + +def normalize_tool_operation(operation: Any) -> str: + return str(operation or "").strip().lower() + + +def is_scheduled_operation(operation: Any) -> bool: + return normalize_tool_operation(operation).startswith("scheduled ") diff --git a/agent/core/cost_estimation.py b/agent/core/cost_estimation.py new file mode 100644 index 000000000..a41ad196e --- /dev/null +++ b/agent/core/cost_estimation.py @@ -0,0 +1,282 @@ +"""Conservative cost estimates for auto-approved infrastructure actions.""" + +import os +import re +import time +from dataclasses import dataclass +from typing import Any + +import httpx + +OPENID_PROVIDER_URL = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co") +JOBS_HARDWARE_URL = f"{OPENID_PROVIDER_URL}/api/jobs/hardware" +JOBS_PRICE_CACHE_TTL_S = 6 * 60 * 60 + +DEFAULT_JOB_TIMEOUT_HOURS = 0.5 +DEFAULT_SANDBOX_RESERVATION_HOURS = 1.0 + +# Static fallback prices are intentionally conservative enough for a budget +# guard. The live /api/jobs/hardware catalog wins whenever it is reachable. +HF_JOBS_PRICE_USD_PER_HOUR: dict[str, float] = { + "cpu-basic": 0.05, + "cpu-upgrade": 0.25, + "cpu-performance": 0.50, + "cpu-xl": 1.00, + "t4-small": 0.60, + "t4-medium": 0.90, + "l4x1": 1.00, + "l4x4": 4.00, + "l40sx1": 2.00, + "l40sx4": 8.00, + "l40sx8": 16.00, + "a10g-small": 1.00, + "a10g-large": 2.00, + "a10g-largex2": 4.00, + "a10g-largex4": 8.00, + "a100-large": 4.00, + "a100x4": 16.00, + "a100x8": 32.00, + "h200": 10.00, + "h200x2": 20.00, + "h200x4": 40.00, + "h200x8": 80.00, + "inf2x6": 6.00, +} + +SPACE_PRICE_USD_PER_HOUR: dict[str, float] = { + "cpu-basic": 0.0, + "cpu-upgrade": 0.05, + "cpu-performance": 0.50, + "cpu-xl": 1.00, + "t4-small": 0.60, + "t4-medium": 0.90, + "l4x1": 1.00, + "l4x4": 4.00, + "l40sx1": 2.00, + "l40sx4": 8.00, + "l40sx8": 16.00, + "a10g-small": 1.00, + "a10g-large": 2.00, + "a10g-largex2": 4.00, + "a10g-largex4": 8.00, + "a100-large": 4.00, + "a100x4": 16.00, + "a100x8": 32.00, + "h200": 10.00, + "h200x2": 20.00, + "h200x4": 40.00, + "h200x8": 80.00, + "inf2x6": 6.00, +} + +_DURATION_RE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([smhd]?)\s*$", re.IGNORECASE) +_PRICE_RE = re.compile(r"(\d+(?:\.\d+)?)") +_jobs_price_cache: tuple[float, dict[str, float]] | None = None + + +@dataclass(frozen=True) +class CostEstimate: + """Estimated cost for a tool call. + + ``estimated_cost_usd=None`` means the call may be billable but we could not + estimate it safely, so auto-approval should fall back to a human decision. + """ + + estimated_cost_usd: float | None + billable: bool + block_reason: str | None = None + label: str | None = None + + +def parse_timeout_hours( + value: Any, *, default_hours: float = DEFAULT_JOB_TIMEOUT_HOURS +) -> float | None: + """Parse HF timeout values into hours. + + Strings accept ``s``, ``m``, ``h``, or ``d`` suffixes. Numeric values are + treated as seconds, matching the Hub client's typed timeout parameter. + """ + if value is None or value == "": + return default_hours + if isinstance(value, bool): + return None + if isinstance(value, int | float): + seconds = float(value) + return seconds / 3600 if seconds > 0 else None + if not isinstance(value, str): + return None + + match = _DURATION_RE.match(value) + if not match: + return None + amount = float(match.group(1)) + unit = match.group(2).lower() or "s" + if amount <= 0: + return None + if unit == "s": + return amount / 3600 + if unit == "m": + return amount / 60 + if unit == "h": + return amount + if unit == "d": + return amount * 24 + return None + + +def _extract_flavor(item: dict[str, Any]) -> str | None: + for key in ("flavor", "name", "id", "value", "hardware", "hardware_flavor"): + value = item.get(key) + if isinstance(value, str) and value: + return value + return None + + +def _coerce_price(value: Any) -> float | None: + if isinstance(value, bool) or value is None: + return None + if isinstance(value, int | float): + return float(value) if value >= 0 else None + if isinstance(value, str): + match = _PRICE_RE.search(value.replace(",", "")) + if match: + return float(match.group(1)) + return None + + +def _extract_hourly_price(item: dict[str, Any]) -> float | None: + for key in ( + "price", + "price_usd", + "priceUsd", + "price_per_hour", + "pricePerHour", + "hourly_price", + "hourlyPrice", + "usd_per_hour", + "usdPerHour", + ): + price = _coerce_price(item.get(key)) + if price is not None: + return price + for key in ("pricing", "billing", "cost"): + nested = item.get(key) + if isinstance(nested, dict): + price = _extract_hourly_price(nested) + if price is not None: + return price + return None + + +def _iter_hardware_items(payload: Any): + if isinstance(payload, list): + for item in payload: + yield from _iter_hardware_items(item) + elif isinstance(payload, dict): + if _extract_flavor(payload): + yield payload + for key in ("hardware", "flavors", "items", "data", "jobs"): + child = payload.get(key) + if child is not None: + yield from _iter_hardware_items(child) + + +def _parse_jobs_price_catalog(payload: Any) -> dict[str, float]: + prices: dict[str, float] = {} + for item in _iter_hardware_items(payload): + flavor = _extract_flavor(item) + price = _extract_hourly_price(item) + if flavor and price is not None: + prices[flavor] = price + return prices + + +async def hf_jobs_price_catalog() -> dict[str, float]: + """Return live HF Jobs hourly prices, falling back to static prices.""" + global _jobs_price_cache + now = time.monotonic() + if _jobs_price_cache and now - _jobs_price_cache[0] < JOBS_PRICE_CACHE_TTL_S: + return dict(_jobs_price_cache[1]) + + prices: dict[str, float] = {} + try: + async with httpx.AsyncClient(timeout=3.0) as client: + response = await client.get(JOBS_HARDWARE_URL) + if response.status_code == 200: + prices = _parse_jobs_price_catalog(response.json()) + except (httpx.HTTPError, ValueError): + prices = {} + + if not prices: + prices = dict(HF_JOBS_PRICE_USD_PER_HOUR) + else: + prices = {**HF_JOBS_PRICE_USD_PER_HOUR, **prices} + + _jobs_price_cache = (now, prices) + return dict(prices) + + +async def estimate_hf_job_cost(args: dict[str, Any]) -> CostEstimate: + flavor = str( + args.get("hardware_flavor") + or args.get("flavor") + or args.get("hardware") + or "cpu-basic" + ) + timeout_hours = parse_timeout_hours(args.get("timeout")) + if timeout_hours is None: + return CostEstimate( + estimated_cost_usd=None, + billable=True, + block_reason=f"Could not parse HF job timeout: {args.get('timeout')!r}.", + label=flavor, + ) + + prices = await hf_jobs_price_catalog() + price = prices.get(flavor) + if price is None: + return CostEstimate( + estimated_cost_usd=None, + billable=True, + block_reason=f"No price is available for HF job hardware '{flavor}'.", + label=flavor, + ) + + return CostEstimate( + estimated_cost_usd=round(price * timeout_hours, 4), + billable=price > 0, + label=flavor, + ) + + +async def estimate_sandbox_cost( + args: dict[str, Any], *, session: Any = None +) -> CostEstimate: + if session is not None and getattr(session, "sandbox", None): + return CostEstimate(estimated_cost_usd=0.0, billable=False, label="existing") + + hardware = str(args.get("hardware") or "cpu-basic") + price = SPACE_PRICE_USD_PER_HOUR.get(hardware) + if price is None: + return CostEstimate( + estimated_cost_usd=None, + billable=True, + block_reason=f"No price is available for sandbox hardware '{hardware}'.", + label=hardware, + ) + + return CostEstimate( + estimated_cost_usd=round(price * DEFAULT_SANDBOX_RESERVATION_HOURS, 4), + billable=price > 0, + label=hardware, + ) + + +async def estimate_tool_cost( + tool_name: str, args: dict[str, Any], *, session: Any = None +) -> CostEstimate: + if tool_name == "sandbox_create": + return await estimate_sandbox_cost(args, session=session) + if tool_name == "hf_jobs": + return await estimate_hf_job_cost(args) + return CostEstimate(estimated_cost_usd=0.0, billable=False) diff --git a/agent/core/doom_loop.py b/agent/core/doom_loop.py index 5050d7550..3b57fe2cc 100644 --- a/agent/core/doom_loop.py +++ b/agent/core/doom_loop.py @@ -17,25 +17,58 @@ @dataclass(frozen=True) class ToolCallSignature: - """Hashable signature for a single tool call (name + args hash).""" + """Hashable signature for a single tool call plus its observed result.""" name: str args_hash: str + result_hash: str | None = None + + +def _normalize_args(args_str: str) -> str: + """Canonicalise a tool-call arguments string before hashing. + + LLMs can emit semantically-identical JSON for the same call with different + key orderings (``{"a": 1, "b": 2}`` vs ``{"b": 2, "a": 1}``) or whitespace + (``{"a":1}`` vs ``{"a": 1}``). Hashing the raw bytes makes the doom-loop + detector miss those repeats. We parse-and-redump with ``sort_keys=True`` + plus the most compact separators so trivially-different spellings collapse + to the same canonical form. + + Falls back to the original string if the input isn't valid JSON (e.g. a + handful of providers occasionally pass a bare string for ``arguments``); + that path keeps the legacy behaviour and never raises. + """ + if not args_str: + return "" + try: + return json.dumps(json.loads(args_str), sort_keys=True, separators=(",", ":")) + except (json.JSONDecodeError, TypeError, ValueError): + return args_str def _hash_args(args_str: str) -> str: - """Return a short hash of the JSON arguments string.""" - return hashlib.md5(args_str.encode()).hexdigest()[:12] + """Return a short hash of the JSON arguments string. + + The input is normalised via :func:`_normalize_args` first so that + semantically-identical tool calls produce the same hash regardless of key + order or whitespace. + """ + return hashlib.md5(_normalize_args(args_str).encode()).hexdigest()[:12] def extract_recent_tool_signatures( messages: list[Message], lookback: int = 30 ) -> list[ToolCallSignature]: - """Extract tool call signatures from recent assistant messages.""" + """Extract tool call signatures from recent assistant messages. + + Includes the immediate tool result hash when present. This prevents + legitimate polling from being classified as a doom loop when the poll + arguments stay constant but the observed result keeps changing. + """ signatures: list[ToolCallSignature] = [] recent = messages[-lookback:] if len(messages) > lookback else messages - for msg in recent: + for idx, msg in enumerate(recent): if getattr(msg, "role", None) != "assistant": continue tool_calls = getattr(msg, "tool_calls", None) @@ -47,7 +80,23 @@ def extract_recent_tool_signatures( continue name = getattr(fn, "name", "") or "" args_str = getattr(fn, "arguments", "") or "" - signatures.append(ToolCallSignature(name=name, args_hash=_hash_args(args_str))) + result_hash = None + for follow in recent[idx + 1 :]: + role = getattr(follow, "role", None) + if role == "tool" and getattr(follow, "tool_call_id", None) == getattr( + tc, "id", None + ): + result_hash = _hash_args(str(getattr(follow, "content", "") or "")) + break + if role in {"assistant", "user"}: + break + signatures.append( + ToolCallSignature( + name=name, + args_hash=_hash_args(args_str), + result_hash=result_hash, + ) + ) return signatures @@ -109,9 +158,13 @@ def check_for_doom_loop(messages: list[Message]) -> str | None: # Check for identical consecutive calls tool_name = detect_identical_consecutive(signatures, threshold=3) if tool_name: - logger.warning("Doom loop detected: %d+ identical consecutive calls to '%s'", 3, tool_name) + logger.warning( + "Repetition guard activated: %d+ identical consecutive calls to '%s'", + 3, + tool_name, + ) return ( - f"[SYSTEM: DOOM LOOP DETECTED] You have called '{tool_name}' with the same " + f"[SYSTEM: REPETITION GUARD] You have called '{tool_name}' with the same " f"arguments multiple times in a row, getting the same result each time. " f"STOP repeating this approach — it is not working. " f"Step back and try a fundamentally different strategy. " @@ -123,9 +176,11 @@ def check_for_doom_loop(messages: list[Message]) -> str | None: pattern = detect_repeating_sequence(signatures) if pattern: pattern_desc = " → ".join(s.name for s in pattern) - logger.warning("Doom loop detected: repeating sequence [%s]", pattern_desc) + logger.warning( + "Repetition guard activated: repeating sequence [%s]", pattern_desc + ) return ( - f"[SYSTEM: DOOM LOOP DETECTED] You are stuck in a repeating cycle of tool calls: " + f"[SYSTEM: REPETITION GUARD] You are stuck in a repeating cycle of tool calls: " f"[{pattern_desc}]. This pattern has repeated multiple times without progress. " f"STOP this cycle and try a fundamentally different approach. " f"Consider: breaking down the problem differently, using alternative tools, " diff --git a/agent/core/effort_probe.py b/agent/core/effort_probe.py index 142feaaa1..583fdd5a9 100644 --- a/agent/core/effort_probe.py +++ b/agent/core/effort_probe.py @@ -22,30 +22,37 @@ import asyncio import logging +import time from dataclasses import dataclass +from typing import Any from litellm import acompletion from agent.core.llm_params import UnsupportedEffortError, _resolve_llm_params +from agent.core.prompt_caching import router_session_id_for, with_prompt_cache_params +from agent.core.yolo_budget import maybe_pause_yolo_after_spend logger = logging.getLogger(__name__) # Cascade: for each user-stated preference, the ordered list of levels to -# try. First success wins. ``max`` / ``xhigh`` are Anthropic-only; providers -# that don't accept them raise ``UnsupportedEffortError`` synchronously (no -# wasted network round-trip) and we advance to the next level. +# try. First success wins. HF Router accepts low/medium/high generically; +# higher preferences are kept in the cascade for future/provider-specific +# support and are skipped synchronously when unsupported. _EFFORT_CASCADE: dict[str, list[str]] = { - "max": ["max", "xhigh", "high", "medium", "low"], - "xhigh": ["xhigh", "high", "medium", "low"], - "high": ["high", "medium", "low"], - "medium": ["medium", "low"], + "max": ["max", "xhigh", "high", "medium", "low"], + "xhigh": ["xhigh", "high", "medium", "low"], + "high": ["high", "medium", "low"], + "medium": ["medium", "low"], "minimal": ["minimal", "low"], - "low": ["low"], + "low": ["low"], } _PROBE_TIMEOUT = 15.0 -_PROBE_MAX_TOKENS = 16 +# Keep the probe cheap, but high enough that frontier reasoning models can +# finish a trivial reply instead of tripping a false "output limit reached" +# error during capability detection. +_PROBE_MAX_TOKENS = 64 class ProbeInconclusive(Exception): @@ -63,6 +70,7 @@ class ProbeOutcome: * str → send this level * None → model doesn't support thinking; strip it """ + effective_effort: str | None attempts: int elapsed_ms: int @@ -72,9 +80,7 @@ class ProbeOutcome: def _is_thinking_unsupported(e: Exception) -> bool: """Model rejected any thinking config. - Matches Anthropic's 'thinking.type.enabled is not supported for this - model' as well as the adaptive variant. Substring-match because the - exact wording shifts across API versions. + Substring-match because exact wording shifts across models and providers. """ s = str(e).lower() return "thinking" in s and "not supported" in s @@ -83,16 +89,12 @@ def _is_thinking_unsupported(e: Exception) -> bool: def _is_invalid_effort(e: Exception) -> bool: """The requested effort level isn't accepted for this model. - Covers both API responses (Anthropic/OpenAI 400 with "invalid", "must - be one of", etc.) and LiteLLM's local validation that fires *before* - the request (e.g. "effort='max' is only supported by Claude Opus 4.6" - — LiteLLM knows max is Opus-4.6-only and raises synchronously). The - cascade walks down on either. + Covers API responses with "invalid", "must be one of", etc. and local + validation that fires *before* the request. The cascade walks down on + either. Explicitly returns False when the message is really about thinking - itself (e.g. Anthropic's 4.7 error mentions ``output_config.effort`` - in its fix hint, but the actual failure is ``thinking.type.enabled`` - being unsupported). That case is caught by ``_is_thinking_unsupported``. + itself. That case is caught by ``_is_thinking_unsupported``. """ if _is_thinking_unsupported(e): return False @@ -102,10 +104,15 @@ def _is_invalid_effort(e: Exception) -> bool: return any( phrase in s for phrase in ( - "invalid", "not supported", "must be one of", "not a valid", - "unrecognized", "unknown", + "invalid", + "not supported", + "must be one of", + "not a valid", + "unrecognized", + "unknown", # LiteLLM's own pre-flight validation phrasing. - "only supported by", "is only supported", + "only supported by", + "is only supported", ) ) @@ -122,11 +129,23 @@ def _is_transient(e: Exception) -> bool: return any( p in s for p in ( - "timeout", "timed out", "429", "rate limit", - "503", "service unavailable", "502", "bad gateway", - "500", "internal server error", "overloaded", "capacity", - "connection reset", "connection refused", "connection error", - "eof", "broken pipe", + "timeout", + "timed out", + "429", + "rate limit", + "503", + "service unavailable", + "502", + "bad gateway", + "500", + "internal server error", + "overloaded", + "capacity", + "connection reset", + "connection refused", + "connection error", + "eof", + "broken pipe", ) ) @@ -135,6 +154,7 @@ async def probe_effort( model_name: str, preference: str | None, hf_token: str | None, + session: Any = None, ) -> ProbeOutcome: """Walk the cascade for ``preference`` on ``model_name``. @@ -143,6 +163,12 @@ async def probe_effort( transient errors (5xx, timeout) — persistent 4xx that aren't thinking/ effort related bubble as the original exception so callers can surface them (auth, model-not-found, quota, etc.). + + ``session`` is optional; when provided, each successful probe attempt + is recorded via ``telemetry.record_llm_call(kind="effort_probe")`` so + the cost shows up in the session's ``total_cost_usd``. Failed probes + (rejected by the provider) typically aren't billed, so we only record + on success. """ loop = asyncio.get_event_loop() start = loop.time() @@ -160,7 +186,14 @@ async def probe_effort( for effort in cascade: try: params = _resolve_llm_params( - model_name, hf_token, reasoning_effort=effort, strict=True, + model_name, + hf_token, + reasoning_effort=effort, + strict=True, + ) + params = with_prompt_cache_params( + params, + session_id=router_session_id_for(session), ) except UnsupportedEffortError: # Provider can't even accept this effort name (e.g. "max" on @@ -169,16 +202,49 @@ async def probe_effort( continue attempts += 1 + probe_messages = [{"role": "user", "content": "ping"}] + params = {**params, "max_tokens": _PROBE_MAX_TOKENS} try: - await asyncio.wait_for( + _t0 = time.monotonic() + response = await asyncio.wait_for( acompletion( - messages=[{"role": "user", "content": "ping"}], - max_tokens=_PROBE_MAX_TOKENS, + messages=probe_messages, stream=False, **params, ), timeout=_PROBE_TIMEOUT, ) + if session is not None: + # Best-effort telemetry — never let a logging blip propagate + # out of the probe and break model switching. + try: + from agent.core import telemetry + + usage = await telemetry.record_llm_call( + session, + model=model_name, + response=response, + latency_ms=int((time.monotonic() - _t0) * 1000), + finish_reason=response.choices[0].finish_reason + if response.choices + else None, + kind="effort_probe", + ) + if await maybe_pause_yolo_after_spend( + session, + spend_kind="effort_probe", + observed_cost_usd=usage.get("cost_usd") + if isinstance(usage, dict) + else None, + ): + return ProbeOutcome( + effective_effort=effort, + attempts=attempts, + elapsed_ms=int((loop.time() - start) * 1000), + note="YOLO budget paused effort probe", + ) + except Exception as _telem_err: + logger.debug("effort_probe telemetry failed: %s", _telem_err) except Exception as e: last_error = e if _is_thinking_unsupported(e): @@ -190,7 +256,9 @@ async def probe_effort( note="model doesn't support reasoning, dropped", ) if _is_invalid_effort(e): - logger.debug("probe: %s rejected effort=%s, trying next", model_name, effort) + logger.debug( + "probe: %s rejected effort=%s, trying next", model_name, effort + ) continue if _is_transient(e): raise ProbeInconclusive(str(e)) from e diff --git a/agent/core/hf_access.py b/agent/core/hf_access.py new file mode 100644 index 000000000..cdfaaf41c --- /dev/null +++ b/agent/core/hf_access.py @@ -0,0 +1,201 @@ +"""Helpers for Hugging Face account / org access decisions. + +HF Jobs are gated by *credits*, not by HF Pro subscriptions. Any user who +has credits — on their personal account or on an org they belong to — can +launch jobs under that namespace. The picker UI lets the caller choose +which wallet to bill. +""" + +from __future__ import annotations + +import asyncio +import os +import re +from dataclasses import dataclass +from typing import Any, Literal + +import httpx + +OPENID_PROVIDER_URL = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co") +HF_BILLING_URL = "https://huggingface.co/settings/billing" +HF_PRO_SUBSCRIBE_URL = "https://huggingface.co/subscribe/pro" + +HfUserPlan = Literal["free", "pro"] + + +@dataclass(frozen=True) +class JobsAccess: + """Namespaces the caller may bill HF Jobs to.""" + + username: str | None + org_names: list[str] + eligible_namespaces: list[str] + default_namespace: str | None + + +class JobsAccessError(Exception): + """Structured jobs-namespace error. + + ``namespace_required`` fires when the caller belongs to more than one + eligible namespace and the UI must prompt them to pick one. There is no + longer an ``upgrade_required`` state — Pro is irrelevant; HF Jobs are + gated on per-wallet credits, surfaced separately when the API returns + a billing error at job-creation time. + """ + + def __init__( + self, + message: str, + *, + access: JobsAccess | None = None, + namespace_required: bool = False, + ) -> None: + super().__init__(message) + self.access = access + self.namespace_required = namespace_required + + +def _extract_username(whoami: dict[str, Any]) -> str | None: + for key in ("name", "user", "preferred_username"): + value = whoami.get(key) + if isinstance(value, str) and value: + return value + return None + + +def _org_names(whoami: dict[str, Any]) -> list[str]: + """All orgs the caller belongs to. + + Plan/tier is ignored — credits live on the namespace itself, so any + org the user belongs to can host a job as long as it has credits. + """ + names: list[str] = [] + orgs = whoami.get("orgs") or [] + if not isinstance(orgs, list): + return names + for org in orgs: + if not isinstance(org, dict): + continue + name = org.get("name") + if isinstance(name, str) and name: + names.append(name) + return sorted(set(names)) + + +def jobs_access_from_whoami(whoami: dict[str, Any]) -> JobsAccess: + username = _extract_username(whoami) + org_names = _org_names(whoami) + eligible: list[str] = [] + if username: + eligible.append(username) + eligible.extend(org_names) + default = username if username else (org_names[0] if org_names else None) + return JobsAccess( + username=username, + org_names=org_names, + eligible_namespaces=eligible, + default_namespace=default, + ) + + +def normalize_hf_user_plan(whoami: Any) -> HfUserPlan | None: + """Normalize a whoami-v2 payload to the supported HF account plan tiers.""" + if not isinstance(whoami, dict): + return None + if whoami.get("isPro") is True: + return "pro" + return "free" + + +async def fetch_whoami_v2(token: str, timeout: float = 5.0) -> dict[str, Any] | None: + if not token: + return None + async with httpx.AsyncClient(timeout=timeout) as client: + try: + response = await client.get( + f"{OPENID_PROVIDER_URL}/api/whoami-v2", + headers={"Authorization": f"Bearer {token}"}, + ) + if response.status_code != 200: + return None + payload = response.json() + return payload if isinstance(payload, dict) else None + except (httpx.HTTPError, ValueError): + return None + + +async def get_jobs_access(token: str) -> JobsAccess | None: + whoami = await fetch_whoami_v2(token) + if whoami is None: + return None + return jobs_access_from_whoami(whoami) + + +async def resolve_jobs_namespace( + token: str, + requested_namespace: str | None = None, +) -> tuple[str, JobsAccess | None]: + """Return the namespace to use for jobs. + + If whoami-v2 is unavailable, fall back to the token owner's username. + """ + access = await get_jobs_access(token) + if access: + if requested_namespace: + if requested_namespace in access.eligible_namespaces: + return requested_namespace, access + raise JobsAccessError( + f"You can only run jobs under your own account or an org you belong to. " + f"Allowed namespaces: {', '.join(access.eligible_namespaces) or '(none)'}", + access=access, + ) + if access.default_namespace: + return access.default_namespace, access + raise JobsAccessError( + "Couldn't resolve a Hugging Face namespace for this token.", + access=access, + ) + + # Fallback: whoami-v2 unavailable. Don't block the call pre-emptively. + from huggingface_hub import HfApi + + username = None + if token: + whoami = await asyncio.to_thread(HfApi(token=token).whoami) + username = whoami.get("name") + if not username: + raise JobsAccessError("No HF token available to resolve a jobs namespace.") + return requested_namespace or username, None + + +_BILLING_PATTERNS = re.compile( + r"\b(insufficient[_\s-]?credits?|out\s+of\s+credits?|" + r"payment\s+required|billing|no\s+credits?|add\s+credits?|requires?\s+credits?|" + r"credits?\s+(?:exhausted|used\s+up|limit))\b", + re.IGNORECASE, +) + +_INFERENCE_BILLING_PATTERNS = re.compile( + r"\b(insufficient[_\s-]?quota|out\s+of\s+monthly\s+credits?|" + r"exhausted\s+monthly\s+credits?|" + r"quota[_\s-]?(?:exceeded|exhausted|limit|insufficient)|" + r"monthly\s+credits?\s+(?:exhausted|used\s+up|limit))\b", + re.IGNORECASE, +) + + +def is_billing_error(message: str) -> bool: + """True if an HF API error message looks like an out-of-credits / billing error.""" + if not message: + return False + if "402" in message: + return True + return bool(_BILLING_PATTERNS.search(message)) + + +def is_inference_billing_error(error: Exception | str) -> bool: + """True if an Inference Providers error looks like exhausted credits.""" + message = str(error) + return is_billing_error(message) or bool( + _INFERENCE_BILLING_PATTERNS.search(message) + ) diff --git a/agent/core/hf_router_catalog.py b/agent/core/hf_router_catalog.py index f6f519d03..1efabb5ed 100644 --- a/agent/core/hf_router_catalog.py +++ b/agent/core/hf_router_catalog.py @@ -7,7 +7,6 @@ • Validate ``/model`` switches with live data instead of a hard-coded allowlist. • Show the user which providers serve a model, at what price, and whether they support tool calls. - • Derive a reasonable context-window limit for any routed model. The listing is cached in-memory for a few minutes so repeated lookups during a session are free. On fetch failure we return stale data if we have it, or an @@ -26,10 +25,12 @@ _CATALOG_URL = "https://router.huggingface.co/v1/models" _CACHE_TTL_SECONDS = 300 +_CACHE_FAILURE_TTL_SECONDS = 15 _HTTP_TIMEOUT_SECONDS = 5.0 _cache: Optional[dict] = None _cache_time: float = 0.0 +_last_fetch_error: Optional[str] = None @dataclass @@ -40,7 +41,6 @@ class ProviderInfo: input_price: Optional[float] output_price: Optional[float] supports_tools: bool - supports_structured_output: bool @dataclass @@ -52,31 +52,29 @@ class ModelInfo: def live_providers(self) -> list[ProviderInfo]: return [p for p in self.providers if p.status == "live"] - @property - def max_context_length(self) -> Optional[int]: - lengths = [p.context_length for p in self.live_providers if p.context_length] - return max(lengths) if lengths else None - @property def any_supports_tools(self) -> bool: return any(p.supports_tools for p in self.live_providers) def _fetch_catalog(force: bool = False) -> dict: - global _cache, _cache_time + global _cache, _cache_time, _last_fetch_error now = time.time() - if not force and _cache is not None and now - _cache_time < _CACHE_TTL_SECONDS: + ttl = _CACHE_FAILURE_TTL_SECONDS if _last_fetch_error else _CACHE_TTL_SECONDS + if not force and _cache is not None and now - _cache_time < ttl: return _cache try: resp = httpx.get(_CATALOG_URL, timeout=_HTTP_TIMEOUT_SECONDS) resp.raise_for_status() _cache = resp.json() _cache_time = now + _last_fetch_error = None except Exception as e: logger.warning("Failed to fetch HF router catalog: %s", e) + _last_fetch_error = str(e) if _cache is None: _cache = {"data": []} - _cache_time = now + _cache_time = now return _cache @@ -92,7 +90,6 @@ def _parse_entry(entry: dict) -> ModelInfo: input_price=pricing.get("input"), output_price=pricing.get("output"), supports_tools=bool(p.get("supports_tools", False)), - supports_structured_output=bool(p.get("supports_structured_output", False)), ) ) return ModelInfo(id=entry.get("id", ""), providers=providers) diff --git a/agent/core/hf_tokens.py b/agent/core/hf_tokens.py new file mode 100644 index 000000000..a7e3df341 --- /dev/null +++ b/agent/core/hf_tokens.py @@ -0,0 +1,77 @@ +"""Hugging Face token resolution helpers.""" + +from __future__ import annotations + +import os +from typing import Any + + +def clean_hf_token(token: str | None) -> str | None: + """Normalize token strings the same way huggingface_hub does.""" + if token is None: + return None + return token.replace("\r", "").replace("\n", "").strip() or None + + +def get_cached_hf_token() -> str | None: + """Return the token from huggingface_hub's normal env/cache lookup.""" + try: + from huggingface_hub import get_token + + return get_token() + except Exception: + return None + + +def resolve_hf_token( + *candidates: str | None, + include_cached: bool = True, +) -> str | None: + """Return the first non-empty explicit token, then optionally HF cache.""" + for token in candidates: + cleaned = clean_hf_token(token) + if cleaned: + return cleaned + if include_cached: + return get_cached_hf_token() + return None + + +def resolve_hf_router_token(session_hf_token: str | None = None) -> str | None: + """Resolve the token used for Hugging Face Router LLM calls. + + App-specific precedence: + 1. session_hf_token: the active user/session token. + 2. huggingface_hub.get_token(): HF_TOKEN/HUGGING_FACE_HUB_TOKEN or + local ``hf auth login`` cache. + """ + return resolve_hf_token(session_hf_token) + + +def bearer_token_from_header(auth_header: str | None) -> str | None: + """Extract a cleaned bearer token from an Authorization header.""" + if not auth_header or not auth_header.startswith("Bearer "): + return None + return clean_hf_token(auth_header[7:]) + + +def resolve_hf_request_token( + request: Any, + *, + include_env_fallback: bool = True, +) -> str | None: + """Resolve a user token from a FastAPI request. + + This intentionally does not use the local ``hf auth login`` cache. Backend + request paths should act as the browser user from Authorization/cookie, or + fall back only to an explicit server ``HF_TOKEN`` in dev/server contexts. + """ + token = bearer_token_from_header(request.headers.get("Authorization", "")) + if token: + return token + token = clean_hf_token(request.cookies.get("hf_access_token")) + if token: + return token + if include_env_fallback: + return clean_hf_token(os.environ.get("HF_TOKEN")) + return None diff --git a/agent/core/hub_artifacts.py b/agent/core/hub_artifacts.py new file mode 100644 index 000000000..8a0b1b5b1 --- /dev/null +++ b/agent/core/hub_artifacts.py @@ -0,0 +1,758 @@ +"""Best-effort Hub metadata for artifacts generated by ML Intern sessions.""" + +import base64 +import logging +import re +import shlex +import tempfile +import textwrap +from datetime import datetime +from pathlib import Path +from typing import Any + +from huggingface_hub import hf_hub_download +from huggingface_hub.repocard import metadata_load, metadata_save +from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError + +logger = logging.getLogger(__name__) + +ML_INTERN_TAG = "ml-intern" +SUPPORTED_REPO_TYPES = {"model", "dataset", "space"} +PROVENANCE_MARKER = "" +_COLLECTION_TITLE_PREFIX = "ml-intern-artifacts" +_COLLECTION_TITLE_MAX_LENGTH = 59 +_UUID_SESSION_ID_RE = re.compile( + r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-" + r"[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$" +) +_KNOWN_ARTIFACTS_ATTR = "_ml_intern_known_hub_artifacts" +_REGISTERED_ARTIFACTS_ATTR = "_ml_intern_registered_hub_artifacts" +_COLLECTION_SLUG_ATTR = "_ml_intern_artifact_collection_slug" +_SESSION_ARTIFACT_SET_FALLBACK: dict[tuple[int, str], set[str]] = {} +_USAGE_HEADING_RE = re.compile( + r"^#{2,6}\s+(usage|how to use|using this (model|dataset)|use this (model|dataset))\b", + re.IGNORECASE | re.MULTILINE, +) +_FRONT_MATTER_RE = re.compile(r"\A---\s*\n.*?\n---\s*\n?", re.DOTALL) + + +def _safe_session_id(session: Any) -> str: + raw = str(getattr(session, "session_id", "") or "unknown-session") + safe = re.sub(r"[^A-Za-z0-9._-]+", "-", raw).strip("-") + return safe or "unknown-session" + + +def session_artifact_date(session: Any) -> str: + """Return the YYYY-MM-DD partition date for a session.""" + raw = getattr(session, "session_start_time", None) + if raw: + try: + return datetime.fromisoformat(str(raw).replace("Z", "+00:00")).strftime( + "%Y-%m-%d" + ) + except ValueError: + logger.debug("Could not parse session_start_time=%r", raw) + return datetime.utcnow().strftime("%Y-%m-%d") + + +def _collection_session_id_fragment(session: Any) -> str: + safe_id = _safe_session_id(session) + if _UUID_SESSION_ID_RE.match(safe_id): + return safe_id[:8] + stem = f"{_COLLECTION_TITLE_PREFIX}-{session_artifact_date(session)}-" + max_id_length = max(1, _COLLECTION_TITLE_MAX_LENGTH - len(stem)) + if len(safe_id) <= max_id_length: + return safe_id + return safe_id[:max_id_length].rstrip("-._") or safe_id[:max_id_length] + + +def artifact_collection_title(session: Any) -> str: + return ( + f"{_COLLECTION_TITLE_PREFIX}-{session_artifact_date(session)}-" + f"{_collection_session_id_fragment(session)}" + ) + + +def _artifact_key(repo_id: str, repo_type: str | None) -> str: + return f"{repo_type or 'model'}:{repo_id}" + + +def _sandbox_space_name_pattern() -> str: + from agent.tools.sandbox_tool import SANDBOX_SPACE_NAME_RE + + return SANDBOX_SPACE_NAME_RE.pattern + + +def is_sandbox_hub_repo(repo_id: str | None, repo_type: str | None) -> bool: + """Return True for ML Intern's ephemeral sandbox Space repos.""" + if (repo_type or "model") != "space" or not repo_id: + return False + repo_name = str(repo_id).rsplit("/", 1)[-1] + return bool(re.fullmatch(_sandbox_space_name_pattern(), repo_name)) + + +def _session_artifact_set(session: Any, attr: str) -> set[str]: + current = getattr(session, attr, None) + if isinstance(current, set): + return current + current = set() + try: + setattr(session, attr, current) + except Exception: + logger.warning( + "Could not attach %s to session; using process-local fallback state", + attr, + ) + return _SESSION_ARTIFACT_SET_FALLBACK.setdefault((id(session), attr), set()) + return current + + +def remember_hub_artifact(session: Any, repo_id: str, repo_type: str | None) -> None: + if session is None or not repo_id: + return + _session_artifact_set(session, _KNOWN_ARTIFACTS_ATTR).add( + _artifact_key(repo_id, repo_type) + ) + + +def is_known_hub_artifact(session: Any, repo_id: str, repo_type: str | None) -> bool: + if session is None or not repo_id: + return False + return _artifact_key(repo_id, repo_type) in _session_artifact_set( + session, _KNOWN_ARTIFACTS_ATTR + ) + + +def _merge_tags(metadata: dict[str, Any], tag: str = ML_INTERN_TAG) -> dict[str, Any]: + merged = dict(metadata) + raw_tags = merged.get("tags") + if raw_tags is None: + tags: list[str] = [] + elif isinstance(raw_tags, str): + tags = [raw_tags] + elif isinstance(raw_tags, list): + tags = [str(item) for item in raw_tags] + else: + tags = [str(raw_tags)] + + if tag not in tags: + tags.append(tag) + merged["tags"] = tags + return merged + + +def _metadata_from_content(content: str) -> dict[str, Any]: + with tempfile.TemporaryDirectory() as tmp_dir: + path = Path(tmp_dir) / "README.md" + path.write_text(content, encoding="utf-8") + return metadata_load(path) or {} + + +def _content_with_metadata(content: str, metadata: dict[str, Any]) -> str: + with tempfile.TemporaryDirectory() as tmp_dir: + path = Path(tmp_dir) / "README.md" + path.write_text(content, encoding="utf-8") + metadata_save(path, metadata) + return path.read_text(encoding="utf-8") + + +def _body_without_metadata(content: str) -> str: + return _FRONT_MATTER_RE.sub("", content, count=1).strip() + + +def _append_section(content: str, section: str) -> str: + base = content.rstrip() + if base: + return f"{base}\n\n{section.strip()}\n" + return f"{section.strip()}\n" + + +def _provenance_section(repo_type: str) -> str: + label = {"model": "model", "dataset": "dataset"}.get(repo_type, "Hub") + return f"""{PROVENANCE_MARKER} +## Generated by ML Intern + +This {label} repository was generated by [ML Intern](https://github.com/huggingface/ml-intern), an agent for machine learning research and development on the Hugging Face Hub. + +- Try ML Intern: https://smolagents-ml-intern.hf.space +- Source code: https://github.com/huggingface/ml-intern +""" + + +def _usage_section(repo_id: str, repo_type: str) -> str: + if repo_type == "dataset": + return f"""## Usage + +```python +from datasets import load_dataset + +dataset = load_dataset("{repo_id}") +``` +""" + + return f"""## Usage + +```python +from transformers import AutoModelForCausalLM, AutoTokenizer + +model_id = "{repo_id}" +tokenizer = AutoTokenizer.from_pretrained(model_id) +model = AutoModelForCausalLM.from_pretrained(model_id) +``` + +For non-causal architectures, replace `AutoModelForCausalLM` with the appropriate `AutoModel` class. +""" + + +def augment_repo_card_content( + content: str | None, + repo_id: str, + repo_type: str = "model", + *, + extra_metadata: dict[str, Any] | None = None, +) -> str: + """Return README content with ML Intern metadata and provenance added.""" + repo_type = repo_type or "model" + content = content or "" + metadata = _metadata_from_content(content) + if extra_metadata: + metadata = {**extra_metadata, **metadata} + metadata = _merge_tags(metadata) + updated = _content_with_metadata(content, metadata) + + if not _body_without_metadata(updated): + updated = _append_section(updated, f"# {repo_id}") + + if repo_type in {"model", "dataset"} and PROVENANCE_MARKER not in updated: + updated = _append_section(updated, _provenance_section(repo_type)) + if not _USAGE_HEADING_RE.search(content): + updated = _append_section(updated, _usage_section(repo_id, repo_type)) + + return updated + + +def _read_remote_readme( + api: Any, + repo_id: str, + repo_type: str, + *, + token: str | bool | None = None, +) -> str: + token_value = token if token is not None else getattr(api, "token", None) + try: + readme_path = hf_hub_download( + repo_id=repo_id, + filename="README.md", + repo_type=repo_type, + token=token_value, + ) + except (EntryNotFoundError, RepositoryNotFoundError): + return "" + return Path(readme_path).read_text(encoding="utf-8") + + +def _update_repo_card( + api: Any, + repo_id: str, + repo_type: str, + *, + token: str | bool | None = None, + extra_metadata: dict[str, Any] | None = None, +) -> None: + current = _read_remote_readme(api, repo_id, repo_type, token=token) + updated = augment_repo_card_content( + current, + repo_id, + repo_type, + extra_metadata=extra_metadata, + ) + if updated == current: + return + api.upload_file( + path_or_fileobj=updated.encode("utf-8"), + path_in_repo="README.md", + repo_id=repo_id, + repo_type=repo_type, + token=token, + commit_message="Update ML Intern artifact metadata", + ) + + +def _ensure_collection_slug( + api: Any, + session: Any, + *, + token: str | bool | None = None, +) -> str | None: + slug = getattr(session, _COLLECTION_SLUG_ATTR, None) + if slug: + return slug + + title = artifact_collection_title(session) + collection = api.create_collection( + title=title, + description=( + f"Artifacts generated by ML Intern session {_safe_session_id(session)} " + f"on {session_artifact_date(session)}." + ), + private=True, + exists_ok=True, + token=token, + ) + slug = getattr(collection, "slug", None) + if slug: + setattr(session, _COLLECTION_SLUG_ATTR, slug) + return slug + + +def _add_to_collection( + api: Any, + session: Any, + repo_id: str, + repo_type: str, + *, + token: str | bool | None = None, +) -> bool: + slug = _ensure_collection_slug(api, session, token=token) + if not slug: + return False + api.add_collection_item( + collection_slug=slug, + item_id=repo_id, + item_type=repo_type, + note=( + f"Generated by ML Intern session {_safe_session_id(session)} " + f"on {session_artifact_date(session)}." + ), + exists_ok=True, + token=token, + ) + return True + + +def register_hub_artifact( + api: Any, + repo_id: str, + repo_type: str = "model", + *, + session: Any = None, + token: str | bool | None = None, + extra_metadata: dict[str, Any] | None = None, + force: bool = False, +) -> bool: + """Tag, card, and collection-register a Hub artifact without raising.""" + if session is None or not repo_id: + return False + repo_type = repo_type or "model" + if repo_type not in SUPPORTED_REPO_TYPES: + return False + if is_sandbox_hub_repo(repo_id, repo_type): + return False + + key = _artifact_key(repo_id, repo_type) + remember_hub_artifact(session, repo_id, repo_type) + registered = _session_artifact_set(session, _REGISTERED_ARTIFACTS_ATTR) + if key in registered and not force: + return True + + token_value = token if token is not None else getattr(api, "token", None) + card_updated = False + collection_updated = False + try: + _update_repo_card( + api, + repo_id, + repo_type, + token=token_value, + extra_metadata=extra_metadata, + ) + card_updated = True + except Exception as e: + logger.debug("ML Intern repo-card update failed for %s: %s", repo_id, e) + + try: + collection_updated = _add_to_collection( + api, + session, + repo_id, + repo_type, + token=token_value, + ) + except Exception as e: + logger.debug("ML Intern collection update failed for %s: %s", repo_id, e) + + if card_updated and collection_updated: + registered.add(key) + return True + return False + + +def build_hub_artifact_sitecustomize(session: Any) -> str: + """Build standalone sitecustomize.py code for HF Jobs Python processes.""" + if session is None or not getattr(session, "session_id", None): + return "" + + session_id = _safe_session_id(session) + session_date = session_artifact_date(session) + collection_title = artifact_collection_title(session) + collection_slug = getattr(session, _COLLECTION_SLUG_ATTR, None) + + return ( + textwrap.dedent( + f""" + # Auto-generated by ML Intern. Best-effort Hub artifact metadata only. + def _install_ml_intern_artifact_hooks(): + import os + import re + import tempfile + from pathlib import Path + + try: + import huggingface_hub as _hub + from huggingface_hub import HfApi, hf_hub_download + from huggingface_hub.repocard import metadata_load, metadata_save + from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError + except Exception: + return + + session_id = {session_id!r} + session_date = {session_date!r} + collection_title = {collection_title!r} + tag = {ML_INTERN_TAG!r} + marker = {PROVENANCE_MARKER!r} + supported = {sorted(SUPPORTED_REPO_TYPES)!r} + sandbox_space_re = re.compile({_sandbox_space_name_pattern()!r}) + registering = False + collection_slug = {collection_slug!r} + registered = set() + usage_re = re.compile( + r"^#{{2,6}}\\s+(usage|how to use|using this (model|dataset)|use this (model|dataset))\\b", + re.IGNORECASE | re.MULTILINE, + ) + front_matter_re = re.compile(r"\\A---\\s*\\n.*?\\n---\\s*\\n?", re.DOTALL) + collection_cache_path = ( + os.environ.get("ML_INTERN_ARTIFACT_COLLECTION_CACHE") + or str( + Path(tempfile.gettempdir()) + / f"ml-intern-artifacts-{{session_id}}.collection" + ) + ) + + def _token(value=None, api=None): + if isinstance(value, str) and value: + return value + api_token = getattr(api, "token", None) + if isinstance(api_token, str) and api_token: + return api_token + return ( + os.environ.get("HF_TOKEN") + or os.environ.get("HUGGINGFACE_HUB_TOKEN") + or None + ) + + def _merge_tags(metadata): + metadata = dict(metadata or {{}}) + raw_tags = metadata.get("tags") + if raw_tags is None: + tags = [] + elif isinstance(raw_tags, str): + tags = [raw_tags] + elif isinstance(raw_tags, list): + tags = [str(item) for item in raw_tags] + else: + tags = [str(raw_tags)] + if tag not in tags: + tags.append(tag) + metadata["tags"] = tags + return metadata + + def _metadata_from_content(content): + with tempfile.TemporaryDirectory() as tmp_dir: + path = Path(tmp_dir) / "README.md" + path.write_text(content or "", encoding="utf-8") + return metadata_load(path) or {{}} + + def _content_with_metadata(content, metadata): + with tempfile.TemporaryDirectory() as tmp_dir: + path = Path(tmp_dir) / "README.md" + path.write_text(content or "", encoding="utf-8") + metadata_save(path, metadata) + return path.read_text(encoding="utf-8") + + def _body_without_metadata(content): + return front_matter_re.sub("", content or "", count=1).strip() + + def _append_section(content, section): + base = (content or "").rstrip() + if base: + return base + "\\n\\n" + section.strip() + "\\n" + return section.strip() + "\\n" + + def _provenance(repo_type): + label = {{"model": "model", "dataset": "dataset"}}.get( + repo_type, "Hub" + ) + return ( + marker + + "\\n## Generated by ML Intern\\n\\n" + + f"This {{label}} repository was generated by [ML Intern](https://github.com/huggingface/ml-intern), an agent for machine learning research and development on the Hugging Face Hub.\\n\\n" + + "- Try ML Intern: https://smolagents-ml-intern.hf.space\\n" + + "- Source code: https://github.com/huggingface/ml-intern\\n" + ) + + def _usage(repo_id, repo_type): + if repo_type == "dataset": + return ( + "## Usage\\n\\n" + "```python\\n" + "from datasets import load_dataset\\n\\n" + f"dataset = load_dataset({{repo_id!r}})\\n" + "```\\n" + ) + return ( + "## Usage\\n\\n" + "```python\\n" + "from transformers import AutoModelForCausalLM, AutoTokenizer\\n\\n" + f"model_id = {{repo_id!r}}\\n" + "tokenizer = AutoTokenizer.from_pretrained(model_id)\\n" + "model = AutoModelForCausalLM.from_pretrained(model_id)\\n" + "```\\n\\n" + "For non-causal architectures, replace `AutoModelForCausalLM` with the appropriate `AutoModel` class.\\n" + ) + + def _augment(content, repo_id, repo_type, extra_metadata=None): + metadata = _metadata_from_content(content or "") + if extra_metadata: + metadata = {{**extra_metadata, **metadata}} + updated = _content_with_metadata(content or "", _merge_tags(metadata)) + if not _body_without_metadata(updated): + updated = _append_section(updated, f"# {{repo_id}}") + if repo_type in {{"model", "dataset"}} and marker not in updated: + updated = _append_section(updated, _provenance(repo_type)) + if not usage_re.search(content or ""): + updated = _append_section(updated, _usage(repo_id, repo_type)) + return updated + + def _readme(api, repo_id, repo_type, token_value): + try: + path = hf_hub_download( + repo_id=repo_id, + filename="README.md", + repo_type=repo_type, + token=token_value, + ) + except (EntryNotFoundError, RepositoryNotFoundError): + return "" + return Path(path).read_text(encoding="utf-8") + + def _ensure_collection(api, token_value): + nonlocal collection_slug + if collection_slug: + return collection_slug + try: + cached_slug = Path(collection_cache_path).read_text( + encoding="utf-8" + ).strip() + if cached_slug: + collection_slug = cached_slug + return collection_slug + except Exception: + pass + collection = api.create_collection( + title=collection_title, + description=( + f"Artifacts generated by ML Intern session {{session_id}} " + f"on {{session_date}}." + ), + private=True, + exists_ok=True, + token=token_value, + ) + collection_slug = getattr(collection, "slug", None) + if collection_slug: + try: + cache_path = Path(collection_cache_path) + cache_path.parent.mkdir(parents=True, exist_ok=True) + cache_path.write_text(collection_slug, encoding="utf-8") + except Exception: + pass + return collection_slug + + def _register( + repo_id, + repo_type="model", + token_value=None, + extra_metadata=None, + force=False, + ): + nonlocal registering + if registering or not repo_id: + return + repo_type = repo_type or "model" + if repo_type not in supported: + return + if _is_sandbox_repo(repo_id, repo_type): + return + key = f"{{repo_type}}:{{repo_id}}" + if key in registered and not force: + return + registering = True + try: + token_value = _token(token_value) + api = HfApi(token=token_value) + card_updated = False + try: + current = _readme(api, repo_id, repo_type, token_value) + updated = _augment( + current, repo_id, repo_type, extra_metadata=extra_metadata + ) + if updated != current: + _original_upload_file( + api, + path_or_fileobj=updated.encode("utf-8"), + path_in_repo="README.md", + repo_id=repo_id, + repo_type=repo_type, + token=token_value, + commit_message="Update ML Intern artifact metadata", + ) + card_updated = True + except Exception: + pass + collection_updated = False + try: + slug = _ensure_collection(api, token_value) + if slug: + api.add_collection_item( + collection_slug=slug, + item_id=repo_id, + item_type=repo_type, + note=( + f"Generated by ML Intern session {{session_id}} " + f"on {{session_date}}." + ), + exists_ok=True, + token=token_value, + ) + collection_updated = True + except Exception: + pass + if card_updated and collection_updated: + registered.add(key) + finally: + registering = False + + _original_create_repo = HfApi.create_repo + _original_upload_file = HfApi.upload_file + _original_upload_folder = getattr(HfApi, "upload_folder", None) + _original_create_commit = getattr(HfApi, "create_commit", None) + + def _repo_id(args, kwargs): + return kwargs.get("repo_id") or (args[0] if args else None) + + def _repo_type(kwargs): + return kwargs.get("repo_type") or "model" + + def _is_sandbox_repo(repo_id, repo_type): + if (repo_type or "model") != "space" or not repo_id: + return False + repo_name = str(repo_id).rsplit("/", 1)[-1] + return bool(sandbox_space_re.fullmatch(repo_name)) + + def _patched_create_repo(self, *args, **kwargs): + result = _original_create_repo(self, *args, **kwargs) + repo_id = _repo_id(args, kwargs) + repo_type = _repo_type(kwargs) + extra = None + if repo_type == "space" and kwargs.get("space_sdk"): + extra = {{"sdk": kwargs.get("space_sdk")}} + _register(repo_id, repo_type, _token(kwargs.get("token"), self), extra) + return result + + def _patched_upload_file(self, *args, **kwargs): + result = _original_upload_file(self, *args, **kwargs) + if not kwargs.get("create_pr"): + force = kwargs.get("path_in_repo") == "README.md" + _register( + kwargs.get("repo_id"), + _repo_type(kwargs), + _token(kwargs.get("token"), self), + force=force, + ) + return result + + def _patched_upload_folder(self, *args, **kwargs): + result = _original_upload_folder(self, *args, **kwargs) + if not kwargs.get("create_pr"): + _register( + kwargs.get("repo_id"), + _repo_type(kwargs), + _token(kwargs.get("token"), self), + force=True, + ) + return result + + def _patched_create_commit(self, *args, **kwargs): + result = _original_create_commit(self, *args, **kwargs) + if not kwargs.get("create_pr"): + _register( + _repo_id(args, kwargs), + _repo_type(kwargs), + _token(kwargs.get("token"), self), + force=True, + ) + return result + + HfApi.create_repo = _patched_create_repo + HfApi.upload_file = _patched_upload_file + if _original_upload_folder is not None: + HfApi.upload_folder = _patched_upload_folder + if _original_create_commit is not None: + HfApi.create_commit = _patched_create_commit + + def _patch_module_func(name, method_name): + original = getattr(_hub, name, None) + if original is None: + return + method = getattr(HfApi, method_name) + + def _patched(*args, **kwargs): + api = HfApi(token=_token(kwargs.get("token"))) + return method(api, *args, **kwargs) + + setattr(_hub, name, _patched) + + _patch_module_func("create_repo", "create_repo") + _patch_module_func("upload_file", "upload_file") + if _original_upload_folder is not None: + _patch_module_func("upload_folder", "upload_folder") + if _original_create_commit is not None: + _patch_module_func("create_commit", "create_commit") + + try: + _install_ml_intern_artifact_hooks() + except Exception: + pass + """ + ).strip() + + "\n" + ) + + +def wrap_shell_command_with_hub_artifact_bootstrap( + command: str, + session: Any, +) -> str: + """Prefix a shell command so child Python processes load Hub hooks.""" + sitecustomize = build_hub_artifact_sitecustomize(session) + if not sitecustomize or not command: + return command + + encoded = base64.b64encode(sitecustomize.encode("utf-8")).decode("ascii") + bootstrap = ( + '_ml_intern_artifacts_dir="$(mktemp -d 2>/dev/null)" ' + f"&& printf %s {shlex.quote(encoded)} | base64 -d " + '> "$_ml_intern_artifacts_dir/sitecustomize.py" ' + '&& export PYTHONPATH="$_ml_intern_artifacts_dir${PYTHONPATH:+:$PYTHONPATH}"' + ) + return f"{bootstrap}; {command}" diff --git a/agent/core/llm_params.py b/agent/core/llm_params.py index d6843df10..d2f821c2b 100644 --- a/agent/core/llm_params.py +++ b/agent/core/llm_params.py @@ -7,75 +7,38 @@ import os - -def _patch_litellm_effort_validation() -> None: - """Neuter LiteLLM 1.83's hardcoded effort-level validation. - - Context: at ``litellm/llms/anthropic/chat/transformation.py:~1443`` the - Anthropic adapter validates ``output_config.effort ∈ {high, medium, - low, max}`` and gates ``max`` behind an ``_is_opus_4_6_model`` check - that only matches the substring ``opus-4-6`` / ``opus_4_6``. Result: - - * ``xhigh`` — valid on Anthropic's real API for Claude 4.7 — is - rejected pre-flight with "Invalid effort value: xhigh". - * ``max`` on Opus 4.7 is rejected with "effort='max' is only supported - by Claude Opus 4.6", even though Opus 4.7 accepts it in practice. - - We don't want to maintain a parallel model table, so we let the - Anthropic API itself be the validator: widen ``_is_opus_4_6_model`` - to also match ``opus-4-7``+ families, and drop the valid-effort-set - check entirely. If Anthropic rejects an effort level, we see a 400 - and the cascade walks down — exactly the behavior we want for any - future model family. - - Removable once litellm ships 1.83.8-stable (which merges PR #25867, - "Litellm day 0 opus 4.7 support") — see commit 0868a82 on their main - branch. Until then, this one-time patch is the escape hatch. - """ - try: - from litellm.llms.anthropic.chat import transformation as _t - except Exception: - return - - cfg = getattr(_t, "AnthropicConfig", None) - if cfg is None: - return - - original = getattr(cfg, "_is_opus_4_6_model", None) - if original is None or getattr(original, "_hf_agent_patched", False): - return - - def _widened(model: str) -> bool: - m = model.lower() - # Original 4.6 match plus any future Opus >= 4.6. We only need this - # to return True for families where "max" / "xhigh" are acceptable - # at the API; the cascade handles the case when they're not. - return any( - v in m for v in ( - "opus-4-6", "opus_4_6", "opus-4.6", "opus_4.6", - "opus-4-7", "opus_4_7", "opus-4.7", "opus_4.7", - ) - ) - - _widened._hf_agent_patched = True # type: ignore[attr-defined] - cfg._is_opus_4_6_model = staticmethod(_widened) - - -_patch_litellm_effort_validation() +from agent.core.hf_tokens import resolve_hf_router_token +from agent.core.local_models import ( + LOCAL_MODEL_API_KEY_DEFAULT, + LOCAL_MODEL_API_KEY_ENV, + LOCAL_MODEL_BASE_URL_ENV, + is_reserved_local_model_id, + local_model_name, + local_model_provider, +) +from agent.core.model_ids import ( + HF_ROUTER_BASE_URL, + strip_huggingface_model_prefix, +) + + +def _resolve_hf_router_token(session_hf_token: str | None = None) -> str | None: + """Backward-compatible private wrapper used by tests and older imports.""" + return resolve_hf_router_token(session_hf_token) # Effort levels accepted on the wire. -# Anthropic (4.6+): low | medium | high | xhigh | max (output_config.effort) -# OpenAI direct: minimal | low | medium | high (reasoning_effort top-level) -# HF router: low | medium | high (extra_body.reasoning_effort) -# -# We validate *shape* here and let the probe cascade walk down on rejection; -# we deliberately do NOT maintain a per-model capability table. -_ANTHROPIC_EFFORTS = {"low", "medium", "high", "xhigh", "max"} -_OPENAI_EFFORTS = {"minimal", "low", "medium", "high"} +# HF Router exposes reasoning controls through the OpenAI-compatible +# ``extra_body`` field. The probe cascade walks down when a provider rejects +# an accepted-looking value, so this stays intentionally small and generic. _HF_EFFORTS = {"low", "medium", "high"} +def _hf_router_effort_level(reasoning_effort: str) -> str: + level = "low" if reasoning_effort == "minimal" else reasoning_effort + return level + + class UnsupportedEffortError(ValueError): """The requested effort isn't valid for this provider's API surface. @@ -84,6 +47,46 @@ class UnsupportedEffortError(ValueError): """ +def _local_api_base(base_url: str) -> str: + base = base_url.strip().rstrip("/") + if base.endswith("/v1"): + return base + return f"{base}/v1" + + +def _resolve_local_model_params( + model_name: str, + reasoning_effort: str | None = None, + strict: bool = False, +) -> dict: + if reasoning_effort and strict: + raise UnsupportedEffortError( + "Local OpenAI-compatible endpoints don't accept reasoning_effort" + ) + + local_name = local_model_name(model_name) + if local_name is None: + raise ValueError(f"Unsupported local model id: {model_name}") + + provider = local_model_provider(model_name) + assert provider is not None + raw_base = ( + os.environ.get(provider["base_url_env"]) + or os.environ.get(LOCAL_MODEL_BASE_URL_ENV) + or provider["base_url_default"] + ) + api_key = ( + os.environ.get(provider["api_key_env"]) + or os.environ.get(LOCAL_MODEL_API_KEY_ENV) + or LOCAL_MODEL_API_KEY_DEFAULT + ) + return { + "model": f"openai/{local_name}", + "api_base": _local_api_base(raw_base), + "api_key": api_key, + } + + def _resolve_llm_params( model_name: str, session_hf_token: str | None = None, @@ -93,30 +96,18 @@ def _resolve_llm_params( """ Build LiteLLM kwargs for a given model id. - • ``anthropic/`` — native thinking config. We bypass LiteLLM's - ``reasoning_effort`` → ``thinking`` mapping (which lags new Claude - releases like 4.7 and sends the wrong API shape). Instead we pass - both ``thinking={"type": "adaptive"}`` and ``output_config= - {"effort": }`` as top-level kwargs — LiteLLM's Anthropic - adapter forwards unknown top-level kwargs into the request body - verbatim (confirmed by live probe; ``extra_body`` does NOT work - here because Anthropic's API rejects it as "Extra inputs are not - permitted"). This is the stable API for 4.6 and 4.7. Older - extended-thinking models that only accept ``thinking.type.enabled`` - will reject this; the probe's cascade catches that and falls back - to no thinking. - - • ``openai/`` — ``reasoning_effort`` forwarded as a top-level - kwarg (GPT-5 / o-series). LiteLLM uses the user's ``OPENAI_API_KEY``. - - • Anything else is treated as a HuggingFace router id. We hit the - auto-routing OpenAI-compatible endpoint at - ``https://router.huggingface.co/v1``. The id can be bare or carry an - HF routing suffix (``:fastest`` / ``:cheapest`` / ``:``). - A leading ``huggingface/`` is stripped. ``reasoning_effort`` is - forwarded via ``extra_body`` (LiteLLM's OpenAI adapter refuses it as - a top-level kwarg for non-OpenAI models). "minimal" normalizes to - "low". + • ``ollama/``, ``vllm/``, ``lm_studio/``, and + ``llamacpp/`` — local OpenAI-compatible endpoints. The id prefix + selects a configurable localhost base URL, and the model suffix is sent + to LiteLLM as ``openai/``. These endpoints don't receive + ``reasoning_effort``. + + • Anything else is treated as an HF Router id. We hit the auto-routing + OpenAI-compatible endpoint at ``https://router.huggingface.co/v1``. + The id can be bare or carry an HF routing suffix (``:fastest`` / + ``:cheapest`` / ``:``). A leading ``huggingface/`` is + stripped. ``reasoning_effort`` is forwarded via ``extra_body``. + "minimal" normalizes to "low". ``strict=True`` raises ``UnsupportedEffortError`` when the requested effort isn't in the provider's accepted set, instead of silently @@ -125,75 +116,32 @@ def _resolve_llm_params( runtime callers leave ``strict=False``, so a stale cached effort can't crash a turn — it just doesn't get sent. - Token precedence (first non-empty wins): - 1. INFERENCE_TOKEN env — shared key on the hosted Space (inference is - free for users, billed to the Space owner via ``X-HF-Bill-To``). - 2. session.hf_token — the user's own token (CLI / OAuth / cache file). - 3. HF_TOKEN env — belt-and-suspenders fallback for CLI users. + Token precedence for HF-router calls (first non-empty wins): + 1. session.hf_token — the user's own token (CLI / OAuth / cache file). + 2. huggingface_hub cache — ``HF_TOKEN`` / ``HUGGING_FACE_HUB_TOKEN`` / + local ``hf auth login`` cache. """ - if model_name.startswith("anthropic/"): - params: dict = {"model": model_name} - if reasoning_effort: - level = reasoning_effort - if level == "minimal": - level = "low" - if level not in _ANTHROPIC_EFFORTS: - if strict: - raise UnsupportedEffortError( - f"Anthropic doesn't accept effort={level!r}" - ) - else: - # Adaptive thinking + output_config.effort is the stable - # Anthropic API for Claude 4.6 / 4.7. Both kwargs are - # passed top-level: LiteLLM forwards unknown params into - # the request body for Anthropic, so ``output_config`` - # reaches the API. ``extra_body`` does NOT work here — - # Anthropic rejects it as "Extra inputs are not - # permitted". - params["thinking"] = {"type": "adaptive"} - params["output_config"] = {"effort": level} - return params - - if model_name.startswith("bedrock/"): - # LiteLLM routes ``bedrock/...`` through the Converse adapter, which - # picks up AWS credentials from the standard env vars - # (``AWS_ACCESS_KEY_ID`` / ``AWS_SECRET_ACCESS_KEY`` / ``AWS_REGION``). - # The Anthropic thinking/effort shape is not forwarded through Converse - # the same way, so we leave it off for now. - return {"model": model_name} - - if model_name.startswith("openai/"): - params = {"model": model_name} - if reasoning_effort: - if reasoning_effort not in _OPENAI_EFFORTS: - if strict: - raise UnsupportedEffortError( - f"OpenAI doesn't accept effort={reasoning_effort!r}" - ) - else: - params["reasoning_effort"] = reasoning_effort - return params - - hf_model = model_name.removeprefix("huggingface/") - api_key = ( - os.environ.get("INFERENCE_TOKEN") - or session_hf_token - or os.environ.get("HF_TOKEN") - ) + normalized_model = strip_huggingface_model_prefix(model_name) or model_name + + if is_reserved_local_model_id(normalized_model): + raise ValueError(f"Unsupported local model id: {normalized_model}") + + if local_model_provider(normalized_model) is not None: + return _resolve_local_model_params(normalized_model, reasoning_effort, strict) + + hf_model = normalized_model + api_key = _resolve_hf_router_token(session_hf_token) params = { "model": f"openai/{hf_model}", - "api_base": "https://router.huggingface.co/v1", + "api_base": HF_ROUTER_BASE_URL, "api_key": api_key, } - if os.environ.get("INFERENCE_TOKEN"): - bill_to = os.environ.get("HF_BILL_TO", "smolagents") - params["extra_headers"] = {"X-HF-Bill-To": bill_to} if reasoning_effort: - hf_level = "low" if reasoning_effort == "minimal" else reasoning_effort + hf_level = _hf_router_effort_level(reasoning_effort) if hf_level not in _HF_EFFORTS: if strict: raise UnsupportedEffortError( - f"HF router doesn't accept effort={hf_level!r}" + f"HF Router doesn't accept effort={hf_level!r}" ) else: params["extra_body"] = {"reasoning_effort": hf_level} diff --git a/agent/core/local_models.py b/agent/core/local_models.py new file mode 100644 index 000000000..9f8a9491d --- /dev/null +++ b/agent/core/local_models.py @@ -0,0 +1,59 @@ +"""Helpers for CLI local OpenAI-compatible model ids.""" + +LOCAL_MODEL_PROVIDERS: dict[str, dict[str, str]] = { + "ollama/": { + "base_url_env": "OLLAMA_BASE_URL", + "base_url_default": "http://localhost:11434", + "api_key_env": "OLLAMA_API_KEY", + }, + "vllm/": { + "base_url_env": "VLLM_BASE_URL", + "base_url_default": "http://localhost:8000", + "api_key_env": "VLLM_API_KEY", + }, + "lm_studio/": { + "base_url_env": "LMSTUDIO_BASE_URL", + "base_url_default": "http://127.0.0.1:1234", + "api_key_env": "LMSTUDIO_API_KEY", + }, + "llamacpp/": { + "base_url_env": "LLAMACPP_BASE_URL", + "base_url_default": "http://localhost:8080", + "api_key_env": "LLAMACPP_API_KEY", + }, +} + +LOCAL_MODEL_PREFIXES = tuple(LOCAL_MODEL_PROVIDERS) +RESERVED_LOCAL_MODEL_PREFIXES = ("openai-compat/",) +LOCAL_MODEL_BASE_URL_ENV = "LOCAL_LLM_BASE_URL" +LOCAL_MODEL_API_KEY_ENV = "LOCAL_LLM_API_KEY" +LOCAL_MODEL_API_KEY_DEFAULT = "sk-local-no-key-required" + + +def local_model_provider(model_id: str) -> dict[str, str] | None: + """Return provider config for a local model id, if it uses a local prefix.""" + for prefix, config in LOCAL_MODEL_PROVIDERS.items(): + if model_id.startswith(prefix): + return config + return None + + +def local_model_name(model_id: str) -> str | None: + """Return the backend model name with the local provider prefix removed.""" + for prefix in LOCAL_MODEL_PREFIXES: + if model_id.startswith(prefix): + name = model_id[len(prefix) :] + return name or None + return None + + +def is_local_model_id(model_id: str) -> bool: + """Return True for non-empty, whitespace-free local model ids.""" + if not model_id or any(char.isspace() for char in model_id): + return False + return local_model_name(model_id) is not None + + +def is_reserved_local_model_id(model_id: str) -> bool: + """Return True for local-style prefixes intentionally not supported.""" + return model_id.startswith(RESERVED_LOCAL_MODEL_PREFIXES) diff --git a/agent/core/model_ids.py b/agent/core/model_ids.py new file mode 100644 index 000000000..778a269e4 --- /dev/null +++ b/agent/core/model_ids.py @@ -0,0 +1,32 @@ +"""Canonical model ids for HF Router inference.""" + +HF_ROUTER_BASE_URL = "https://router.huggingface.co/v1" + +# Keep these as verbatim HF Router ids; version punctuation differs by model. +CLAUDE_OPUS_48_MODEL_ID = "anthropic/claude-opus-4.8:fal-ai" +GPT_55_MODEL_ID = "openai/gpt-5.5:fal-ai" +KIMI_K27_CODE_MODEL_ID = "moonshotai/Kimi-K2.7-Code:novita" +MINIMAX_M3_MODEL_ID = "MiniMaxAI/MiniMax-M3:novita" +GLM_52_MODEL_ID = "zai-org/GLM-5.2:novita" +DEEPSEEK_V4_PRO_MODEL_ID = "deepseek-ai/DeepSeek-V4-Pro:novita" + +HOSTED_MODEL_IDS = { + CLAUDE_OPUS_48_MODEL_ID, + GPT_55_MODEL_ID, + KIMI_K27_CODE_MODEL_ID, + MINIMAX_M3_MODEL_ID, + GLM_52_MODEL_ID, + DEEPSEEK_V4_PRO_MODEL_ID, +} + + +def strip_huggingface_model_prefix(model_id: str | None) -> str | None: + """Return model ids without LiteLLM's optional ``huggingface/`` prefix.""" + if not model_id: + return model_id + return model_id.removeprefix("huggingface/") + + +def is_known_router_model_id(model_id: str | None) -> bool: + normalized = strip_huggingface_model_prefix(model_id) + return bool(normalized and normalized in HOSTED_MODEL_IDS) diff --git a/agent/core/model_switcher.py b/agent/core/model_switcher.py index afb8d52c6..5ece764d8 100644 --- a/agent/core/model_switcher.py +++ b/agent/core/model_switcher.py @@ -15,41 +15,69 @@ from __future__ import annotations +import asyncio + +from litellm import acompletion + from agent.core.effort_probe import ProbeInconclusive, probe_effort +from agent.core.llm_params import _resolve_llm_params +from agent.core.local_models import ( + LOCAL_MODEL_PREFIXES, + is_local_model_id, + is_reserved_local_model_id, +) +from agent.core.model_ids import ( + CLAUDE_OPUS_48_MODEL_ID, + DEEPSEEK_V4_PRO_MODEL_ID, + GLM_52_MODEL_ID, + GPT_55_MODEL_ID, + KIMI_K27_CODE_MODEL_ID, + MINIMAX_M3_MODEL_ID, + strip_huggingface_model_prefix, +) # Suggested models shown by `/model` (not a gate). Users can paste any HF -# model id (e.g. "MiniMaxAI/MiniMax-M2.7") or an `anthropic/` / `openai/` -# prefix for direct API access. For HF ids, append ":fastest" / -# ":cheapest" / ":preferred" / ":" to override the default -# routing policy (auto = fastest with failover). +# Router model id (e.g. "MiniMaxAI/MiniMax-M3:novita"). Append ":fastest", +# ":cheapest", ":preferred", or ":" to override the default routing +# policy (auto = fastest with failover). SUGGESTED_MODELS = [ - {"id": "bedrock/us.anthropic.claude-opus-4-7", "label": "Claude Opus 4.7"}, - {"id": "bedrock/us.anthropic.claude-opus-4-6-v1", "label": "Claude Opus 4.6"}, - {"id": "MiniMaxAI/MiniMax-M2.7", "label": "MiniMax M2.7"}, - {"id": "moonshotai/Kimi-K2.6", "label": "Kimi K2.6"}, - {"id": "zai-org/GLM-5.1", "label": "GLM 5.1"}, + {"id": CLAUDE_OPUS_48_MODEL_ID, "label": "Claude Opus 4.8"}, + {"id": GPT_55_MODEL_ID, "label": "GPT-5.5"}, + {"id": MINIMAX_M3_MODEL_ID, "label": "MiniMax M3"}, + {"id": KIMI_K27_CODE_MODEL_ID, "label": "Kimi K2.7 Code"}, + {"id": GLM_52_MODEL_ID, "label": "GLM 5.2"}, + {"id": DEEPSEEK_V4_PRO_MODEL_ID, "label": "DeepSeek V4 Pro"}, ] _ROUTING_POLICIES = {"fastest", "cheapest", "preferred"} +_LOCAL_PROBE_TIMEOUT = 15.0 def is_valid_model_id(model_id: str) -> bool: """Loose format check — lets users pick any model id. Accepts: - • anthropic/ - • openai/ + • ollama/, vllm/, lm_studio/, llamacpp/ • /[:] (HF router; tag = provider or policy) - • huggingface//[:] (same, accepts legacy prefix) + • huggingface//[:] (same, optional LiteLLM prefix) Actual availability is verified against the HF router catalog on switch, and by the provider on the probe's ping call. """ - if not model_id or "/" not in model_id: + if not model_id: + return False + normalized_model_id = strip_huggingface_model_prefix(model_id) or model_id + if is_local_model_id(normalized_model_id): + return True + if is_reserved_local_model_id(normalized_model_id): + return False + if any(normalized_model_id.startswith(prefix) for prefix in LOCAL_MODEL_PREFIXES): return False - head = model_id.split(":", 1)[0] + if "/" not in normalized_model_id: + return False + head = normalized_model_id.split(":", 1)[0] parts = head.split("/") return len(parts) >= 2 and all(parts) @@ -60,10 +88,11 @@ def _print_hf_routing_info(model_id: str, console) -> bool: proceed with the switch, ``False`` to indicate a hard problem the user should notice before we fire the effort probe. - Anthropic / OpenAI ids return ``True`` without printing anything — - the probe below covers "does this model exist". + Local ids return ``True`` without printing anything. Router ids are checked + against the router catalog when possible; the probe below covers provider + availability for uncataloged ids. """ - if model_id.startswith(("anthropic/", "openai/")): + if is_local_model_id(model_id): return True from agent.core import hf_router_catalog as cat @@ -118,9 +147,7 @@ def _print_hf_routing_info(model_id: str, console) -> bool: ) ctx = f"{p.context_length:,} ctx" if p.context_length else "ctx n/a" tools = "tools" if p.supports_tools else "no tools" - console.print( - f" [dim]{p.provider}: {price}, {ctx}, {tools}[/dim]" - ) + console.print(f" [dim]{p.provider}: {price}, {ctx}, {tools}[/dim]") return True @@ -134,9 +161,10 @@ def print_model_listing(config, console) -> None: marker = " [dim]<-- current[/dim]" if m["id"] == current else "" console.print(f" {m['id']} [dim]({m['label']})[/dim]{marker}") console.print( - "\n[dim]Paste any HF model id (e.g. 'MiniMaxAI/MiniMax-M2.7').\n" + "\n[dim]Paste any HF model id (e.g. 'MiniMaxAI/MiniMax-M3:novita').\n" "Add ':fastest', ':cheapest', ':preferred', or ':' to override routing.\n" - "Use 'anthropic/' or 'openai/' for direct API access.[/dim]" + "Use 'ollama/', 'vllm/', 'lm_studio/', or " + "'llamacpp/' for local OpenAI-compatible endpoints.[/dim]" ) @@ -145,8 +173,20 @@ def print_invalid_id(arg: str, console) -> None: console.print( "[dim]Expected:\n" " • /[:tag] (HF router — paste from huggingface.co)\n" - " • anthropic/\n" - " • openai/[/dim]" + " • ollama/ | vllm/ | lm_studio/ | llamacpp/[/dim]" + ) + + +async def _probe_local_model(model_id: str) -> None: + params = _resolve_llm_params(model_id) + await asyncio.wait_for( + acompletion( + messages=[{"role": "user", "content": "ping"}], + max_tokens=1, + stream=False, + **params, + ), + timeout=_LOCAL_PROBE_TIMEOUT, ) @@ -168,9 +208,26 @@ async def probe_and_switch_model( * ✗ hard error (auth, model-not-found, quota) — we reject the switch and keep the current model so the user isn't stranded - Transient errors (5xx, timeout) complete the switch with a yellow - warning; the next real call re-surfaces the error if it's persistent. + For non-local models, transient errors (5xx, timeout) complete the switch + with a yellow warning; the next real call re-surfaces the error if it's + persistent. Local models reject every probe error, including timeouts, and + keep the current model. """ + if is_local_model_id(model_id): + console.print(f"[dim]checking local model {model_id}...[/dim]") + try: + await _probe_local_model(model_id) + except Exception as e: + console.print(f"[bold red]Switch failed:[/bold red] {e}") + console.print(f"[dim]Keeping current model: {config.model_name}[/dim]") + return + + _commit_switch(model_id, config, session, effective=None, cache=True) + console.print( + f"[green]Model switched to {model_id}[/green] [dim](effort: off)[/dim]" + ) + return + preference = config.reasoning_effort if not _print_hf_routing_info(model_id, console): return @@ -179,12 +236,14 @@ async def probe_and_switch_model( # Nothing to validate with a ping that we couldn't validate on the # first real call just as cheaply. Skip the probe entirely. _commit_switch(model_id, config, session, effective=None, cache=False) - console.print(f"[green]Model switched to {model_id}[/green] [dim](effort: off)[/dim]") + console.print( + f"[green]Model switched to {model_id}[/green] [dim](effort: off)[/dim]" + ) return console.print(f"[dim]checking {model_id} (effort: {preference})...[/dim]") try: - outcome = await probe_effort(model_id, preference, hf_token) + outcome = await probe_effort(model_id, preference, hf_token, session=session) except ProbeInconclusive as e: _commit_switch(model_id, config, session, effective=None, cache=False) console.print( @@ -199,8 +258,11 @@ async def probe_and_switch_model( return _commit_switch( - model_id, config, session, - effective=outcome.effective_effort, cache=True, + model_id, + config, + session, + effective=outcome.effective_effort, + cache=True, ) effort_label = outcome.effective_effort or "off" suffix = f" — {outcome.note}" if outcome.note else "" diff --git a/agent/core/prompt_caching.py b/agent/core/prompt_caching.py index 56685a5ef..f445a2d37 100644 --- a/agent/core/prompt_caching.py +++ b/agent/core/prompt_caching.py @@ -1,59 +1,219 @@ -"""Anthropic prompt caching breakpoints for outgoing LLM requests. +"""Prompt-cache helpers for HF Router FAL requests. -Caching is GA on Anthropic's API and natively supported by litellm >=1.83 -via ``cache_control`` blocks. We apply two breakpoints (out of 4 allowed): +The HF Router/OpenRouter path uses provider-native prompt caching. Anthropic +models keep explicit JSON ``cache_control`` content blocks for compatibility, +and also need the top-level ``cache_control`` hint on the OpenAI-compatible HF +Router path; the explicit markers alone are accepted there but do not produce +cache writes. OpenAI models cache eligible prefixes automatically and accept +routing/retention hints in the body. +Headers like ``X-OpenRouter-Cache`` control response caching, not prompt +caching through this route. +""" - 1. The tool block — caches all tool definitions as a single prefix. - 2. The system message — caches the rendered system prompt. +from typing import Any -Together these cover the ~4-5K static tokens that were being re-billed on -every turn. Subsequent turns within the 5-minute TTL hit cache_read pricing -(~10% of input cost) instead of full input. +from agent.core.model_ids import HF_ROUTER_BASE_URL -Non-Anthropic models (HF router, OpenAI) are passed through unchanged. -""" +_CACHE_CONTROL = {"type": "ephemeral"} +_CACHEABLE_ROLES = {"system", "user"} +_HF_ROUTER_SESSION_ID_MAX_LENGTH = 256 +HF_ROUTER_SESSION_ID_HEADER = "X-HF-Session-id" -from typing import Any + +def router_session_id_for(session: Any) -> str | None: + """Return the usage-window-scoped Router session ID for a runtime session.""" + billing_session_id = getattr(session, "inference_billing_session_id", None) + if isinstance(billing_session_id, str) and billing_session_id: + return billing_session_id + session_id = getattr(session, "session_id", None) + if isinstance(session_id, str) and session_id: + return session_id + return None + + +def _is_hf_router_request(llm_params: dict[str, Any]) -> bool: + api_base = str(llm_params.get("api_base") or "").rstrip("/") + return api_base == HF_ROUTER_BASE_URL + + +def _is_fal_router_request(llm_params: dict[str, Any]) -> bool: + return _is_hf_router_request(llm_params) and ":fal" in _router_model(llm_params) + + +def _router_model(llm_params: dict[str, Any]) -> str: + model = str(llm_params.get("model") or "") + return model.removeprefix("openai/") + + +def _uses_explicit_cache_control(llm_params: dict[str, Any]) -> bool: + if not _is_fal_router_request(llm_params): + return False + return _router_model(llm_params).startswith("anthropic/") + + +def _is_openai_gpt55(llm_params: dict[str, Any]) -> bool: + if not _is_fal_router_request(llm_params): + return False + return _router_model(llm_params).startswith("openai/gpt-5.5") + + +def _merge_extra_body( + llm_params: dict[str, Any], updates: dict[str, Any] +) -> dict[str, Any]: + if not updates: + return llm_params + + cached_params = dict(llm_params) + extra_body = dict(cached_params.get("extra_body") or {}) + extra_body.update(updates) + cached_params["extra_body"] = extra_body + return cached_params + + +def _merge_extra_headers( + llm_params: dict[str, Any], updates: dict[str, str] +) -> dict[str, Any]: + if not updates: + return llm_params + + cached_params = dict(llm_params) + extra_headers = dict(cached_params.get("extra_headers") or {}) + extra_headers.update(updates) + cached_params["extra_headers"] = extra_headers + return cached_params + + +def with_prompt_cache_params( + llm_params: dict[str, Any], + *, + session_id: str | None = None, +) -> dict[str, Any]: + """Return LiteLLM params with provider-native prompt-cache body hints.""" + updates: dict[str, Any] = {} + headers: dict[str, str] = {} + if session_id and _is_hf_router_request(llm_params): + stable_session_id = session_id[:_HF_ROUTER_SESSION_ID_MAX_LENGTH] + headers[HF_ROUTER_SESSION_ID_HEADER] = stable_session_id + if _is_openai_gpt55(llm_params): + updates["prompt_cache_key"] = stable_session_id + + if _uses_explicit_cache_control(llm_params): + updates["cache_control"] = dict(_CACHE_CONTROL) + + if _is_openai_gpt55(llm_params): + updates["prompt_cache_retention"] = "24h" + + return _merge_extra_headers(_merge_extra_body(llm_params, updates), headers) + + +def _message_role(message: Any) -> str | None: + if isinstance(message, dict): + role = message.get("role") + else: + role = getattr(message, "role", None) + return role if isinstance(role, str) else None + + +def _message_content(message: Any) -> Any: + if isinstance(message, dict): + return message.get("content") + return getattr(message, "content", None) + + +def _message_to_dict(message: Any) -> dict[str, Any]: + if isinstance(message, dict): + return dict(message) + if hasattr(message, "model_dump"): + return message.model_dump(exclude_none=True) + raise TypeError(f"Unsupported message type for prompt caching: {type(message)!r}") + + +def _has_cacheable_text(content: Any) -> bool: + if isinstance(content, str): + return bool(content) + if not isinstance(content, list): + return False + return any( + isinstance(block, dict) + and block.get("type") == "text" + and isinstance(block.get("text"), str) + and bool(block.get("text")) + for block in content + ) + + +def _cache_target_index(messages: list[Any]) -> int | None: + if len(messages) < 2: + return None + + for idx in range(len(messages) - 2, -1, -1): + message = messages[idx] + if _message_role(message) not in _CACHEABLE_ROLES: + continue + if _has_cacheable_text(_message_content(message)): + return idx + return None + + +def _content_with_cache_control(content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [ + {"type": "text", "text": content, "cache_control": dict(_CACHE_CONTROL)} + ] + + blocks = [dict(block) if isinstance(block, dict) else block for block in content] + for idx in range(len(blocks) - 1, -1, -1): + block = blocks[idx] + if ( + isinstance(block, dict) + and block.get("type") == "text" + and isinstance(block.get("text"), str) + and bool(block.get("text")) + ): + cached = dict(block) + cached["cache_control"] = dict(_CACHE_CONTROL) + blocks[idx] = cached + break + return blocks + + +def _tools_with_cache_control(tools: list[dict] | None) -> list[dict] | None: + if not tools: + return tools + + cached_tools = list(tools) + last_tool = dict(cached_tools[-1]) + last_tool["cache_control"] = dict(_CACHE_CONTROL) + cached_tools[-1] = last_tool + return cached_tools def with_prompt_caching( messages: list[Any], tools: list[dict] | None, - model_name: str | None, + llm_params: dict[str, Any], ) -> tuple[list[Any], list[dict] | None]: - """Return (messages, tools) with cache_control breakpoints for Anthropic. + """Return outgoing messages with explicit cache breakpoints when needed. - No-op for non-Anthropic models. Original objects are not mutated; a fresh - list with replaced first message and last tool is returned, so callers - that share the underlying ``ContextManager.items`` list don't see their - persisted history rewritten. + The newest message is treated as dynamic. For Anthropic FAL models, the + cache breakpoint is placed on the closest earlier system/user text block so + provider-side caching covers the stable prefix without changing persisted + conversation history. The final tool spec is also marked so stable tool + definitions are cached. """ - if not model_name or "anthropic" not in model_name: + if not _uses_explicit_cache_control(llm_params): return messages, tools - if tools: - new_tools = list(tools) - last = dict(new_tools[-1]) - last["cache_control"] = {"type": "ephemeral"} - new_tools[-1] = last - tools = new_tools - - if messages: - first = messages[0] - role = first.get("role") if isinstance(first, dict) else getattr(first, "role", None) - if role == "system": - content = ( - first.get("content") - if isinstance(first, dict) - else getattr(first, "content", None) - ) - if isinstance(content, str) and content: - cached_block = [{ - "type": "text", - "text": content, - "cache_control": {"type": "ephemeral"}, - }] - new_first = {"role": "system", "content": cached_block} - messages = [new_first] + list(messages[1:]) - - return messages, tools + cached_tools = _tools_with_cache_control(tools) + idx = _cache_target_index(messages) + if idx is None: + return messages, cached_tools + + cached_message = _message_to_dict(messages[idx]) + cached_message["content"] = _content_with_cache_control( + cached_message.get("content") + ) + + cached_messages = list(messages) + cached_messages[idx] = cached_message + return cached_messages, cached_tools diff --git a/agent/core/redact.py b/agent/core/redact.py new file mode 100644 index 000000000..a91bcb280 --- /dev/null +++ b/agent/core/redact.py @@ -0,0 +1,66 @@ +"""Secret scrubbing for session trajectories before upload. + +Users frequently paste HF / API / GitHub tokens into the chat, or scripts echo +them via env dumps. This module applies regex-based redaction to any string +value found recursively in a trajectory payload. The goal is best-effort — +strict formats are matched; we won't catch free-form leaks like "my password +is hunter2". +""" + +from __future__ import annotations + +import re +from typing import Any + +# Each entry: (compiled regex, replacement placeholder). +# Patterns are conservative: they only match tokens with the canonical prefix +# and a minimum body length so we don't paint over normal text. +_PATTERNS: list[tuple[re.Pattern, str]] = [ + # Hugging Face tokens: hf_[A-Za-z0-9]{30,} + (re.compile(r"hf_[A-Za-z0-9]{30,}"), "[REDACTED_HF_TOKEN]"), + # Provider API keys with common sk-* prefixes. + (re.compile(r"sk-ant-[A-Za-z0-9_\-]{20,}"), "[REDACTED_PROVIDER_API_KEY]"), + (re.compile(r"sk-(?!ant-)[A-Za-z0-9_\-]{40,}"), "[REDACTED_PROVIDER_API_KEY]"), + # GitHub classic PATs: ghp_, gho_, ghu_, ghs_, ghr_ followed by 36+ chars + (re.compile(r"gh[pousr]_[A-Za-z0-9]{36,}"), "[REDACTED_GITHUB_TOKEN]"), + # GitHub fine-grained PATs: github_pat_ + (re.compile(r"github_pat_[A-Za-z0-9_]{36,}"), "[REDACTED_GITHUB_TOKEN]"), + # AWS access key IDs: AKIA / ASIA + 16 uppercase alnum + (re.compile(r"\b(?:AKIA|ASIA)[A-Z0-9]{16}\b"), "[REDACTED_AWS_KEY_ID]"), + # Generic 'Bearer ' header values + (re.compile(r"(?i)bearer\s+[A-Za-z0-9_\-\.=]{20,}"), "Bearer [REDACTED]"), +] + +# Env-var-like exports: we scrub the value but keep the name so callers can +# still see which secret was referenced. Covers `KEY=value` and `KEY: value` +# when the key looks secret-y. +_SECRETY_NAMES = re.compile( + r"(?i)\b([A-Z0-9_]*(?:TOKEN|API_KEY|SECRET|PASSWORD|ACCESS_KEY_ID))" + r"\s*[:=]\s*([^\s\"']+)" +) + + +def scrub_string(s: str) -> str: + """Apply all redaction patterns to a single string. Safe on non-strings.""" + if not isinstance(s, str) or not s: + return s + out = s + for pat, repl in _PATTERNS: + out = pat.sub(repl, out) + out = _SECRETY_NAMES.sub(lambda m: f"{m.group(1)}=[REDACTED]", out) + return out + + +def scrub(obj: Any) -> Any: + """Recursively scrub every string value in a nested dict/list structure. + + Returns a new object — inputs are not mutated.""" + if isinstance(obj, str): + return scrub_string(obj) + if isinstance(obj, dict): + return {k: scrub(v) for k, v in obj.items()} + if isinstance(obj, list): + return [scrub(v) for v in obj] + if isinstance(obj, tuple): + return tuple(scrub(v) for v in obj) + return obj diff --git a/agent/core/session.py b/agent/core/session.py index 4b6390d84..3a7f50c39 100644 --- a/agent/core/session.py +++ b/agent/core/session.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import os import subprocess import sys import uuid @@ -10,23 +11,49 @@ from pathlib import Path from typing import Any, Optional +from litellm import Message + from agent.config import Config from agent.context_manager.manager import ContextManager +from agent.messaging.gateway import NotificationGateway +from agent.messaging.models import NotificationRequest +from agent.core.usage_thresholds import ( + USAGE_THRESHOLD_TOOL_NAME, + USAGE_WARNING_FIRST_THRESHOLD_USD, +) logger = logging.getLogger(__name__) _DEFAULT_MAX_TOKENS = 200_000 +_TURN_COMPLETE_NOTIFICATION_CHARS = 39000 + +DEFAULT_SESSION_LOG_DIR = Path("session_logs") + + +def _format_usd(value: Any) -> str: + if isinstance(value, bool): + return "$0.00" + try: + amount = float(value) + except (TypeError, ValueError): + amount = 0.0 + return f"${amount:.2f}" + + +def _approval_tools_are_usage_thresholds(tools: Any) -> bool: + if not isinstance(tools, list) or len(tools) != 1: + return False + tool = tools[0] + return isinstance(tool, dict) and tool.get("tool") == USAGE_THRESHOLD_TOOL_NAME def _get_max_tokens_safe(model_name: str) -> int: """Return the max input-context tokens for a model. - Primary source: ``litellm.get_model_info(model)['max_input_tokens']`` — - LiteLLM maintains an upstream catalog that knows Claude Opus 4.6 is - 1M, GPT-5 is 272k, Sonnet 4.5 is 200k, and so on. Strips any HF routing - suffix / huggingface/ prefix so tagged ids ('moonshotai/Kimi-K2.6:cheapest') - look up the bare model. Falls back to a conservative 200k default for - models not in the catalog (typically HF-router-only models). + Primary source: ``litellm.get_model_info(model)['max_input_tokens']``. + Strips any HF routing suffix / huggingface/ prefix so tagged ids + ('moonshotai/Kimi-K2.7-Code:novita') look up the bare model. Falls back to a + conservative 200k default for models not in the catalog. """ from litellm import get_model_info @@ -44,7 +71,8 @@ def _get_max_tokens_safe(model_name: str) -> int: continue logger.info( "No litellm.get_model_info entry for %s, falling back to %d", - model_name, _DEFAULT_MAX_TOKENS, + model_name, + _DEFAULT_MAX_TOKENS, ) return _DEFAULT_MAX_TOKENS @@ -52,9 +80,10 @@ def _get_max_tokens_safe(model_name: str) -> int: class OpType(Enum): USER_INPUT = "user_input" EXEC_APPROVAL = "exec_approval" - INTERRUPT = "interrupt" UNDO = "undo" COMPACT = "compact" + NEW = "new" + RESUME = "resume" SHUTDOWN = "shutdown" @@ -62,6 +91,7 @@ class OpType(Enum): class Event: event_type: str data: Optional[dict[str, Any]] = None + seq: Optional[int] = None class Session: @@ -73,16 +103,33 @@ class Session: def __init__( self, event_queue: asyncio.Queue, - config: Config | None = None, + config: Config, tool_router=None, context_manager: ContextManager | None = None, hf_token: str | None = None, local_mode: bool = False, + autonomous_mode: bool = False, stream: bool = True, + notification_gateway: NotificationGateway | None = None, + notification_destinations: list[str] | None = None, + defer_turn_complete_notification: bool = False, + session_id: str | None = None, + user_id: str | None = None, + hf_username: str | None = None, + user_plan: str | None = None, + persistence_store: Any | None = None, ): self.hf_token: Optional[str] = hf_token + self.user_id: Optional[str] = user_id + self.hf_username: Optional[str] = hf_username + self.user_plan: str | None = user_plan + self.local_mode = local_mode + self.autonomous_mode = autonomous_mode + self.persistence_store = persistence_store self.tool_router = tool_router self.stream = stream + if config is None: + raise ValueError("Session requires a Config") tool_specs = tool_router.get_tool_specs_for_llm() if tool_router else [] self.context_manager = context_manager or ContextManager( model_max_tokens=_get_max_tokens_safe(config.model_name), @@ -90,24 +137,47 @@ def __init__( untouched_messages=5, tool_specs=tool_specs, hf_token=hf_token, + hf_username=hf_username, local_mode=local_mode, + autonomous_mode=autonomous_mode, ) self.event_queue = event_queue - self.session_id = str(uuid.uuid4()) - self.config = config or Config( - model_name="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - ) + self.session_id = session_id or str(uuid.uuid4()) + self.inference_billing_session_id: str | None = None + self.config = config self.is_running = True + self.current_plan: list[dict[str, str]] = [] self._cancelled = asyncio.Event() self.pending_approval: Optional[dict[str, Any]] = None self.sandbox = None + self.sandbox_hardware: Optional[str] = None + self.sandbox_preload_task: Optional[asyncio.Task] = None + self.sandbox_preload_error: Optional[str] = None + self.sandbox_preload_cancel_event: Any | None = None self._running_job_ids: set[str] = set() # HF job IDs currently executing + self.notification_gateway = notification_gateway + self.notification_destinations = list(notification_destinations or []) + self.defer_turn_complete_notification = defer_turn_complete_notification + self.auto_approval_enabled: bool = False + self.auto_approval_cost_cap_usd: float | None = None + self.auto_approval_estimated_spend_usd: float = 0.0 + self._yolo_budget_reservations: dict[str, Any] = {} + self.usage_warning_next_threshold_usd: float = USAGE_WARNING_FIRST_THRESHOLD_USD + self.usage_threshold_checker: Any | None = None + self.yolo_budget_checker: Any | None = None + self.usage_hf_billing_snapshot: dict[str, Any] | None = None + self.usage_metrics: dict[str, Any] | None = None # Session trajectory logging self.logged_events: list[dict] = [] - self.session_start_time = datetime.now().isoformat() + self.session_start_time = datetime.now().astimezone().isoformat() self.turn_count: int = 0 self.last_auto_save_turn: int = 0 + # Stable local save path so heartbeat saves overwrite one file instead + # of spamming session_logs/. ``_last_heartbeat_ts`` is owned by + # ``agent.core.telemetry.HeartbeatSaver`` and lazily initialised there. + self._local_save_path: Optional[str] = None + self._last_heartbeat_ts: Optional[float] = None # Per-model probed reasoning-effort cache. Populated by the probe # on /model switch, read by ``effective_effort_for`` below. Keys are @@ -118,19 +188,176 @@ def __init__( # thinking params at all # Key absent → not probed yet; fall back to the raw preference. self.model_effective_effort: dict[str, str | None] = {} + self.context_manager.on_message_added = self._schedule_trace_message async def send_event(self, event: Event) -> None: """Send event back to client and log to trajectory""" - await self.event_queue.put(event) - # Log event to trajectory self.logged_events.append( { - "timestamp": datetime.now().isoformat(), + "timestamp": datetime.now().astimezone().isoformat(), "event_type": event.event_type, "data": event.data, } ) + if self.persistence_store is not None: + try: + event.seq = await self.persistence_store.append_event( + self.session_id, event.event_type, event.data + ) + except Exception as e: + logger.debug("Event persistence failed for %s: %s", self.session_id, e) + + await self.event_queue.put(event) + await self._enqueue_auto_notification_requests(event) + + # Mid-turn heartbeat flush (owned by telemetry module). + from agent.core.telemetry import HeartbeatSaver + + HeartbeatSaver.maybe_fire(self) + + def _schedule_trace_message(self, message: Any) -> None: + """Best-effort append-only trace save for SFT/KPI export.""" + if self.persistence_store is None: + return + try: + payload = message.model_dump(mode="json") + except Exception: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + source = str(payload.get("role") or "message") + loop.create_task( + self.persistence_store.append_trace_message( + self.session_id, payload, source=source + ) + ) + + def set_notification_destinations(self, destinations: list[str]) -> None: + """Replace the session's opted-in auto-notification destinations.""" + deduped: list[str] = [] + seen: set[str] = set() + for destination in destinations: + if destination not in seen: + deduped.append(destination) + seen.add(destination) + self.notification_destinations = deduped + + async def send_deferred_turn_complete_notification(self, event: Event) -> None: + if event.event_type != "turn_complete": + return + await self._enqueue_auto_notification_requests( + event, + include_deferred_turn_complete=True, + ) + + async def _enqueue_auto_notification_requests( + self, + event: Event, + include_deferred_turn_complete: bool = False, + ) -> None: + if self.notification_gateway is None: + return + if not self.notification_destinations: + return + auto_events = set(self.config.messaging.auto_event_types) + if event.event_type not in auto_events: + return + if ( + self.defer_turn_complete_notification + and event.event_type == "turn_complete" + and not include_deferred_turn_complete + ): + return + + requests = self._build_auto_notification_requests(event) + for request in requests: + await self.notification_gateway.enqueue(request) + + def _build_auto_notification_requests( + self, event: Event + ) -> list[NotificationRequest]: + metadata = { + "session_id": self.session_id, + "model": self.config.model_name, + "event_type": event.event_type, + } + + title: str | None = None + message: str | None = None + severity = "info" + data = event.data or {} + if event.event_type == "approval_required": + tools = data.get("tools", []) + if _approval_tools_are_usage_thresholds(tools): + tool = tools[0] + args = tool.get("arguments") if isinstance(tool, dict) else {} + args = args if isinstance(args, dict) else {} + current = _format_usd(args.get("current_spend_usd")) + threshold = _format_usd(args.get("threshold_usd")) + next_threshold = _format_usd(args.get("next_threshold_usd")) + title = "Usage approval required" + message = ( + f"Session {self.session_id} reached {current} in current-session " + f"usage, crossing the {threshold} warning threshold." + ) + if next_threshold: + message += f" The next warning is at {next_threshold}." + severity = "warning" + else: + tools = data.get("tools", []) + tool_names = [] + for tool in tools if isinstance(tools, list) else []: + if isinstance(tool, dict): + tool_name = str(tool.get("tool") or "").strip() + if tool_name and tool_name not in tool_names: + tool_names.append(tool_name) + count = len(tools) if isinstance(tools, list) else 0 + title = "Agent approval required" + message = ( + f"Session {self.session_id} is waiting for approval " + f"for {count} tool call(s)." + ) + if tool_names: + message += " Tools: " + ", ".join(tool_names) + severity = "warning" + elif event.event_type == "error": + title = "Agent error" + error = str(data.get("error") or "Unknown error") + message = f"Session {self.session_id} hit an error.\n{error[:500]}" + severity = "error" + elif event.event_type == "turn_complete": + title = "Agent task complete" + summary = str(data.get("final_response") or "").strip() + if summary: + summary = summary[:_TURN_COMPLETE_NOTIFICATION_CHARS] + message = ( + f"Session {self.session_id} completed successfully.\n{summary}" + ) + else: + message = f"Session {self.session_id} completed successfully." + severity = "success" + + if message is None: + return [] + + requests: list[NotificationRequest] = [] + for destination in self.notification_destinations: + if not self.config.messaging.can_auto_send(destination): + continue + requests.append( + NotificationRequest( + destination=destination, + title=title, + message=message, + severity=severity, + metadata=metadata, + event_type=event.event_type, + ) + ) + return requests def cancel(self) -> None: """Signal cancellation to the running agent loop.""" @@ -146,8 +373,45 @@ def is_cancelled(self) -> bool: def update_model(self, model_name: str) -> None: """Switch the active model and update the context window limit.""" - self.config.model_name = model_name - self.context_manager.model_max_tokens = _get_max_tokens_safe(model_name) + from agent.core.model_ids import strip_huggingface_model_prefix + + normalized = strip_huggingface_model_prefix(model_name) or model_name + self.config.model_name = normalized + self.context_manager.model_max_tokens = _get_max_tokens_safe(normalized) + + def set_auto_approval_policy( + self, *, enabled: bool, cost_cap_usd: float | None + ) -> None: + self.auto_approval_enabled = bool(enabled) + self.auto_approval_cost_cap_usd = cost_cap_usd + + def add_auto_approval_estimated_spend(self, amount_usd: float | None) -> None: + if amount_usd is None or amount_usd <= 0: + return + self.auto_approval_estimated_spend_usd = round( + self.auto_approval_estimated_spend_usd + float(amount_usd), 4 + ) + + @property + def auto_approval_remaining_usd(self) -> float | None: + if self.auto_approval_cost_cap_usd is None: + return None + return round( + max( + 0.0, + self.auto_approval_cost_cap_usd + - self.auto_approval_estimated_spend_usd, + ), + 4, + ) + + def auto_approval_policy_summary(self) -> dict[str, Any]: + return { + "enabled": self.auto_approval_enabled, + "cost_cap_usd": self.auto_approval_cost_cap_usd, + "estimated_spend_usd": round(self.auto_approval_estimated_spend_usd, 4), + "remaining_usd": self.auto_approval_remaining_usd, + } def effective_effort_for(self, model_name: str) -> str | None: """Resolve the effort level to actually send for ``model_name``. @@ -166,6 +430,88 @@ def increment_turn(self) -> None: """Increment turn counter (called after each user interaction)""" self.turn_count += 1 + def start_new_conversation(self) -> dict[str, Any]: + """Rotate this runtime into a fresh conversation. + + The tool router, model/config choices, user identity, and external + resources stay attached to the CLI process. Conversation-specific state + gets reset so later saves do not merge with the prior chat. Warm runtime + resources such as the sandbox, in-flight job tracking, and probed + model-effort cache are deliberately preserved. + """ + previous_session_id = self.session_id + previous_turn_count = self.turn_count + previous_message_count = len(self.context_manager.items) + previous_non_system_count = sum( + 1 + for item in self.context_manager.items + if getattr(item, "role", None) != "system" + ) + + saved_path: str | None = None + if self.config.save_sessions and previous_non_system_count: + saved_path = self.save_and_upload_detached(self.config.session_dataset_repo) + + from agent.tools.plan_tool import reset_current_plan + + self.current_plan = [] + reset_current_plan() + + system_msg = self._fresh_system_message() + self.context_manager.items = [system_msg] if system_msg is not None else [] + self.context_manager.running_context_usage = 0 + + self.session_id = str(uuid.uuid4()) + self.inference_billing_session_id = None + self.session_start_time = datetime.now().astimezone().isoformat() + self.turn_count = 0 + self.last_auto_save_turn = 0 + self.logged_events = [] + self._local_save_path = None + self._last_heartbeat_ts = None + self.pending_approval = None + self.auto_approval_estimated_spend_usd = 0.0 + self._yolo_budget_reservations = {} + self.usage_hf_billing_snapshot = None + self.usage_metrics = None + self.reset_cancel() + + # Previous-session metadata is intentionally included for event + # consumers and telemetry, even though the CLI currently prints only + # the optional save path. + return { + "session_id": self.session_id, + "previous_session_id": previous_session_id, + "previous_turn_count": previous_turn_count, + "previous_message_count": previous_message_count, + "saved_path": saved_path, + } + + def _fresh_system_message(self) -> Message | None: + existing = ( + self.context_manager.items[0] + if self.context_manager.items + and getattr(self.context_manager.items[0], "role", None) == "system" + else None + ) + refresh = getattr(self.context_manager, "refresh_system_prompt", None) + if refresh is None: + return existing + try: + tool_specs = ( + self.tool_router.get_tool_specs_for_llm() if self.tool_router else [] + ) + return refresh( + tool_specs=tool_specs, + hf_token=self.hf_token, + hf_username=self.hf_username, + local_mode=self.local_mode, + autonomous_mode=self.autonomous_mode, + ) + except Exception as e: + logger.warning("Failed to refresh system prompt for new chat: %s", e) + return existing + async def auto_save_if_needed(self) -> None: """Check if auto-save should trigger and save if so (completely non-blocking)""" if not self.config.save_sessions: @@ -184,18 +530,49 @@ async def auto_save_if_needed(self) -> None: def get_trajectory(self) -> dict: """Serialize complete session trajectory for logging""" + tools: list = [] + if self.tool_router is not None: + try: + tools = self.tool_router.get_tool_specs_for_llm() or [] + except Exception: + tools = [] + # Sum per-call cost from llm_call events so analyzers don't have to + # walk the events array themselves. Each `llm_call` event already + # carries cost_usd from `agent.core.telemetry.record_llm_call`. + total_cost_usd = sum( + float((e.get("data") or {}).get("cost_usd") or 0.0) + for e in self.logged_events + if e.get("event_type") == "llm_call" + ) + try: + from agent.core.usage_metrics import summarize_usage_events + + usage_metrics = summarize_usage_events( + self.logged_events, + session_id=self.session_id, + hf_billing_snapshot=self.usage_hf_billing_snapshot, + ) + self.usage_metrics = usage_metrics + except Exception as e: + logger.debug("Usage metrics summary failed for %s: %s", self.session_id, e) + usage_metrics = self.usage_metrics or {} return { "session_id": self.session_id, + "user_id": self.user_id, + "hf_username": self.hf_username, "session_start_time": self.session_start_time, "session_end_time": datetime.now().isoformat(), "model_name": self.config.model_name, + "total_cost_usd": total_cost_usd, + "usage_metrics": usage_metrics, "messages": [msg.model_dump() for msg in self.context_manager.items], "events": self.logged_events, + "tools": tools, } def save_trajectory_local( self, - directory: str = "session_logs", + directory: str = str(DEFAULT_SESSION_LOG_DIR), upload_status: str = "pending", dataset_url: Optional[str] = None, ) -> Optional[str]: @@ -216,98 +593,217 @@ def save_trajectory_local( trajectory = self.get_trajectory() + # Scrub secrets at save time so session_logs/ never holds raw + # tokens on disk — a log aggregator, crash dump, or filesystem + # snapshot between heartbeats would otherwise leak them. + try: + from agent.core.redact import scrub + + for key in ("messages", "events", "tools"): + if key in trajectory: + trajectory[key] = scrub(trajectory[key]) + except Exception as _e: + logger.debug("Redact-on-save failed (non-fatal): %s", _e) + # Add upload metadata trajectory["upload_status"] = upload_status trajectory["upload_url"] = dataset_url trajectory["last_save_time"] = datetime.now().isoformat() - filename = f"session_{self.session_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" - filepath = log_dir / filename - - with open(filepath, "w") as f: + # Reuse one stable path per session so heartbeat saves overwrite + # the same file instead of creating a new timestamped file every + # minute. The timestamp in the filename is kept for first-save + # ordering; subsequent saves just rewrite that file. + if self._local_save_path and Path(self._local_save_path).parent == log_dir: + filepath = Path(self._local_save_path) + else: + filename = ( + f"session_{self.session_id}_" + f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" + ) + filepath = log_dir / filename + self._local_save_path = str(filepath) + + # Atomic-ish write: stage to .tmp then rename so a crash mid-write + # doesn't leave a truncated JSON that breaks the retry scanner. + tmp_path = filepath.with_suffix(filepath.suffix + ".tmp") + with open(tmp_path, "w") as f: json.dump(trajectory, f, indent=2) + tmp_path.replace(filepath) return str(filepath) except Exception as e: logger.error(f"Failed to save session locally: {e}") return None - def update_local_save_status( - self, filepath: str, upload_status: str, dataset_url: Optional[str] = None - ) -> bool: - """Update the upload status of an existing local save file""" - try: - with open(filepath, "r") as f: - data = json.load(f) + def _personal_trace_repo_id(self) -> Optional[str]: + """Resolve the per-user trace repo id from config + HF username. - data["upload_status"] = upload_status - data["upload_url"] = dataset_url - data["last_save_time"] = datetime.now().isoformat() + Returns ``None`` when sharing is disabled, the user is anonymous, + or the template is missing — caller skips the personal upload in + those cases. + """ + if not getattr(self.config, "share_traces", False): + return None + hf_user = self.hf_username or self.user_id + if not hf_user: + return None + template = getattr(self.config, "personal_trace_repo_template", None) + if not template: + return None + try: + return template.format(hf_user=hf_user) + except (KeyError, IndexError): + logger.debug("personal_trace_repo_template format failed: %r", template) + return None - with open(filepath, "w") as f: - json.dump(data, f, indent=2) + def _spawn_uploader( + self, + action: str, + target: str, + repo_id: str, + *, + format: str, + token_env: Optional[str], + private: bool, + token_value: Optional[str] = None, + ) -> None: + """Fire-and-forget spawn of ``session_uploader.py`` with the given args.""" + try: + uploader_script = Path(__file__).parent / "session_uploader.py" + cmd = [ + sys.executable, + str(uploader_script), + action, + target, + repo_id, + "--format", + format, + "--private", + "true" if private else "false", + ] + if token_env: + cmd.extend(["--token-env", token_env]) + + env = os.environ.copy() + if token_value: + env["_ML_INTERN_PERSONAL_TOKEN"] = token_value - return True + subprocess.Popen( + cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + env=env, + start_new_session=True, # Detach from parent + ) except Exception as e: - logger.error(f"Failed to update local save status: {e}") - return False + logger.warning(f"Failed to spawn upload subprocess: {e}") def save_and_upload_detached(self, repo_id: str) -> Optional[str]: """ - Save session locally and spawn detached subprocess for upload (fire-and-forget) + Save session locally and spawn detached subprocess(es) for upload + (fire-and-forget). + + Always uploads to the shared org dataset (``repo_id``) in the + single-row format used by the KPI scheduler. When + ``config.share_traces`` is enabled and a username is known, also + uploads to the user's personal private dataset in Claude Code JSONL + format so the HF Agent Trace Viewer auto-renders it. Args: - repo_id: HuggingFace dataset repo ID + repo_id: HuggingFace dataset repo ID for the org/KPI upload. Returns: Path to local save file """ - # Save locally first (fast, synchronous) local_path = self.save_trajectory_local(upload_status="pending") if not local_path: return None - # Spawn detached subprocess for upload (fire-and-forget) - try: - uploader_script = Path(__file__).parent / "session_uploader.py" + self._spawn_uploader( + "upload", + local_path, + repo_id, + format="row", + token_env=None, # default org token chain + private=False, + ) - # Use Popen with detached process - subprocess.Popen( - [sys.executable, str(uploader_script), "upload", local_path, repo_id], - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, # Detach from parent + personal_repo = self._personal_trace_repo_id() + if personal_repo: + # User's own HF_TOKEN write-scoped to their namespace. + self._spawn_uploader( + "upload", + local_path, + personal_repo, + format="claude_code", + token_env="HF_TOKEN", + token_value=self.hf_token, + private=True, ) - except Exception as e: - logger.warning(f"Failed to spawn upload subprocess: {e}") return local_path @staticmethod def retry_failed_uploads_detached( - directory: str = "session_logs", repo_id: Optional[str] = None + directory: str = str(DEFAULT_SESSION_LOG_DIR), + repo_id: Optional[str] = None, + *, + personal_repo_id: Optional[str] = None, ) -> None: """ - Spawn detached subprocess to retry failed/pending uploads (fire-and-forget) + Spawn detached subprocess(es) to retry failed/pending uploads + (fire-and-forget). Args: directory: Directory containing session logs - repo_id: Target dataset repo ID + repo_id: Target dataset repo ID for the shared org/KPI upload. + personal_repo_id: Per-user dataset for Claude-Code-format + retries. ``None`` skips the personal retry pass. """ - if not repo_id: + if not repo_id and not personal_repo_id: return try: uploader_script = Path(__file__).parent / "session_uploader.py" - # Spawn detached subprocess for retry - subprocess.Popen( - [sys.executable, str(uploader_script), "retry", directory, repo_id], - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, # Detach from parent - ) + if repo_id: + subprocess.Popen( + [ + sys.executable, + str(uploader_script), + "retry", + directory, + repo_id, + "--format", + "row", + ], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + + if personal_repo_id: + subprocess.Popen( + [ + sys.executable, + str(uploader_script), + "retry", + directory, + personal_repo_id, + "--format", + "claude_code", + "--token-env", + "HF_TOKEN", + "--private", + "true", + ], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) except Exception as e: logger.warning(f"Failed to spawn retry subprocess: {e}") diff --git a/agent/core/session_persistence.py b/agent/core/session_persistence.py new file mode 100644 index 000000000..760b59340 --- /dev/null +++ b/agent/core/session_persistence.py @@ -0,0 +1,520 @@ +"""Optional durable session persistence for the hosted backend. + +The public CLI must keep working without MongoDB. This module therefore +exposes one small async store interface and returns a no-op implementation +unless ``MONGODB_URI`` is configured and reachable. +""" + +from __future__ import annotations + +import logging +import os +from datetime import UTC, datetime +from typing import Any + +from bson import BSON +from pymongo import AsyncMongoClient, DeleteMany, ReturnDocument, UpdateOne +from pymongo.errors import InvalidDocument, PyMongoError + +logger = logging.getLogger(__name__) + +SCHEMA_VERSION = 1 +MAX_BSON_BYTES = 15 * 1024 * 1024 +USAGE_EVENT_TYPES = ( + "llm_call", + "hf_job_complete", + "sandbox_create", + "sandbox_destroy", +) + + +def _now() -> datetime: + return datetime.now(UTC) + + +def _doc_id(session_id: str, idx: int) -> str: + return f"{session_id}:{idx}" + + +def _safe_message_doc(message: dict[str, Any]) -> dict[str, Any]: + """Return a Mongo-safe message document payload. + + Mongo's hard document limit is 16 MB. We stay below that and store an + explicit marker rather than failing the whole snapshot for one huge tool log. + """ + try: + if len(BSON.encode({"message": message})) <= MAX_BSON_BYTES: + return message + except (InvalidDocument, OverflowError): + pass + return { + "role": "tool", + "content": ( + "[SYSTEM: A single persisted message exceeded MongoDB's document " + "size/encoding limit and was replaced by this marker.]" + ), + "ml_intern_persistence_error": "message_too_large_or_invalid", + } + + +class NoopSessionStore: + """Async no-op store used when Mongo is not configured.""" + + enabled = False + + async def init(self) -> None: + return None + + async def close(self) -> None: + return None + + async def upsert_session(self, **_: Any) -> None: + return None + + async def save_snapshot(self, **_: Any) -> None: + return None + + async def load_session(self, *_: Any, **__: Any) -> dict[str, Any] | None: + return None + + async def list_sessions(self, *_: Any, **__: Any) -> list[dict[str, Any]]: + return [] + + async def soft_delete_session(self, *_: Any, **__: Any) -> None: + return None + + async def update_session_fields(self, *_: Any, **__: Any) -> None: + return None + + async def append_event(self, *_: Any, **__: Any) -> int | None: + return None + + async def load_events_after(self, *_: Any, **__: Any) -> list[dict[str, Any]]: + return [] + + async def load_usage_events(self, *_: Any, **__: Any) -> list[dict[str, Any]]: + return [] + + async def append_trace_message(self, *_: Any, **__: Any) -> int | None: + return None + + async def mark_pro_seen(self, *_: Any, **__: Any) -> dict[str, Any] | None: + return None + + +class MongoSessionStore(NoopSessionStore): + """MongoDB-backed session store.""" + + enabled = True + + def __init__(self, uri: str, db_name: str) -> None: + self.uri = uri + self.db_name = db_name + self.enabled = False + self.client: AsyncMongoClient | None = None + self.db = None + + async def init(self) -> None: + try: + self.client = AsyncMongoClient(self.uri, serverSelectionTimeoutMS=3000) + self.db = self.client[self.db_name] + await self.client.admin.command("ping") + await self._create_indexes() + self.enabled = True + logger.info("Mongo session persistence enabled (db=%s)", self.db_name) + except Exception as e: + logger.warning("Mongo session persistence disabled: %s", e) + self.enabled = False + if self.client is not None: + await self.client.close() + self.client = None + self.db = None + + async def close(self) -> None: + if self.client is not None: + await self.client.close() + self.client = None + self.db = None + + async def _create_indexes(self) -> None: + if self.db is None: + return + await self.db.sessions.create_index( + [("user_id", 1), ("visibility", 1), ("updated_at", -1)] + ) + await self.db.sessions.create_index( + [("visibility", 1), ("status", 1), ("last_active_at", -1)] + ) + await self.db.session_messages.create_index( + [("session_id", 1), ("idx", 1)], unique=True + ) + await self.db.session_events.create_index( + [("session_id", 1), ("seq", 1)], unique=True + ) + await self.db.session_events.create_index( + [("session_id", 1), ("created_at", 1), ("event_type", 1)] + ) + await self.db.session_trace_messages.create_index( + [("session_id", 1), ("seq", 1)], unique=True + ) + await self.db.session_trace_messages.create_index([("created_at", -1)]) + await self.db.pro_users.create_index([("first_seen_pro_at", -1)]) + + def _ready(self) -> bool: + return bool(self.enabled and self.db is not None) + + async def upsert_session( + self, + *, + session_id: str, + user_id: str, + model: str, + title: str | None = None, + surface: str = "frontend", + created_at: datetime | None = None, + usage_window_started_at: datetime | None = None, + inference_billing_session_id: str | None = None, + runtime_state: str = "idle", + status: str = "active", + message_count: int = 0, + turn_count: int = 0, + pending_approval: list[dict[str, Any]] | None = None, + notification_destinations: list[str] | None = None, + auto_approval_enabled: bool = False, + auto_approval_cost_cap_usd: float | None = None, + auto_approval_estimated_spend_usd: float = 0.0, + usage_warning_next_threshold_usd: float = 5.0, + ) -> None: + if not self._ready(): + return + now = _now() + await self.db.sessions.update_one( + {"_id": session_id}, + { + "$setOnInsert": { + "_id": session_id, + "session_id": session_id, + "user_id": user_id, + "surface": surface, + "created_at": created_at or now, + "schema_version": SCHEMA_VERSION, + "visibility": "live", + }, + "$set": { + "title": title, + "model": model, + "usage_window_started_at": ( + usage_window_started_at or created_at or now + ), + "inference_billing_session_id": inference_billing_session_id, + "status": status, + "runtime_state": runtime_state, + "updated_at": now, + "last_active_at": now, + "message_count": message_count, + "turn_count": turn_count, + "pending_approval": pending_approval or [], + "notification_destinations": notification_destinations or [], + "auto_approval_enabled": auto_approval_enabled, + "auto_approval_cost_cap_usd": auto_approval_cost_cap_usd, + "auto_approval_estimated_spend_usd": auto_approval_estimated_spend_usd, + "usage_warning_next_threshold_usd": usage_warning_next_threshold_usd, + }, + }, + upsert=True, + ) + + async def save_snapshot( + self, + *, + session_id: str, + user_id: str, + model: str, + messages: list[dict[str, Any]], + title: str | None = None, + runtime_state: str = "idle", + status: str = "active", + turn_count: int = 0, + pending_approval: list[dict[str, Any]] | None = None, + created_at: datetime | None = None, + usage_window_started_at: datetime | None = None, + inference_billing_session_id: str | None = None, + notification_destinations: list[str] | None = None, + auto_approval_enabled: bool = False, + auto_approval_cost_cap_usd: float | None = None, + auto_approval_estimated_spend_usd: float = 0.0, + usage_warning_next_threshold_usd: float = 5.0, + raise_on_error: bool = False, + ) -> None: + if not self._ready(): + if raise_on_error: + raise RuntimeError("session store not ready") + return + now = _now() + await self.upsert_session( + session_id=session_id, + user_id=user_id, + model=model, + title=title, + created_at=created_at, + runtime_state=runtime_state, + status=status, + message_count=len(messages), + turn_count=turn_count, + pending_approval=pending_approval, + notification_destinations=notification_destinations, + usage_window_started_at=usage_window_started_at, + inference_billing_session_id=inference_billing_session_id, + auto_approval_enabled=auto_approval_enabled, + auto_approval_cost_cap_usd=auto_approval_cost_cap_usd, + auto_approval_estimated_spend_usd=auto_approval_estimated_spend_usd, + usage_warning_next_threshold_usd=usage_warning_next_threshold_usd, + ) + ops: list[Any] = [] + for idx, raw in enumerate(messages): + ops.append( + UpdateOne( + {"_id": _doc_id(session_id, idx)}, + { + "$set": { + "session_id": session_id, + "idx": idx, + "message": _safe_message_doc(raw), + "updated_at": now, + }, + "$setOnInsert": {"created_at": now}, + }, + upsert=True, + ) + ) + ops.append( + DeleteMany({"session_id": session_id, "idx": {"$gte": len(messages)}}) + ) + try: + if ops: + await self.db.session_messages.bulk_write(ops, ordered=False) + except PyMongoError as e: + # Best-effort by default, but the reaper passes raise_on_error so a + # silent message-write failure doesn't let it evict a session whose + # latest messages never made it to Mongo. + if raise_on_error: + raise + logger.warning("Failed to persist session %s snapshot: %s", session_id, e) + + async def load_session( + self, session_id: str, *, include_deleted: bool = False + ) -> dict[str, Any] | None: + if not self._ready(): + return None + meta = await self.db.sessions.find_one({"_id": session_id}) + if not meta: + return None + if meta.get("visibility") == "deleted" and not include_deleted: + return None + cursor = self.db.session_messages.find({"session_id": session_id}).sort( + "idx", 1 + ) + messages = [row.get("message") async for row in cursor] + return {"metadata": meta, "messages": messages} + + async def list_sessions( + self, user_id: str, *, include_deleted: bool = False + ) -> list[dict[str, Any]]: + if not self._ready(): + return [] + query: dict[str, Any] = {"user_id": user_id} + if user_id == "dev": + query = {} + if not include_deleted: + query["visibility"] = {"$ne": "deleted"} + cursor = self.db.sessions.find(query).sort("updated_at", -1) + return [row async for row in cursor] + + async def soft_delete_session(self, session_id: str) -> None: + if not self._ready(): + return + await self.db.sessions.update_one( + {"_id": session_id}, + { + "$set": { + "visibility": "deleted", + "runtime_state": "idle", + "updated_at": _now(), + } + }, + ) + + async def update_session_fields(self, session_id: str, **fields: Any) -> None: + if not self._ready() or not fields: + return + fields["updated_at"] = _now() + await self.db.sessions.update_one({"_id": session_id}, {"$set": fields}) + + async def _next_seq(self, counter_id: str) -> int: + doc = await self.db.counters.find_one_and_update( + {"_id": counter_id}, + {"$inc": {"seq": 1}}, + upsert=True, + return_document=ReturnDocument.AFTER, + ) + return int(doc["seq"]) + + async def append_event( + self, session_id: str, event_type: str, data: dict[str, Any] | None + ) -> int | None: + if not self._ready(): + return None + try: + seq = await self._next_seq(f"event:{session_id}") + await self.db.session_events.insert_one( + { + "_id": _doc_id(session_id, seq), + "session_id": session_id, + "seq": seq, + "event_type": event_type, + "data": data or {}, + "created_at": _now(), + } + ) + return seq + except PyMongoError as e: + logger.debug("Failed to append event for %s: %s", session_id, e) + return None + + async def load_events_after( + self, session_id: str, after_seq: int = 0 + ) -> list[dict[str, Any]]: + if not self._ready(): + return [] + cursor = self.db.session_events.find( + {"session_id": session_id, "seq": {"$gt": int(after_seq or 0)}} + ).sort("seq", 1) + return [row async for row in cursor] + + async def load_usage_events( + self, + user_id: str, + *, + session_id: str | None = None, + start: datetime | None = None, + end: datetime | None = None, + ) -> list[dict[str, Any]]: + if not self._ready(): + return [] + session_query: dict[str, Any] = {"visibility": {"$ne": "deleted"}} + if user_id != "dev": + session_query["user_id"] = user_id + if session_id is not None: + session_query["_id"] = session_id + + session_cursor = self.db.sessions.find(session_query, {"_id": 1}) + session_ids = [str(row.get("_id")) async for row in session_cursor] + if not session_ids: + return [] + + event_query: dict[str, Any] = { + "session_id": {"$in": session_ids}, + "event_type": {"$in": list(USAGE_EVENT_TYPES)}, + } + if start is not None or end is not None: + created_at: dict[str, datetime] = {} + if start is not None: + created_at["$gte"] = start + if end is not None: + created_at["$lt"] = end + event_query["created_at"] = created_at + + event_cursor = self.db.session_events.find(event_query).sort("created_at", 1) + return [row async for row in event_cursor] + + async def append_trace_message( + self, session_id: str, message: dict[str, Any], source: str = "message" + ) -> int | None: + if not self._ready(): + return None + try: + seq = await self._next_seq(f"trace:{session_id}") + await self.db.session_trace_messages.insert_one( + { + "_id": _doc_id(session_id, seq), + "session_id": session_id, + "seq": seq, + "role": message.get("role"), + "message": _safe_message_doc(message), + "source": source, + "created_at": _now(), + } + ) + return seq + except PyMongoError as e: + logger.debug("Failed to append trace message for %s: %s", session_id, e) + return None + + async def mark_pro_seen( + self, user_id: str, *, is_pro: bool + ) -> dict[str, Any] | None: + """Track per-user Pro state and detect free→Pro conversions. + + Returns ``{"converted": True, "first_seen_at": ..."}`` exactly once + per user — the first time we see them as Pro after having recorded + them as non-Pro at least once. Otherwise returns ``None``. + + Storing ``ever_non_pro`` lets us distinguish "user joined as Pro" + (no conversion) from "user upgraded" (conversion). The atomic + ``find_one_and_update`` on a guarded filter makes the conversion + emit at-most-once even under concurrent requests. + """ + if not self._ready() or not user_id: + return None + now = _now() + set_fields: dict[str, Any] = {"last_seen_at": now, "is_pro": bool(is_pro)} + if not is_pro: + set_fields["ever_non_pro"] = True + try: + await self.db.pro_users.update_one( + {"_id": user_id}, + { + "$setOnInsert": {"_id": user_id, "first_seen_at": now}, + "$set": set_fields, + }, + upsert=True, + ) + except PyMongoError as e: + logger.debug("mark_pro_seen upsert failed for %s: %s", user_id, e) + return None + + if not is_pro: + return None + + try: + doc = await self.db.pro_users.find_one_and_update( + { + "_id": user_id, + "ever_non_pro": True, + "first_seen_pro_at": {"$exists": False}, + }, + {"$set": {"first_seen_pro_at": now}}, + return_document=ReturnDocument.AFTER, + ) + except PyMongoError as e: + logger.debug("mark_pro_seen conversion check failed for %s: %s", user_id, e) + return None + + if not doc: + return None + return { + "converted": True, + "first_seen_at": (doc.get("first_seen_at") or now).isoformat(), + } + + +_store: NoopSessionStore | MongoSessionStore | None = None + + +def get_session_store() -> NoopSessionStore | MongoSessionStore: + global _store + if _store is None: + uri = os.environ.get("MONGODB_URI") + db_name = os.environ.get("MONGODB_DB", "ml-intern") + _store = MongoSessionStore(uri, db_name) if uri else NoopSessionStore() + return _store diff --git a/agent/core/session_resume.py b/agent/core/session_resume.py new file mode 100644 index 000000000..ac7d335f6 --- /dev/null +++ b/agent/core/session_resume.py @@ -0,0 +1,289 @@ +"""Reload a previously saved session log into the active CLI session.""" + +from __future__ import annotations + +import json +import logging +import re +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any + +from litellm import Message + +from agent.core.model_ids import strip_huggingface_model_prefix +from agent.core.model_switcher import is_valid_model_id +from agent.core.session import DEFAULT_SESSION_LOG_DIR + +logger = logging.getLogger(__name__) + +_REDACTED_MARKER = re.compile(r"\[REDACTED_[A-Z_]+\]") + + +@dataclass +class SessionLogEntry: + """Metadata for a locally saved session log.""" + + path: Path + session_id: str + session_start_time: str | None + session_end_time: str | None + model_name: str | None + message_count: int + preview: str + mtime: float + + +def _message_preview(content: Any, max_chars: int = 72) -> str: + """Return a one-line preview for string or OpenAI-style block content.""" + if isinstance(content, str): + text = content + elif isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, dict): + value = block.get("text") or block.get("content") + if isinstance(value, str): + parts.append(value) + elif isinstance(block, str): + parts.append(block) + text = " ".join(parts) + else: + text = "" + text = " ".join(text.split()) + if len(text) > max_chars: + return text[: max_chars - 1].rstrip() + "…" + return text + + +def _first_user_preview(messages: list[Any]) -> str: + for raw in messages: + if isinstance(raw, dict) and raw.get("role") == "user": + preview = _message_preview(raw.get("content")) + if preview: + return preview + return "(no user prompt preview)" + + +def list_session_logs( + directory: Path = DEFAULT_SESSION_LOG_DIR, +) -> list[SessionLogEntry]: + """Return readable session logs under ``directory``, newest first.""" + if not directory.exists(): + return [] + + entries: list[SessionLogEntry] = [] + for path in directory.glob("*.json"): + try: + with open(path) as f: + data = json.load(f) + except Exception: + continue + + messages = data.get("messages") or [] + if not isinstance(messages, list): + continue + + session_id = data.get("session_id") + if not isinstance(session_id, str) or not session_id: + session_id = path.stem + + stat = path.stat() + entries.append( + SessionLogEntry( + path=path, + session_id=session_id, + session_start_time=data.get("session_start_time"), + session_end_time=data.get("session_end_time"), + model_name=data.get("model_name"), + message_count=len(messages), + preview=_first_user_preview(messages), + mtime=stat.st_mtime, + ) + ) + + entries.sort(key=lambda item: item.mtime, reverse=True) + return entries + + +def format_session_log_entry(index: int, entry: SessionLogEntry) -> str: + timestamp = entry.session_end_time or entry.session_start_time + label = "unknown time" + if isinstance(timestamp, str) and timestamp: + try: + label = datetime.fromisoformat(timestamp).strftime("%Y-%m-%d %H:%M") + except ValueError: + label = timestamp[:16] + short_id = entry.session_id[:8] + model = entry.model_name or "unknown model" + return ( + f"{index:>2}. {label} {short_id} " + f"{entry.message_count} msgs {model}\n" + f" {entry.preview}" + ) + + +def resolve_session_log_arg( + arg: str, + entries: list[SessionLogEntry], + directory: Path = DEFAULT_SESSION_LOG_DIR, +) -> Path | None: + """Resolve ``/resume `` as index, path, filename, or session id prefix.""" + value = arg.strip() + if not value: + return None + + if value.isdigit(): + idx = int(value) + if 1 <= idx <= len(entries): + return entries[idx - 1].path + + candidate = Path(value).expanduser() + candidates = [candidate] + if not candidate.is_absolute(): + candidates.append(directory / candidate) + if candidate.suffix != ".json": + candidates.append(directory / f"{value}.json") + + for path in candidates: + if path.exists() and path.is_file(): + return path + + matches = [ + entry.path + for entry in entries + if entry.session_id.startswith(value) or entry.path.name.startswith(value) + ] + if len(matches) == 1: + return matches[0] + return None + + +def _turn_count_from_messages(messages: list[Any]) -> int: + return sum( + 1 for raw in messages if isinstance(raw, dict) and raw.get("role") == "user" + ) + + +def _has_redacted_content(messages: list[Any]) -> bool: + """Whether any message body contains a ``[REDACTED_*]`` marker.""" + for raw in messages: + if not isinstance(raw, dict): + continue + content = raw.get("content") + if isinstance(content, str) and _REDACTED_MARKER.search(content): + return True + if isinstance(content, list): + for block in content: + if isinstance(block, dict): + text = block.get("text") or block.get("content") + if isinstance(text, str) and _REDACTED_MARKER.search(text): + return True + return False + + +def restore_session_from_log(session: Any, path: Path) -> dict[str, Any]: + """Replace the active session context with messages from ``path``. + + Continues the saved session (reusing its id and on-disk save path) when + the log's ``user_id`` matches the current session, and forks otherwise: + the caller's session id stays put and future heartbeat saves go to a + fresh file rather than overwriting the source log. + + Returns metadata for the ``resume_complete`` event. + """ + with open(path) as f: + data = json.load(f) + + raw_messages = data.get("messages") + if not isinstance(raw_messages, list): + raise ValueError("Selected log does not contain a messages array") + + restored_messages: list[Message] = [] + dropped_count = 0 + for raw in raw_messages: + if not isinstance(raw, dict) or raw.get("role") == "system": + continue + try: + restored_messages.append(Message.model_validate(raw)) + except Exception as e: + dropped_count += 1 + logger.warning("Dropping malformed message from %s: %s", path, e) + + if not restored_messages: + raise ValueError("Selected log has no restorable non-system messages") + + cm = session.context_manager + system_msg = cm.items[0] if cm.items and cm.items[0].role == "system" else None + cm.items = ([system_msg] if system_msg else []) + restored_messages + + # Validate the saved model id before switching. ``update_model`` doesn't + # check availability; an unrecognised id silently sticks and the next LLM + # call fails with a cryptic routing error. Logs from a different + # deployment, an older catalog, or a removed model land here. + saved_model = data.get("model_name") + invalid_saved_model: str | None = None + if isinstance(saved_model, str) and saved_model: + normalized_model = strip_huggingface_model_prefix(saved_model) + if normalized_model and is_valid_model_id(normalized_model): + session.update_model(normalized_model) + else: + invalid_saved_model = saved_model + logger.warning( + "Saved log model %r failed format validation; keeping %r", + saved_model, + session.config.model_name, + ) + + cm._recompute_usage(session.config.model_name) + + saved_session_id = data.get("session_id") + saved_user_id = data.get("user_id") + is_continuation = saved_user_id == session.user_id + + if is_continuation: + if isinstance(saved_session_id, str) and saved_session_id: + session.session_id = saved_session_id + session.session_start_time = ( + data.get("session_start_time") or session.session_start_time + ) + + # Always fork the on-disk save path. The source log is treated as an + # immutable snapshot: ``logged_events`` is reset to a single + # ``resumed_from`` marker below for cost accounting, so reusing the + # source path would let the next heartbeat save destroy the original + # ``llm_call``/event history on disk. The next save will pick a fresh + # filename instead. + session._local_save_path = None + + saved_event_count = ( + len(data.get("events", [])) if isinstance(data.get("events"), list) else 0 + ) + session.logged_events = [ + { + "timestamp": datetime.now().isoformat(), + "event_type": "resumed_from", + "data": { + "path": str(path), + "original_session_id": ( + saved_session_id if isinstance(saved_session_id, str) else None + ), + "original_event_count": saved_event_count, + "forked": not is_continuation, + }, + } + ] + session.turn_count = _turn_count_from_messages(raw_messages) + session.last_auto_save_turn = session.turn_count + session.pending_approval = None + + return { + "path": str(path), + "restored_count": len(restored_messages), + "dropped_count": dropped_count, + "model_name": session.config.model_name, + "invalid_saved_model": invalid_saved_model, + "forked": not is_continuation, + "had_redacted_content": _has_redacted_content(raw_messages), + } diff --git a/agent/core/session_uploader.py b/agent/core/session_uploader.py index ef2f9496d..268c84596 100644 --- a/agent/core/session_uploader.py +++ b/agent/core/session_uploader.py @@ -3,32 +3,479 @@ Standalone script for uploading session trajectories to HuggingFace. This runs as a separate process to avoid blocking the main agent. Uses individual file uploads to avoid race conditions. + +Two formats are supported: + +* ``row`` — single-line JSONL row used by the existing org telemetry/KPI + pipeline (``smolagents/ml-intern-sessions``). Compatible with + ``backend/kpis_scheduler.py``. +* ``claude_code`` — one event per line in the Claude Code JSONL schema, + auto-detected by the HF Agent Trace Viewer + (https://huggingface.co/changelog/agent-trace-viewer). Used for the + per-user private dataset (default ``{hf_user}/ml-intern-sessions``). """ +import argparse +import hashlib import json import os import sys from datetime import datetime from pathlib import Path +from typing import Any from dotenv import load_dotenv +from agent.core.usage_metrics import ( + summarize_usage_events, + usage_metric_scalar_fields, +) + load_dotenv() -# Token for session uploads — loaded from env var (never hardcode tokens in source) -_SESSION_TOKEN = os.environ.get("HF_SESSION_UPLOAD_TOKEN", "") +# Token resolution for the org KPI dataset. Fallback chain (least-privilege +# first) — matches backend/kpis_scheduler.py so one write-scoped token on the +# Space covers every telemetry dataset. Never hardcode tokens in source. +_ORG_TOKEN_FALLBACK_CHAIN = ( + "HF_SESSION_UPLOAD_TOKEN", + "HF_TOKEN", + "HF_ADMIN_TOKEN", +) +_PERSONAL_TOKEN_ENV = "_ML_INTERN_PERSONAL_TOKEN" + + +def _resolve_token(token_env: str | None) -> str: + """Resolve an HF token from env. ``token_env`` overrides the fallback chain.""" + if token_env == "HF_TOKEN": + try: + from agent.core.hf_tokens import resolve_hf_token + + return ( + resolve_hf_token( + os.environ.get(_PERSONAL_TOKEN_ENV), + os.environ.get("HF_TOKEN"), + ) + or "" + ) + except Exception: + token = os.environ.get(_PERSONAL_TOKEN_ENV) or os.environ.get("HF_TOKEN") + return token or "" + + if token_env: + return os.environ.get(token_env, "") or "" + for var in _ORG_TOKEN_FALLBACK_CHAIN: + val = os.environ.get(var) + if val: + return val + return "" + + +def _scrub(obj: Any) -> Any: + """Best-effort regex scrub for HF tokens / API keys before upload.""" + try: + from agent.core.redact import scrub # type: ignore + except Exception: + # Fallback for environments where the agent package isn't importable + # (shouldn't happen in our subprocess, but be defensive). + import importlib.util + + _spec = importlib.util.spec_from_file_location( + "_redact", + Path(__file__).parent / "redact.py", + ) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) # type: ignore + scrub = _mod.scrub + return scrub(obj) + + +def _msg_uuid(session_id: str, role: str, idx: int) -> str: + """Deterministic UUID-shaped id for a Claude Code message. + + Uses sha1 of ``session_id::role::idx`` so re-uploads/heartbeats keep the + parent/child chain stable. Same convention as the example dataset + https://huggingface.co/datasets/clem/hf-coding-tools-traces. + """ + digest = hashlib.sha1(f"{session_id}::{role}::{idx}".encode("utf-8")).hexdigest() + # Format like a UUID for visual familiarity (32 hex chars w/ dashes). + return ( + f"{digest[0:8]}-{digest[8:12]}-{digest[12:16]}-{digest[16:20]}-{digest[20:32]}" + ) + + +def _content_to_text(content: Any) -> str: + """Best-effort flatten of a litellm/openai content field to plain text.""" + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, dict): + text = block.get("text") + if isinstance(text, str): + parts.append(text) + else: + # Unknown content block — keep round-trippable representation. + parts.append(json.dumps(block, default=str)) + else: + parts.append(str(block)) + return "\n".join(parts) + return str(content) + + +def _parse_tool_args(raw: Any) -> Any: + """Tool call arguments arrive as a JSON-encoded string from LLMs.""" + if isinstance(raw, dict): + return raw + if isinstance(raw, str): + try: + return json.loads(raw) + except (json.JSONDecodeError, TypeError): + return {"_raw": raw} + return raw + + +def to_claude_code_jsonl(trajectory: dict) -> list[dict]: + """Convert an internal trajectory dict to Claude Code JSONL events. + + Schema reference (per the HF Agent Trace Viewer auto-detector): + + {"type":"user","message":{"role":"user","content":"..."}, + "uuid":"...","parentUuid":null,"sessionId":"...","timestamp":"..."} + {"type":"assistant", + "message":{"role":"assistant","model":"...", + "content":[{"type":"text","text":"..."}, + {"type":"tool_use","id":"...","name":"...","input":{...}}]}, + "uuid":"...","parentUuid":"","sessionId":"...","timestamp":"..."} + {"type":"user","message":{"role":"user", + "content":[{"type":"tool_result", + "tool_use_id":"...","content":"..."}]}, + "uuid":"...","parentUuid":"","sessionId":"...","timestamp":"..."} + + System messages are skipped (they're not part of the viewer schema and + contain large prompts that pollute the trace viewer UI). + """ + session_id = trajectory["session_id"] + model_name = trajectory.get("model_name") or "" + fallback_timestamp = ( + trajectory.get("session_start_time") or datetime.now().isoformat() + ) + messages: list[dict] = trajectory.get("messages") or [] + + out: list[dict] = [] + parent_uuid: str | None = None + + for idx, msg in enumerate(messages): + if not isinstance(msg, dict): + continue + role = msg.get("role") + if role == "system": + continue + timestamp = msg.get("timestamp") or fallback_timestamp + + if role == "user": + content = _content_to_text(msg.get("content")) + event_uuid = _msg_uuid(session_id, "user", idx) + out.append( + { + "type": "user", + "message": {"role": "user", "content": content}, + "uuid": event_uuid, + "parentUuid": parent_uuid, + "sessionId": session_id, + "timestamp": timestamp, + } + ) + parent_uuid = event_uuid + + elif role == "assistant": + content_text = _content_to_text(msg.get("content")) + content_blocks: list[dict] = [] + if content_text: + content_blocks.append({"type": "text", "text": content_text}) + for tc in msg.get("tool_calls") or []: + if not isinstance(tc, dict): + continue + fn = tc.get("function") or {} + content_blocks.append( + { + "type": "tool_use", + "id": tc.get("id") or "", + "name": fn.get("name") or "", + "input": _parse_tool_args(fn.get("arguments")), + } + ) + if not content_blocks: + # Edge case: empty assistant turn (shouldn't normally happen, + # but skip rather than emit an empty content array which + # confuses the viewer). + continue + event_uuid = _msg_uuid(session_id, "assistant", idx) + out.append( + { + "type": "assistant", + "message": { + "role": "assistant", + "model": model_name, + "content": content_blocks, + }, + "uuid": event_uuid, + "parentUuid": parent_uuid, + "sessionId": session_id, + "timestamp": timestamp, + } + ) + parent_uuid = event_uuid + + elif role == "tool": + tool_call_id = msg.get("tool_call_id") or "" + content_text = _content_to_text(msg.get("content")) + event_uuid = _msg_uuid(session_id, "tool", idx) + out.append( + { + "type": "user", + "message": { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": tool_call_id, + "content": content_text, + } + ], + }, + "uuid": event_uuid, + "parentUuid": parent_uuid, + "sessionId": session_id, + "timestamp": timestamp, + } + ) + parent_uuid = event_uuid + + return out + + +def _scrub_session_for_upload(data: dict) -> dict: + """Best-effort scrub of transcript fields before any upload temp file.""" + scrubbed = dict(data) + scrubbed["messages"] = _scrub(data.get("messages") or []) + scrubbed["events"] = _scrub(data.get("events") or []) + scrubbed["tools"] = _scrub(data.get("tools") or []) + return scrubbed + + +def _usage_metrics_for_row(data: dict) -> dict: + metrics = data.get("usage_metrics") + if isinstance(metrics, str): + try: + parsed = json.loads(metrics) + metrics = parsed if isinstance(parsed, dict) else None + except (json.JSONDecodeError, TypeError): + metrics = None + if isinstance(metrics, dict): + return metrics + events = data.get("events") + return summarize_usage_events( + events if isinstance(events, list) else [], + session_id=data.get("session_id"), + ) + + +def _write_row_payload(data: dict, tmp_path: str) -> None: + """Single-row JSONL (existing format) — used by KPI scheduler.""" + scrubbed = _scrub_session_for_upload(data) + usage_metrics = _usage_metrics_for_row(data) + session_row = { + "session_id": data["session_id"], + "user_id": data.get("user_id"), + "session_start_time": data["session_start_time"], + "session_end_time": data["session_end_time"], + "model_name": data["model_name"], + "total_cost_usd": data.get("total_cost_usd"), + "messages": json.dumps(scrubbed["messages"]), + "events": json.dumps(scrubbed["events"]), + "tools": json.dumps(scrubbed["tools"]), + "usage_metrics": json.dumps(_scrub(usage_metrics)), + } + session_row.update(usage_metric_scalar_fields(usage_metrics)) + + with open(tmp_path, "w") as tmp: + json.dump(session_row, tmp) + + +def _write_claude_code_payload(data: dict, tmp_path: str) -> None: + """Multi-line JSONL in Claude Code schema for the HF trace viewer.""" + # Scrub before conversion so secrets never reach the upload temp file. + scrubbed = _scrub_session_for_upload(data) + events = to_claude_code_jsonl(scrubbed) + with open(tmp_path, "w") as tmp: + for event in events: + tmp.write(json.dumps(event)) + tmp.write("\n") + + +def _status_field(format: str) -> str: + """Per-format upload status field on the local trajectory file.""" + return "personal_upload_status" if format == "claude_code" else "upload_status" + + +def _url_field(format: str) -> str: + return "personal_upload_url" if format == "claude_code" else "upload_url" + + +def _read_session_file(session_file: str) -> dict: + """Read a local session file while respecting uploader file locks.""" + import fcntl + + with open(session_file, "r") as f: + fcntl.flock(f, fcntl.LOCK_SH) + try: + return json.load(f) + finally: + fcntl.flock(f, fcntl.LOCK_UN) + + +def _update_upload_status( + session_file: str, + status_key: str, + url_key: str, + status: str, + dataset_url: str | None = None, +) -> None: + """Atomically update only this uploader's status fields. + + The org and personal uploaders run as separate processes against the same + local session JSON file. Re-read under an exclusive lock so one uploader + cannot clobber fields written by the other. + """ + import fcntl + + with open(session_file, "r+") as f: + fcntl.flock(f, fcntl.LOCK_EX) + try: + data = json.load(f) + data[status_key] = status + if dataset_url is not None: + data[url_key] = dataset_url + data["last_save_time"] = datetime.now().isoformat() + f.seek(0) + json.dump(data, f, indent=2) + f.truncate() + f.flush() + os.fsync(f.fileno()) + finally: + fcntl.flock(f, fcntl.LOCK_UN) + + +def dataset_card_readme(repo_id: str) -> str: + """Dataset card for personal ML Intern session trace repos.""" + return """--- +pretty_name: "ML Intern Session Traces" +language: +- en +license: other +task_categories: +- text-generation +tags: +- agent-traces +- coding-agent +- ml-intern +- session-traces +- claude-code +- hf-agent-trace-viewer +configs: +- config_name: default + data_files: + - split: train + path: "sessions/**/*.jsonl" +--- + +# ML Intern session traces + +This dataset contains ML Intern coding agent session traces uploaded from local +ML Intern runs. The traces are stored as JSON Lines files under `sessions/`, +with one file per session. + +## Links + +- ML Intern demo: https://smolagents-ml-intern.hf.space +- ML Intern CLI: https://github.com/huggingface/ml-intern + +## Data description + +Each `*.jsonl` file contains a single ML Intern session converted to a +Claude-Code-style event stream for the Hugging Face Agent Trace Viewer. Entries +can include user messages, assistant messages, tool calls, tool results, model +metadata, and timestamps. + +Session files are written to paths of the form: + +```text +sessions/YYYY-MM-DD/.jsonl +``` + +## Redaction and review + +**WARNING: no comprehensive redaction or human review has been performed for this dataset.** + +ML Intern applies automated best-effort scrubbing for common secret patterns +such as Hugging Face, GitHub, AWS, and provider API tokens before upload. +This is not a privacy guarantee. + +These traces may contain sensitive information, including prompts, code, +terminal output, file paths, repository names, private task context, tool +outputs, or other data from the local development environment. Treat every +session as potentially sensitive. + +Do not make this dataset public unless you have manually inspected the uploaded +sessions and are comfortable sharing their full contents. + +## Limitations + +Coding agent transcripts can include private or off-topic content, failed +experiments, credentials accidentally pasted by a user, and outputs copied from +local files or services. Use with appropriate caution, especially before +changing repository visibility. +""" + + +def _upload_dataset_card(api: Any, repo_id: str, token: str, format: str) -> None: + """Create/update a README for personal trace datasets.""" + if format != "claude_code": + return + + api.upload_file( + path_or_fileobj=dataset_card_readme(repo_id).encode("utf-8"), + path_in_repo="README.md", + repo_id=repo_id, + repo_type="dataset", + token=token, + commit_message="Update dataset card", + ) def upload_session_as_file( - session_file: str, repo_id: str, max_retries: int = 3 + session_file: str, + repo_id: str, + max_retries: int = 3, + format: str = "row", + token_env: str | None = None, + private: bool = False, ) -> bool: - """ - Upload a single session as an individual JSONL file (no race conditions) + """Upload a single session as an individual JSONL file (no race conditions). Args: session_file: Path to local session JSON file repo_id: HuggingFace dataset repo ID max_retries: Number of retry attempts + format: ``row`` (default, KPI-compatible) or ``claude_code`` (HF + Agent Trace Viewer compatible). + token_env: Name of the env var holding the HF token. ``None`` falls + back to the org-token chain (``HF_SESSION_UPLOAD_TOKEN`` → + ``HF_TOKEN`` → ``HF_ADMIN_TOKEN``). + private: When creating the repo for the first time, mark it private. Returns: True if successful, False otherwise @@ -39,72 +486,60 @@ def upload_session_as_file( print("Error: huggingface_hub library not available", file=sys.stderr) return False + status_key = _status_field(format) + url_key = _url_field(format) + try: - # Load session data - with open(session_file, "r") as f: - data = json.load(f) + data = _read_session_file(session_file) - # Check if already uploaded - upload_status = data.get("upload_status") - if upload_status == "success": + # Skip if already uploaded for this format. + if data.get(status_key) == "success": return True - # Use dedicated session upload token (write-only access to session dataset) - hf_token = _SESSION_TOKEN + hf_token = _resolve_token(token_env) if not hf_token: - # Update status to failed - data["upload_status"] = "failed" - with open(session_file, "w") as f: - json.dump(data, f, indent=2) + _update_upload_status(session_file, status_key, url_key, "failed") return False - # Prepare JSONL content (single line) - # Store messages and events as JSON strings to avoid schema conflicts - session_row = { - "session_id": data["session_id"], - "session_start_time": data["session_start_time"], - "session_end_time": data["session_end_time"], - "model_name": data["model_name"], - "messages": json.dumps(data["messages"]), - "events": json.dumps(data["events"]), - } - - # Create temporary JSONL file + # Build temp upload payload in the requested format. import tempfile with tempfile.NamedTemporaryFile( mode="w", suffix=".jsonl", delete=False ) as tmp: - json.dump(session_row, tmp) # Single line JSON tmp_path = tmp.name try: - # Generate unique path in repo: sessions/YYYY-MM-DD/session_id.jsonl + if format == "claude_code": + _write_claude_code_payload(data, tmp_path) + else: + _write_row_payload(data, tmp_path) + session_id = data["session_id"] date_str = datetime.fromisoformat(data["session_start_time"]).strftime( "%Y-%m-%d" ) repo_path = f"sessions/{date_str}/{session_id}.jsonl" - # Upload with retries api = HfApi() for attempt in range(max_retries): try: - # Try to create repo if it doesn't exist (idempotent) + # Idempotent create — visibility is set on first creation + # only. Existing repos keep whatever the user picked via + # /share-traces. try: api.create_repo( repo_id=repo_id, repo_type="dataset", - private=False, + private=private, token=hf_token, - exist_ok=True, # Don't fail if already exists + exist_ok=True, ) - except Exception: - # Repo might already exist, continue pass - # Upload the session file + _upload_dataset_card(api, repo_id, hf_token, format) + api.upload_file( path_or_fileobj=tmp_path, path_in_repo=repo_path, @@ -114,12 +549,13 @@ def upload_session_as_file( commit_message=f"Add session {session_id}", ) - # Update local status to success - data["upload_status"] = "success" - data["upload_url"] = f"https://huggingface.co/datasets/{repo_id}" - with open(session_file, "w") as f: - json.dump(data, f, indent=2) - + _update_upload_status( + session_file, + status_key, + url_key, + "success", + f"https://huggingface.co/datasets/{repo_id}", + ) return True except Exception: @@ -129,14 +565,12 @@ def upload_session_as_file( wait_time = 2**attempt time.sleep(wait_time) else: - # Final attempt failed - data["upload_status"] = "failed" - with open(session_file, "w") as f: - json.dump(data, f, indent=2) + _update_upload_status( + session_file, status_key, url_key, "failed" + ) return False finally: - # Clean up temp file try: os.unlink(tmp_path) except Exception: @@ -147,56 +581,102 @@ def upload_session_as_file( return False -def retry_failed_uploads(directory: str, repo_id: str): - """Retry all failed/pending uploads in a directory""" +def retry_failed_uploads( + directory: str, + repo_id: str, + format: str = "row", + token_env: str | None = None, + private: bool = False, +): + """Retry all failed/pending uploads in a directory for the given format.""" log_dir = Path(directory) if not log_dir.exists(): return + status_key = _status_field(format) session_files = list(log_dir.glob("session_*.json")) for filepath in session_files: try: - with open(filepath, "r") as f: - data = json.load(f) - - upload_status = data.get("upload_status", "unknown") - - # Only retry pending or failed uploads - if upload_status in ["pending", "failed"]: - upload_session_as_file(str(filepath), repo_id) + data = _read_session_file(str(filepath)) + + # Only retry pending or failed uploads. Files predating this + # field don't have it; treat unknown as "not yet attempted" for + # the row format (legacy behavior) and "skip" for claude_code + # so we don't suddenly re-upload pre-existing sessions to a + # newly-introduced personal repo. + status = data.get(status_key, "unknown") + if format == "claude_code" and status_key not in data: + continue + + if status in ("pending", "failed", "unknown"): + upload_session_as_file( + str(filepath), + repo_id, + format=format, + token_env=token_env, + private=private, + ) except Exception: pass +def _str2bool(v: str) -> bool: + return str(v).strip().lower() in {"1", "true", "yes", "on"} + + if __name__ == "__main__": - if len(sys.argv) < 3: - print("Usage: session_uploader.py ") - sys.exit(1) - - command = sys.argv[1] - - if command == "upload": - # python session_uploader.py upload - if len(sys.argv) < 4: - print("Usage: session_uploader.py upload ") - sys.exit(1) - session_file = sys.argv[2] - repo_id = sys.argv[3] - success = upload_session_as_file(session_file, repo_id) - sys.exit(0 if success else 1) - - elif command == "retry": - # python session_uploader.py retry - if len(sys.argv) < 4: - print("Usage: session_uploader.py retry ") - sys.exit(1) - directory = sys.argv[2] - repo_id = sys.argv[3] - retry_failed_uploads(directory, repo_id) + parser = argparse.ArgumentParser(prog="session_uploader.py") + sub = parser.add_subparsers(dest="command", required=True) + + p_upload = sub.add_parser("upload") + p_upload.add_argument("session_file") + p_upload.add_argument("repo_id") + p_upload.add_argument( + "--format", + choices=["row", "claude_code"], + default="row", + ) + p_upload.add_argument( + "--token-env", + default=None, + help="Env var name holding the HF token (default: org fallback chain).", + ) + p_upload.add_argument("--private", default="false") + + p_retry = sub.add_parser("retry") + p_retry.add_argument("directory") + p_retry.add_argument("repo_id") + p_retry.add_argument( + "--format", + choices=["row", "claude_code"], + default="row", + ) + p_retry.add_argument("--token-env", default=None) + p_retry.add_argument("--private", default="false") + + args = parser.parse_args() + + if args.command == "upload": + ok = upload_session_as_file( + args.session_file, + args.repo_id, + format=args.format, + token_env=args.token_env, + private=_str2bool(args.private), + ) + sys.exit(0 if ok else 1) + + if args.command == "retry": + retry_failed_uploads( + args.directory, + args.repo_id, + format=args.format, + token_env=args.token_env, + private=_str2bool(args.private), + ) sys.exit(0) - else: - print(f"Unknown command: {command}") - sys.exit(1) + parser.print_help() + sys.exit(1) diff --git a/agent/core/telemetry.py b/agent/core/telemetry.py new file mode 100644 index 000000000..ef6623dd9 --- /dev/null +++ b/agent/core/telemetry.py @@ -0,0 +1,439 @@ +"""All agent observability in one module. + +Every telemetry signal the agent emits — LLM-call usage / cost, hf_jobs +lifecycle, sandbox lifecycle, user feedback, mid-turn heartbeat saves — is +defined here so business-logic files stay free of instrumentation noise. + +Callsites are one-liners:: + + await telemetry.record_llm_call(session, model=..., response=r, ...) + await telemetry.record_hf_job_submit(session, job, args, image=..., job_type="Python") + HeartbeatSaver.maybe_fire(session) + +All ``record_*`` functions emit a single ``Event`` via ``session.send_event`` +and never raise — telemetry is best-effort and must not break the agent. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from typing import Any + +from agent.core.cost_estimation import hf_jobs_price_catalog + +logger = logging.getLogger(__name__) + + +# ── usage extraction ──────────────────────────────────────────────────────── + + +def extract_usage(response_or_chunk: Any) -> dict: + """Flat usage dict from a litellm response or final-chunk usage object. + + Normalizes cache-token details across provider response shapes. Exposed + under the stable keys ``cache_read_tokens`` / ``cache_creation_tokens``. + """ + u = getattr(response_or_chunk, "usage", None) + if u is None and isinstance(response_or_chunk, dict): + u = response_or_chunk.get("usage") + if u is None: + return {} + + def _g(name, default=0): + if isinstance(u, dict): + return u.get(name, default) or default + return getattr(u, name, default) or default + + prompt = _g("prompt_tokens") + completion = _g("completion_tokens") + total = _g("total_tokens") or (prompt + completion) + + cache_read = _g("cache_read_input_tokens") + cache_creation = _g("cache_creation_input_tokens") + details = _g("prompt_tokens_details", None) + + if not cache_read and details is not None: + if isinstance(details, dict): + cache_read = details.get("cached_tokens", 0) or 0 + else: + cache_read = getattr(details, "cached_tokens", 0) or 0 + if not cache_creation and details is not None: + if isinstance(details, dict): + cache_creation = details.get("cache_write_tokens", 0) or 0 + else: + cache_creation = getattr(details, "cache_write_tokens", 0) or 0 + + return { + "prompt_tokens": int(prompt), + "completion_tokens": int(completion), + "total_tokens": int(total), + "cache_read_tokens": int(cache_read), + "cache_creation_tokens": int(cache_creation), + } + + +# ── llm_call ──────────────────────────────────────────────────────────────── + + +async def record_llm_call( + session: Any, + *, + model: str, + response: Any = None, + latency_ms: int, + finish_reason: str | None, + kind: str = "main", +) -> dict: + """Emit an ``llm_call`` event and return the extracted usage dict so + callers can stash it on their result object if they want. + + ``kind`` tags the call site so downstream analytics can break spend + down by category. Values currently emitted by the codebase: + + * ``main`` — agent loop turn (user-facing reply or tool follow-up) + * ``research`` — research sub-agent inner loop (3 call sites) + * ``compaction`` — context-window summary on overflow + * ``effort_probe``— effort cascade walk on rejection / model switch + * ``restore`` — session re-seed summary after a Space restart + + Pre-2026-04-29 only ``main`` calls were instrumented; observed gap on + Cost Explorer was ~67%, with the other 5 call sites accounting for + the rest. Tagging lets us split the dataset's ``total_cost_usd`` by + category and validate against billing data. + + The ``/title`` and ``/health/llm`` diagnostic call sites are intentionally + not instrumented because they have no session context and are tiny. + """ + usage = extract_usage(response) if response is not None else {} + cost_usd = 0.0 + if response is not None: + try: + from litellm import completion_cost + + cost_usd = float(completion_cost(completion_response=response) or 0.0) + except Exception: + cost_usd = 0.0 + from agent.core.session import Event # local import to avoid cycle + + try: + payload = { + "model": model, + "latency_ms": latency_ms, + "finish_reason": finish_reason, + "cost_usd": cost_usd, + "kind": kind, + **usage, + } + await session.send_event( + Event( + event_type="llm_call", + data=payload, + ) + ) + except Exception as e: + logger.debug("record_llm_call failed (non-fatal): %s", e) + return {"cost_usd": cost_usd, **usage} + + +# ── hf_jobs ──────────────────────────────────────────────────────────────── + + +def _infer_push_to_hub(script_or_cmd: Any) -> bool: + if not isinstance(script_or_cmd, str): + return False + return ( + "push_to_hub=True" in script_or_cmd + or "push_to_hub=true" in script_or_cmd + or "hub_model_id" in script_or_cmd + ) + + +async def record_hf_job_submit( + session: Any, + job: Any, + args: dict, + *, + image: str, + job_type: str, +) -> float: + """Emit ``hf_job_submit``. Returns the monotonic start timestamp so the + caller can pass it back into :func:`record_hf_job_complete`.""" + from agent.core.session import Event + + t_start = time.monotonic() + try: + script_text = args.get("script") or args.get("command") or "" + await session.send_event( + Event( + event_type="hf_job_submit", + data={ + "job_id": getattr(job, "id", None), + "job_url": getattr(job, "url", None), + "flavor": args.get("hardware_flavor", "cpu-basic"), + "timeout": args.get("timeout", "30m"), + "job_type": job_type, + "image": image, + "namespace": args.get("namespace"), + "push_to_hub": _infer_push_to_hub(script_text), + }, + ) + ) + except Exception as e: + logger.debug("record_hf_job_submit failed (non-fatal): %s", e) + return t_start + + +async def record_hf_job_complete( + session: Any, + job: Any, + *, + flavor: str, + final_status: str, + submit_ts: float, +) -> dict: + from agent.core.session import Event + + try: + wall_time_s = int(time.monotonic() - submit_ts) + billable_seconds = max(0, wall_time_s) + price_usd_per_hour = None + estimated_cost_usd = None + cost_estimate_source = "unknown_price" + prices = await hf_jobs_price_catalog() + if flavor in prices: + price_usd_per_hour = float(prices[flavor]) + estimated_cost_usd = round( + price_usd_per_hour * (billable_seconds / 3600), + 4, + ) + cost_estimate_source = "runtime_price_catalog" + payload = { + "job_id": getattr(job, "id", None), + "flavor": flavor, + "final_status": final_status, + "wall_time_s": wall_time_s, + "billable_seconds_estimate": billable_seconds, + "price_usd_per_hour": price_usd_per_hour, + "estimated_cost_usd": estimated_cost_usd, + "cost_estimate_source": cost_estimate_source, + } + await session.send_event( + Event( + event_type="hf_job_complete", + data=payload, + ) + ) + return payload + except Exception as e: + logger.debug("record_hf_job_complete failed (non-fatal): %s", e) + return {} + + +# ── sandbox ───────────────────────────────────────────────────────────────── + + +async def record_sandbox_create( + session: Any, + sandbox: Any, + *, + hardware: str, + create_latency_s: int, +) -> None: + from agent.core.session import Event + + try: + # Pin created-at on the session so record_sandbox_destroy can diff. + session._sandbox_created_at = time.monotonic() - create_latency_s + await session.send_event( + Event( + event_type="sandbox_create", + data={ + "sandbox_id": getattr(sandbox, "space_id", None), + "hardware": hardware, + "create_latency_s": int(create_latency_s), + }, + ) + ) + except Exception as e: + logger.debug("record_sandbox_create failed (non-fatal): %s", e) + + +async def record_sandbox_destroy(session: Any, sandbox: Any) -> dict: + from agent.core.session import Event + + try: + created = getattr(session, "_sandbox_created_at", None) + lifetime_s = int(time.monotonic() - created) if created else None + hardware = getattr(session, "sandbox_hardware", None) or "cpu-basic" + estimated_cost_usd = None + try: + from agent.core.cost_estimation import SPACE_PRICE_USD_PER_HOUR + + price_usd_per_hour = SPACE_PRICE_USD_PER_HOUR.get(str(hardware)) + if price_usd_per_hour is not None and lifetime_s is not None: + estimated_cost_usd = round( + float(price_usd_per_hour) * (max(0, lifetime_s) / 3600), + 4, + ) + except Exception: + estimated_cost_usd = None + payload = { + "sandbox_id": getattr(sandbox, "space_id", None), + "hardware": hardware, + "lifetime_s": lifetime_s, + "estimated_cost_usd": estimated_cost_usd, + } + await session.send_event( + Event( + event_type="sandbox_destroy", + data=payload, + ) + ) + return payload + except Exception as e: + logger.debug("record_sandbox_destroy failed (non-fatal): %s", e) + return {} + + +# ── feedback ─────────────────────────────────────────────────────────────── + + +async def record_feedback( + session: Any, + *, + rating: str, + turn_index: int | None = None, + message_id: str | None = None, + comment: str | None = None, +) -> None: + from agent.core.session import Event + + try: + await session.send_event( + Event( + event_type="feedback", + data={ + "rating": rating, + "turn_index": turn_index, + "message_id": message_id, + "comment": (comment or "")[:500], + }, + ) + ) + except Exception as e: + logger.debug("record_feedback failed (non-fatal): %s", e) + + +async def record_pro_cta_click( + session: Any, + *, + source: str, + target: str = "pro_pricing", +) -> None: + from agent.core.session import Event + + try: + await session.send_event( + Event( + event_type="pro_cta_click", + data={"source": source, "target": target}, + ) + ) + except Exception as e: + logger.debug("record_pro_cta_click failed (non-fatal): %s", e) + + +async def record_pro_conversion( + session: Any, + *, + first_seen_at: str | None = None, +) -> None: + """Emit a ``pro_conversion`` event for a user we've previously observed + as non-Pro and now see as Pro for the first time. Detected upstream in + ``MongoSessionStore.mark_pro_seen``; fired into the user's first Pro + session so the rollup picks it up alongside other event-driven KPIs.""" + from agent.core.session import Event + + try: + await session.send_event( + Event( + event_type="pro_conversion", + data={"first_seen_at": first_seen_at}, + ) + ) + except Exception as e: + logger.debug("record_pro_conversion failed (non-fatal): %s", e) + + +async def record_credits_topped_up( + session: Any, + *, + namespace: str | None = None, +) -> None: + """Emit a ``credits_topped_up`` event when an hf_job submits successfully + in a session that previously hit ``jobs_access_blocked`` — i.e. the user + came back from the HF billing top-up flow and unblocked themselves. + Caller is responsible for firing this at most once per session.""" + from agent.core.session import Event + + try: + await session.send_event( + Event( + event_type="credits_topped_up", + data={"namespace": namespace}, + ) + ) + except Exception as e: + logger.debug("record_credits_topped_up failed (non-fatal): %s", e) + + +# ── heartbeat ────────────────────────────────────────────────────────────── + +# Module-level reference set for fire-and-forget heartbeat tasks. asyncio only +# keeps *weak* references to tasks, so the returned Task would otherwise be +# eligible for GC before running — the task gets discarded and the upload +# silently never happens. Hold strong refs until the task completes. +_heartbeat_tasks: set[asyncio.Task] = set() + + +class HeartbeatSaver: + """Time-gated mid-turn flush. + + Called from ``Session.send_event`` after every event. Fires + ``save_and_upload_detached`` in a worker thread at most once per + ``heartbeat_interval_s`` (default 60s). Guards against losing trace data + on long-running turns that crash before ``turn_complete``. + """ + + @staticmethod + def maybe_fire(session: Any) -> None: + if not getattr(session.config, "save_sessions", False): + return + interval = getattr(session.config, "heartbeat_interval_s", 0) or 0 + if interval <= 0: + return + now = time.monotonic() + last = getattr(session, "_last_heartbeat_ts", None) + if last is None: + # Initialise on first event; no save yet. + session._last_heartbeat_ts = now + return + if now - last < interval: + return + session._last_heartbeat_ts = now + repo_id = session.config.session_dataset_repo + try: + task = asyncio.get_running_loop().create_task( + asyncio.to_thread(session.save_and_upload_detached, repo_id) + ) + # Hold a strong reference until the task finishes so asyncio can't + # GC it. ``set.discard`` is a no-op on missing keys → safe callback. + _heartbeat_tasks.add(task) + task.add_done_callback(_heartbeat_tasks.discard) + except RuntimeError: + try: + session.save_and_upload_detached(repo_id) + except Exception as e: + logger.debug("Heartbeat save failed (non-fatal): %s", e) diff --git a/agent/core/tools.py b/agent/core/tools.py index 9bbf91d79..0ac100fae 100644 --- a/agent/core/tools.py +++ b/agent/core/tools.py @@ -8,8 +8,6 @@ from dataclasses import dataclass from typing import Any, Awaitable, Callable, Optional -logger = logging.getLogger(__name__) - from fastmcp import Client from fastmcp.exceptions import ToolError from mcp.types import EmbeddedResource, ImageContent, TextContent @@ -46,22 +44,20 @@ hf_repo_git_handler, ) from agent.tools.jobs_tool import HF_JOBS_TOOL_SPEC, hf_jobs_handler +from agent.tools.notify_tool import NOTIFY_TOOL_SPEC, notify_handler from agent.tools.papers_tool import HF_PAPERS_TOOL_SPEC, hf_papers_handler from agent.tools.plan_tool import PLAN_TOOL_SPEC, plan_tool_handler from agent.tools.research_tool import RESEARCH_TOOL_SPEC, research_handler from agent.tools.sandbox_tool import get_sandbox_tools - -# NOTE: Private HF repo tool disabled - replaced by hf_repo_files and hf_repo_git -# from agent.tools.private_hf_repo_tools import ( -# PRIVATE_HF_REPO_TOOL_SPEC, -# private_hf_repo_handler, -# ) +from agent.tools.web_search_tool import WEB_SEARCH_TOOL_SPEC, web_search_handler # Suppress aiohttp deprecation warning warnings.filterwarnings( "ignore", category=DeprecationWarning, module="aiohttp.connector" ) +logger = logging.getLogger(__name__) + NOT_ALLOWED_TOOL_NAMES = ["hf_jobs", "hf_doc_search", "hf_doc_fetch", "hf_whoami"] @@ -129,7 +125,12 @@ class ToolRouter: Based on codex-rs/core/src/tools/router.rs """ - def __init__(self, mcp_servers: dict[str, MCPServerConfig], hf_token: str | None = None, local_mode: bool = False): + def __init__( + self, + mcp_servers: dict[str, MCPServerConfig], + hf_token: str | None = None, + local_mode: bool = False, + ): self.tools: dict[str, ToolSpec] = {} self.mcp_servers: dict[str, dict[str, Any]] = {} @@ -142,7 +143,9 @@ def __init__(self, mcp_servers: dict[str, MCPServerConfig], hf_token: str | None for name, server in mcp_servers.items(): data = server.model_dump() if hf_token: - data.setdefault("headers", {})["Authorization"] = f"Bearer {hf_token}" + data.setdefault("headers", {})["Authorization"] = ( + f"Bearer {hf_token}" + ) mcp_servers_payload[name] = data self.mcp_client = Client({"mcpServers": mcp_servers_payload}) self._mcp_initialized = False @@ -216,7 +219,9 @@ async def __aenter__(self) -> "ToolRouter": await self.register_mcp_tools() self._mcp_initialized = True except Exception as e: - logger.warning("MCP connection failed, continuing without MCP tools: %s", e) + logger.warning( + "MCP connection failed, continuing without MCP tools: %s", e + ) self.mcp_client = None await self.register_openapi_tool() @@ -310,6 +315,12 @@ def create_builtin_tools(local_mode: bool = False) -> list[ToolSpec]: parameters=HF_PAPERS_TOOL_SPEC["parameters"], handler=hf_papers_handler, ), + ToolSpec( + name=WEB_SEARCH_TOOL_SPEC["name"], + description=WEB_SEARCH_TOOL_SPEC["description"], + parameters=WEB_SEARCH_TOOL_SPEC["parameters"], + handler=web_search_handler, + ), # Dataset inspection tool (unified) ToolSpec( name=HF_INSPECT_DATASET_TOOL_SPEC["name"], @@ -324,6 +335,12 @@ def create_builtin_tools(local_mode: bool = False) -> list[ToolSpec]: parameters=PLAN_TOOL_SPEC["parameters"], handler=plan_tool_handler, ), + ToolSpec( + name=NOTIFY_TOOL_SPEC["name"], + description=NOTIFY_TOOL_SPEC["description"], + parameters=NOTIFY_TOOL_SPEC["parameters"], + handler=notify_handler, + ), ToolSpec( name=HF_JOBS_TOOL_SPEC["name"], description=HF_JOBS_TOOL_SPEC["description"], @@ -366,6 +383,7 @@ def create_builtin_tools(local_mode: bool = False) -> list[ToolSpec]: # Sandbox or local tools (highest priority) if local_mode: from agent.tools.local_tools import get_local_tools + tools = get_local_tools() + tools else: tools = get_sandbox_tools() + tools diff --git a/agent/core/usage_metrics.py b/agent/core/usage_metrics.py new file mode 100644 index 000000000..f93c0a42d --- /dev/null +++ b/agent/core/usage_metrics.py @@ -0,0 +1,448 @@ +"""Pure usage/billing summaries for session trajectory analytics.""" + +from collections import Counter, defaultdict +from datetime import UTC, datetime, timedelta +from math import isfinite +from typing import Any + +from agent.core.cost_estimation import SPACE_PRICE_USD_PER_HOUR + +USAGE_METRICS_VERSION = 1 +BILLING_SCOPE_ACCOUNT_WINDOW_DELTA = "account_window_delta" + +_USAGE_SCALAR_KEYS = ( + "usage_total_usd", + "usage_total_usd_source", + "usage_app_total_usd", + "usage_hf_billing_total_usd", + "usage_llm_calls", + "usage_total_tokens", + "usage_hf_job_submits", + "usage_hf_job_status_snapshots", + "usage_sandbox_creates", + "usage_sandbox_pairs", +) + + +def _coerce_float(value: Any) -> float: + if isinstance(value, bool) or value is None: + return 0.0 + try: + parsed = float(value) + except (TypeError, ValueError): + return 0.0 + return parsed if isfinite(parsed) else 0.0 + + +def _coerce_optional_float(value: Any) -> float | None: + if isinstance(value, bool) or value is None: + return None + try: + parsed = float(value) + except (TypeError, ValueError): + return None + return parsed if isfinite(parsed) else None + + +def _coerce_int(value: Any) -> int: + if isinstance(value, bool) or value is None: + return 0 + try: + return int(value) + except (TypeError, ValueError): + return 0 + + +def _round_usd(value: Any) -> float: + return round(_coerce_float(value), 6) + + +def _parse_timestamp(value: Any) -> datetime | None: + if isinstance(value, datetime): + dt = value + elif isinstance(value, str) and value: + try: + dt = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + else: + return None + if dt.tzinfo is None: + return dt.replace(tzinfo=UTC) + return dt.astimezone(UTC) + + +def event_created_at(event: dict[str, Any]) -> datetime | None: + return _parse_timestamp(event.get("created_at") or event.get("timestamp")) + + +def _event_data(event: dict[str, Any]) -> dict[str, Any]: + data = event.get("data") or {} + return data if isinstance(data, dict) else {} + + +def _has_number(value: Any) -> bool: + return _coerce_optional_float(value) is not None + + +def _counter_dict(counter: Counter[str]) -> dict[str, int]: + return dict(sorted(counter.items())) + + +def _empty_app_bucket(session_id: str | None) -> dict[str, Any]: + return { + "session_id": session_id, + "total_usd": 0.0, + "inference_usd": 0.0, + "hf_jobs_estimated_usd": 0.0, + "sandbox_estimated_usd": 0.0, + "llm_calls": 0, + "hf_jobs_count": 0, + "sandbox_count": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "cache_read_tokens": 0, + "cache_creation_tokens": 0, + "total_tokens": 0, + "hf_jobs_billable_seconds_estimate": 0, + "sandbox_billable_seconds_estimate": 0, + } + + +def _sandbox_id(event: dict[str, Any]) -> str | None: + sandbox_id = _event_data(event).get("sandbox_id") + return sandbox_id if isinstance(sandbox_id, str) and sandbox_id else None + + +def _sandbox_duration_seconds( + create_event: dict[str, Any], + destroy_event: dict[str, Any], +) -> int: + create_data = _event_data(create_event) + destroy_data = _event_data(destroy_event) + lifetime_s = _coerce_int(destroy_data.get("lifetime_s")) + if lifetime_s > 0: + return lifetime_s + + create_at = event_created_at(create_event) + destroy_at = event_created_at(destroy_event) + if create_at is None or destroy_at is None: + return 0 + create_latency_s = max(0, _coerce_int(create_data.get("create_latency_s"))) + interval_start = create_at - timedelta(seconds=create_latency_s) + if destroy_at <= interval_start: + return 0 + return int((destroy_at - interval_start).total_seconds()) + + +def summarize_sandbox_lifecycle( + lifecycle_events: list[tuple[int, dict[str, Any]]], +) -> dict[str, Any]: + """Pair sandbox lifecycle events and estimate billed usage. + + Shared by dataset usage metrics and backend usage responses so sandbox + pricing and create/destroy pairing semantics cannot drift. + """ + ordered_events = [ + event + for _, event in sorted( + lifecycle_events, + key=lambda indexed: ( + event_created_at(indexed[1]) is None, + event_created_at(indexed[1]) or datetime.min.replace(tzinfo=UTC), + indexed[0], + ), + ) + ] + active_creates: dict[str, list[dict[str, Any]]] = defaultdict(list) + matched_pairs = 0 + unpaired_destroys = 0 + estimated_usd = 0.0 + billable_seconds = 0 + + for event in ordered_events: + event_type = event.get("event_type") + sandbox_id = _sandbox_id(event) + if sandbox_id is None: + continue + if event_type == "sandbox_create": + active_creates[sandbox_id].append(event) + continue + if event_type != "sandbox_destroy": + continue + + creates = active_creates.get(sandbox_id) + if not creates: + unpaired_destroys += 1 + continue + + create_event = creates.pop() + if not creates: + active_creates.pop(sandbox_id, None) + + hardware = str(_event_data(create_event).get("hardware") or "cpu-basic") + seconds = _sandbox_duration_seconds(create_event, event) + price_usd_per_hour = _coerce_float(SPACE_PRICE_USD_PER_HOUR.get(hardware)) + matched_pairs += 1 + if price_usd_per_hour > 0: + billable_seconds += seconds + estimated_usd += price_usd_per_hour * (seconds / 3600) + + return { + "matched_pairs": matched_pairs, + "unpaired_creates": sum(len(events) for events in active_creates.values()), + "unpaired_destroys": unpaired_destroys, + "estimated_usd": _round_usd(estimated_usd), + "billable_seconds_estimate": billable_seconds, + } + + +def normalize_hf_billing_snapshot(snapshot: dict[str, Any] | None) -> dict[str, Any]: + """Return a dataset-safe HF billing snapshot. + + Only current-session window rollups are retained. Monthly account totals, + credit limits, and any caller-provided extra fields are intentionally + dropped before the snapshot can be serialized into session artifacts. + """ + hf_billing = snapshot.get("hf_billing") if isinstance(snapshot, dict) else None + hf_billing = hf_billing if isinstance(hf_billing, dict) else {} + current_session = hf_billing.get("current_session") + current_session = current_session if isinstance(current_session, dict) else None + + sanitized_current = None + if current_session is not None: + sanitized_current = { + "window_start": current_session.get("window_start"), + "window_end": current_session.get("window_end"), + "timezone": current_session.get("timezone"), + "total_usd": _round_usd(current_session.get("total_usd")), + "inference_providers_usd": _round_usd( + current_session.get("inference_providers_usd") + ), + "hf_jobs_usd": _round_usd(current_session.get("hf_jobs_usd")), + "inference_provider_requests": _coerce_int( + current_session.get("inference_provider_requests") + ), + "hf_jobs_minutes": round( + _coerce_float(current_session.get("hf_jobs_minutes")), 3 + ), + } + + available = bool(hf_billing.get("available") and sanitized_current is not None) + return { + "billing_scope": BILLING_SCOPE_ACCOUNT_WINDOW_DELTA, + "hf_billing": { + "source": str(hf_billing.get("source") or "hf_billing_usage_v2"), + "available": available, + "error": None if available else hf_billing.get("error"), + "current_session": sanitized_current if available else None, + }, + } + + +def summarize_usage_events( + events: list[dict[str, Any]], + *, + session_id: str | None = None, + hf_billing_snapshot: dict[str, Any] | None = None, +) -> dict[str, Any]: + app = _empty_app_bucket(session_id) + llm_by_kind: Counter[str] = Counter() + llm_by_model: Counter[str] = Counter() + job_statuses: Counter[str] = Counter() + job_submit_flavors: Counter[str] = Counter() + job_status_flavors: Counter[str] = Counter() + sandbox_hardware: Counter[str] = Counter() + lifecycle_events: list[tuple[int, dict[str, Any]]] = [] + + event_count = 0 + events_without_timestamp = 0 + llm_calls_with_cost_usd = 0 + llm_calls_with_nonzero_cost_usd = 0 + job_submits = 0 + job_status_snapshots = 0 + job_snapshots_with_estimated_cost = 0 + job_snapshots_with_nonzero_estimated_cost = 0 + sandbox_creates = 0 + sandbox_destroys = 0 + turn_complete_count = 0 + assistant_stream_end_count = 0 + + for index, event in enumerate(events or []): + if not isinstance(event, dict): + continue + event_count += 1 + if event_created_at(event) is None: + events_without_timestamp += 1 + + event_type = event.get("event_type") + data = _event_data(event) + if event_type == "llm_call": + app["llm_calls"] += 1 + if "cost_usd" in data: + llm_calls_with_cost_usd += 1 + cost_usd = _coerce_float(data.get("cost_usd")) + if cost_usd > 0: + llm_calls_with_nonzero_cost_usd += 1 + app["inference_usd"] += cost_usd + + prompt_tokens = _coerce_int(data.get("prompt_tokens")) + completion_tokens = _coerce_int(data.get("completion_tokens")) + cache_read_tokens = _coerce_int(data.get("cache_read_tokens")) + cache_creation_tokens = _coerce_int(data.get("cache_creation_tokens")) + total_tokens = _coerce_int(data.get("total_tokens")) or ( + prompt_tokens + + completion_tokens + + cache_read_tokens + + cache_creation_tokens + ) + app["prompt_tokens"] += prompt_tokens + app["completion_tokens"] += completion_tokens + app["cache_read_tokens"] += cache_read_tokens + app["cache_creation_tokens"] += cache_creation_tokens + app["total_tokens"] += total_tokens + llm_by_kind[str(data.get("kind") or "unknown")] += 1 + llm_by_model[str(data.get("model") or "unknown")] += 1 + elif event_type == "hf_job_submit": + job_submits += 1 + job_submit_flavors[str(data.get("flavor") or "unknown")] += 1 + elif event_type == "hf_job_complete": + job_status_snapshots += 1 + app["hf_jobs_count"] += 1 + estimated_cost = _coerce_float(data.get("estimated_cost_usd")) + app["hf_jobs_estimated_usd"] += estimated_cost + app["hf_jobs_billable_seconds_estimate"] += _coerce_int( + data.get("billable_seconds_estimate") or data.get("wall_time_s") + ) + if _has_number(data.get("estimated_cost_usd")): + job_snapshots_with_estimated_cost += 1 + if estimated_cost > 0: + job_snapshots_with_nonzero_estimated_cost += 1 + job_statuses[str(data.get("final_status") or "unknown")] += 1 + job_status_flavors[str(data.get("flavor") or "unknown")] += 1 + elif event_type == "sandbox_create": + sandbox_creates += 1 + sandbox_hardware[str(data.get("hardware") or "cpu-basic")] += 1 + lifecycle_events.append((index, event)) + elif event_type == "sandbox_destroy": + sandbox_destroys += 1 + lifecycle_events.append((index, event)) + elif event_type == "turn_complete": + turn_complete_count += 1 + elif event_type == "assistant_stream_end": + assistant_stream_end_count += 1 + + sandbox = summarize_sandbox_lifecycle(lifecycle_events) + app["sandbox_count"] = sandbox["matched_pairs"] + app["sandbox_estimated_usd"] = sandbox["estimated_usd"] + app["sandbox_billable_seconds_estimate"] = sandbox["billable_seconds_estimate"] + app["inference_usd"] = _round_usd(app["inference_usd"]) + app["hf_jobs_estimated_usd"] = _round_usd(app["hf_jobs_estimated_usd"]) + app["total_usd"] = _round_usd( + app["inference_usd"] + + app["hf_jobs_estimated_usd"] + + app["sandbox_estimated_usd"] + ) + + billing = normalize_hf_billing_snapshot(hf_billing_snapshot) + current_billing = billing["hf_billing"]["current_session"] + hf_billing_total = None + if billing["hf_billing"]["available"] and current_billing is not None: + hf_billing_total = _round_usd(current_billing.get("total_usd")) + usage_total = _round_usd(hf_billing_total + app["sandbox_estimated_usd"]) + usage_total_source = "hf_billing_plus_sandbox_estimate" + else: + usage_total = app["total_usd"] + usage_total_source = "app_telemetry_fallback" + + job_flavors = job_submit_flavors + job_status_flavors + + return { + "version": USAGE_METRICS_VERSION, + "session_id": session_id, + "billing_scope": BILLING_SCOPE_ACCOUNT_WINDOW_DELTA, + "total_usd": usage_total, + "total_usd_source": usage_total_source, + "app_total_usd": app["total_usd"], + "hf_billing_total_usd": hf_billing_total, + "app_telemetry": app, + "hf_billing": billing["hf_billing"], + "llm": { + "calls": app["llm_calls"], + "calls_by_kind": _counter_dict(llm_by_kind), + "calls_by_model": _counter_dict(llm_by_model), + "prompt_tokens": app["prompt_tokens"], + "completion_tokens": app["completion_tokens"], + "cache_read_tokens": app["cache_read_tokens"], + "cache_creation_tokens": app["cache_creation_tokens"], + "total_tokens": app["total_tokens"], + }, + "turns": { + "turn_complete_count": turn_complete_count, + "assistant_stream_end_count": assistant_stream_end_count, + }, + "hf_jobs": { + "submits": job_submits, + "status_snapshots": job_status_snapshots, + "statuses": _counter_dict(job_statuses), + "flavors": _counter_dict(job_flavors), + "submit_flavors": _counter_dict(job_submit_flavors), + "status_snapshot_flavors": _counter_dict(job_status_flavors), + "estimated_usd": app["hf_jobs_estimated_usd"], + "billable_seconds_estimate": app["hf_jobs_billable_seconds_estimate"], + "snapshots_with_estimated_cost": job_snapshots_with_estimated_cost, + "snapshots_with_nonzero_estimated_cost": ( + job_snapshots_with_nonzero_estimated_cost + ), + }, + "sandboxes": { + "creates": sandbox_creates, + "destroys": sandbox_destroys, + "matched_pairs": sandbox["matched_pairs"], + "unpaired_creates": sandbox["unpaired_creates"], + "unpaired_destroys": sandbox["unpaired_destroys"], + "hardware": _counter_dict(sandbox_hardware), + "estimated_usd": app["sandbox_estimated_usd"], + "billable_seconds_estimate": app["sandbox_billable_seconds_estimate"], + }, + "data_quality": { + "event_count": event_count, + "events_without_timestamp": events_without_timestamp, + "llm_calls_with_cost_usd": llm_calls_with_cost_usd, + "llm_calls_with_nonzero_cost_usd": llm_calls_with_nonzero_cost_usd, + "job_snapshots_with_estimated_cost": job_snapshots_with_estimated_cost, + "job_snapshots_missing_estimated_cost": ( + job_status_snapshots - job_snapshots_with_estimated_cost + ), + }, + } + + +def usage_metric_scalar_fields(metrics: dict[str, Any]) -> dict[str, Any]: + app = metrics.get("app_telemetry") if isinstance(metrics, dict) else {} + llm = metrics.get("llm") if isinstance(metrics, dict) else {} + jobs = metrics.get("hf_jobs") if isinstance(metrics, dict) else {} + sandboxes = metrics.get("sandboxes") if isinstance(metrics, dict) else {} + values = { + "usage_total_usd": metrics.get("total_usd"), + "usage_total_usd_source": metrics.get("total_usd_source"), + "usage_app_total_usd": metrics.get("app_total_usd"), + "usage_hf_billing_total_usd": metrics.get("hf_billing_total_usd"), + "usage_llm_calls": app.get("llm_calls") if isinstance(app, dict) else None, + "usage_total_tokens": llm.get("total_tokens") + if isinstance(llm, dict) + else None, + "usage_hf_job_submits": ( + jobs.get("submits") if isinstance(jobs, dict) else None + ), + "usage_hf_job_status_snapshots": ( + jobs.get("status_snapshots") if isinstance(jobs, dict) else None + ), + "usage_sandbox_creates": ( + sandboxes.get("creates") if isinstance(sandboxes, dict) else None + ), + "usage_sandbox_pairs": ( + sandboxes.get("matched_pairs") if isinstance(sandboxes, dict) else None + ), + } + return {key: values.get(key) for key in _USAGE_SCALAR_KEYS} diff --git a/agent/core/usage_thresholds.py b/agent/core/usage_thresholds.py new file mode 100644 index 000000000..effba12ba --- /dev/null +++ b/agent/core/usage_thresholds.py @@ -0,0 +1,55 @@ +"""Helpers for session usage-threshold approval warnings.""" + +from typing import Any + +USAGE_THRESHOLD_TOOL_NAME = "usage_threshold" +USAGE_WARNING_FIRST_THRESHOLD_USD = 5.0 +USAGE_WARNING_MULTIPLIER = 2.0 + + +def normalize_usage_threshold(value: Any) -> float: + """Return a usable positive threshold, defaulting to the first warning.""" + if isinstance(value, bool): + return USAGE_WARNING_FIRST_THRESHOLD_USD + try: + threshold = float(value) + except (TypeError, ValueError): + return USAGE_WARNING_FIRST_THRESHOLD_USD + if threshold <= 0: + return USAGE_WARNING_FIRST_THRESHOLD_USD + return threshold + + +def next_usage_warning_threshold( + current_spend_usd: float, + acknowledged_threshold_usd: float, +) -> float: + """Advance the next threshold until it is above the current spend.""" + threshold = normalize_usage_threshold(acknowledged_threshold_usd) + current = max(0.0, float(current_spend_usd or 0.0)) + while threshold <= current: + threshold *= USAGE_WARNING_MULTIPLIER + return round(threshold, 4) + + +def is_usage_threshold_pending(pending_approval: Any) -> bool: + return ( + isinstance(pending_approval, dict) + and pending_approval.get("kind") == USAGE_THRESHOLD_TOOL_NAME + ) + + +def usage_threshold_pending_to_tool(pending_approval: dict[str, Any]) -> dict[str, Any]: + """Represent a synthetic usage approval as the existing pending-tool shape.""" + tool_call_id = str(pending_approval.get("tool_call_id") or "") + arguments = { + "threshold_usd": pending_approval.get("threshold_usd"), + "current_spend_usd": pending_approval.get("current_spend_usd"), + "next_threshold_usd": pending_approval.get("next_threshold_usd"), + "billing_source": pending_approval.get("billing_source"), + } + return { + "tool": USAGE_THRESHOLD_TOOL_NAME, + "tool_call_id": tool_call_id, + "arguments": arguments, + } diff --git a/agent/core/yolo_budget.py b/agent/core/yolo_budget.py new file mode 100644 index 000000000..c8ed1684e --- /dev/null +++ b/agent/core/yolo_budget.py @@ -0,0 +1,403 @@ +"""Session-scoped YOLO budget guardrails.""" + +import uuid +from dataclasses import dataclass +from typing import Any + +from agent.core.cost_estimation import CostEstimate + +YOLO_BUDGET_TOOL_NAME = "yolo_budget" + + +@dataclass(frozen=True) +class BudgetReservation: + reservation_id: str + amount_usd: float + spend_kind: str + + +@dataclass(frozen=True) +class BudgetDecision: + allowed: bool + estimated_cost_usd: float | None = None + remaining_cap_usd: float | None = None + block_reason: str | None = None + billable: bool = False + reservation: BudgetReservation | None = None + + +def session_yolo_enabled(session: Any | None) -> bool: + return bool(session and getattr(session, "auto_approval_enabled", False)) + + +def session_spend_usd(session: Any | None) -> float: + if not session: + return 0.0 + return max( + 0.0, + float(getattr(session, "auto_approval_estimated_spend_usd", 0.0) or 0.0), + ) + + +def session_remaining_usd( + session: Any | None, reserved_spend_usd: float = 0.0 +) -> float | None: + if not session or getattr(session, "auto_approval_cost_cap_usd", None) is None: + return None + cap = float(getattr(session, "auto_approval_cost_cap_usd") or 0.0) + return round(max(0.0, cap - session_spend_usd(session) - reserved_spend_usd), 4) + + +def _set_session_spend(session: Any, amount_usd: float) -> None: + session.auto_approval_estimated_spend_usd = round(max(0.0, amount_usd), 4) + + +def add_session_spend(session: Any, amount_usd: float | None) -> None: + if amount_usd is None or amount_usd <= 0: + return + if hasattr(session, "add_auto_approval_estimated_spend"): + session.add_auto_approval_estimated_spend(amount_usd) + else: + _set_session_spend(session, session_spend_usd(session) + float(amount_usd)) + + +def adjust_session_spend(session: Any, delta_usd: float | None) -> None: + if delta_usd is None or delta_usd == 0: + return + _set_session_spend(session, session_spend_usd(session) + float(delta_usd)) + + +def seed_session_spend(session: Any, amount_usd: float | None) -> None: + if amount_usd is None: + return + _set_session_spend(session, max(session_spend_usd(session), float(amount_usd))) + + +def _cap_usd(session: Any | None) -> float | None: + if not session or getattr(session, "auto_approval_cost_cap_usd", None) is None: + return None + return max(0.0, float(getattr(session, "auto_approval_cost_cap_usd") or 0.0)) + + +def _reservation_store(session: Any) -> dict[str, BudgetReservation]: + store = getattr(session, "_yolo_budget_reservations", None) + if not isinstance(store, dict): + store = {} + setattr(session, "_yolo_budget_reservations", store) + return store + + +def _coerce_cost(value: Any) -> float | None: + if isinstance(value, bool) or value is None: + return None + try: + return max(0.0, float(value)) + except (TypeError, ValueError): + return None + + +def check_session_budget( + session: Any | None, + estimate: CostEstimate, + *, + reserved_spend_usd: float = 0.0, +) -> BudgetDecision: + if not session_yolo_enabled(session) or not estimate.billable: + return BudgetDecision( + allowed=True, + estimated_cost_usd=estimate.estimated_cost_usd, + billable=estimate.billable, + ) + + remaining = session_remaining_usd(session, reserved_spend_usd=reserved_spend_usd) + amount = _coerce_cost(estimate.estimated_cost_usd) + if amount is None: + return BudgetDecision( + allowed=False, + estimated_cost_usd=None, + remaining_cap_usd=remaining, + block_reason=estimate.block_reason + or "Could not estimate this session spend safely.", + billable=True, + ) + if remaining is not None and amount > remaining: + return BudgetDecision( + allowed=False, + estimated_cost_usd=round(amount, 4), + remaining_cap_usd=remaining, + block_reason=( + f"Estimated cost ${amount:.2f} exceeds remaining YOLO cap " + f"${remaining:.2f}." + ), + billable=True, + ) + return BudgetDecision( + allowed=True, + estimated_cost_usd=round(amount, 4), + remaining_cap_usd=remaining, + billable=True, + ) + + +def reserve_session_budget( + session: Any | None, + estimate: CostEstimate, + *, + spend_kind: str, + reservation_id: str | None = None, +) -> BudgetDecision: + decision = check_session_budget(session, estimate) + if not session or not session_yolo_enabled(session) or not decision.billable: + return decision + if not decision.allowed: + return decision + amount = _coerce_cost(decision.estimated_cost_usd) + if amount is None or amount <= 0: + return decision + + add_session_spend(session, amount) + rid = reservation_id or f"{spend_kind}-{uuid.uuid4().hex[:10]}" + reservation = BudgetReservation( + reservation_id=rid, + amount_usd=round(amount, 4), + spend_kind=spend_kind, + ) + _reservation_store(session)[rid] = reservation + return BudgetDecision( + allowed=True, + estimated_cost_usd=round(amount, 4), + remaining_cap_usd=session_remaining_usd(session), + billable=True, + reservation=reservation, + ) + + +def release_budget_reservation(session: Any | None, reservation_id: str | None) -> None: + if not session or not reservation_id: + return + reservation = _reservation_store(session).pop(reservation_id, None) + if reservation is None: + return + adjust_session_spend(session, -reservation.amount_usd) + + +def reconcile_budget_reservation( + session: Any | None, + reservation_id: str | None, + actual_cost_usd: Any, + *, + allow_zero_actual: bool = False, +) -> None: + if not session or not reservation_id: + return + reservation = _reservation_store(session).pop(reservation_id, None) + if reservation is None: + return + actual = _coerce_cost(actual_cost_usd) + if actual is None or (actual == 0 and not allow_zero_actual): + return + adjust_session_spend(session, actual - reservation.amount_usd) + + +def is_yolo_budget_pending(pending_approval: Any) -> bool: + return ( + isinstance(pending_approval, dict) + and pending_approval.get("kind") == YOLO_BUDGET_TOOL_NAME + ) + + +def yolo_budget_pending_to_tool(pending_approval: dict[str, Any]) -> dict[str, Any]: + tool_call_id = str(pending_approval.get("tool_call_id") or "") + arguments = { + "cap_usd": pending_approval.get("cap_usd"), + "current_spend_usd": pending_approval.get("current_spend_usd"), + "remaining_cap_usd": pending_approval.get("remaining_cap_usd"), + "estimated_next_usd": pending_approval.get("estimated_next_usd"), + "spend_kind": pending_approval.get("spend_kind"), + "reason": pending_approval.get("reason"), + } + return { + "tool": YOLO_BUDGET_TOOL_NAME, + "tool_call_id": tool_call_id, + "arguments": arguments, + "auto_approval_blocked": True, + "block_reason": pending_approval.get("reason"), + "estimated_cost_usd": pending_approval.get("estimated_next_usd"), + "remaining_cap_usd": pending_approval.get("remaining_cap_usd"), + } + + +async def request_yolo_budget_approval( + session: Any, + decision: BudgetDecision, + *, + spend_kind: str, + current_spend_usd: float | None = None, + cap_usd: float | None = None, + billing_source: str | None = None, + continuation: str | None = None, + final_response: str | None = None, + history_size: int | None = None, +) -> bool: + if session.pending_approval: + return False + from agent.core.session import Event + + current_spend = ( + session_spend_usd(session) + if current_spend_usd is None + else max(0.0, float(current_spend_usd)) + ) + cap = getattr(session, "auto_approval_cost_cap_usd", None) + if cap_usd is not None: + cap = max(0.0, float(cap_usd)) + pending = { + "kind": YOLO_BUDGET_TOOL_NAME, + "tool_call_id": f"yolo-budget-{uuid.uuid4().hex[:10]}", + "cap_usd": cap, + "current_spend_usd": round(current_spend, 6), + "remaining_cap_usd": decision.remaining_cap_usd, + "estimated_next_usd": decision.estimated_cost_usd, + "spend_kind": spend_kind, + "reason": decision.block_reason or "YOLO budget requires confirmation.", + "history_size": history_size + if history_size is not None + else len(session.context_manager.items), + } + if billing_source: + pending["billing_source"] = billing_source + if continuation: + pending["continuation"] = continuation + if isinstance(final_response, str): + pending["final_response"] = final_response + session.pending_approval = pending + tool = yolo_budget_pending_to_tool(pending) + await session.send_event( + Event( + event_type="approval_required", + data={ + "tools": [tool], + "count": 1, + "yolo_budget": True, + "auto_approval_blocked": True, + "block_reason": pending["reason"], + "estimated_cost_usd": pending["estimated_next_usd"], + "remaining_cap_usd": pending["remaining_cap_usd"], + }, + ) + ) + return True + + +async def request_yolo_budget_exceeded_approval( + session: Any, + *, + spend_kind: str, + current_spend_usd: float, + cap_usd: float, + billing_source: str | None = None, + reason: str | None = None, + continuation: str | None = None, + final_response: str | None = None, + history_size: int | None = None, +) -> bool: + current_spend = max(0.0, float(current_spend_usd)) + cap = max(0.0, float(cap_usd)) + seed_session_spend(session, current_spend) + if not session_yolo_enabled(session) or current_spend < cap: + return False + decision = BudgetDecision( + allowed=False, + estimated_cost_usd=None, + remaining_cap_usd=round(max(0.0, cap - current_spend), 4), + block_reason=reason + or ( + "YOLO cap paused session usage after " + f"{spend_kind}: current session spend ${current_spend:.2f} " + f"has reached the ${cap:.2f} cap." + ), + billable=True, + ) + return await request_yolo_budget_approval( + session, + decision, + spend_kind=spend_kind, + current_spend_usd=current_spend, + cap_usd=cap, + billing_source=billing_source, + continuation=continuation, + final_response=final_response, + history_size=history_size, + ) + + +async def maybe_pause_yolo_after_spend( + session: Any | None, + *, + spend_kind: str, + observed_cost_usd: Any = None, + continuation: str | None = None, + final_response: str | None = None, +) -> bool: + if not session or not session_yolo_enabled(session) or session.pending_approval: + return False + + observed = _coerce_cost(observed_cost_usd) + if observed is not None and observed > 0: + add_session_spend(session, observed) + + checker = getattr(session, "yolo_budget_checker", None) + if checker is not None: + try: + return bool( + await checker( + { + "spend_kind": spend_kind, + "observed_cost_usd": observed, + "continuation": continuation, + "final_response": final_response, + "history_size": len(session.context_manager.items), + } + ) + ) + except Exception: + pass + + cap = _cap_usd(session) + current_spend = session_spend_usd(session) + if cap is None or current_spend < cap: + return False + return await request_yolo_budget_exceeded_approval( + session, + spend_kind=spend_kind, + current_spend_usd=current_spend, + cap_usd=cap, + continuation=continuation, + final_response=final_response, + history_size=len(session.context_manager.items), + ) + + +def yolo_budget_can_resume( + session: Any, pending: dict[str, Any] +) -> tuple[bool, str | None]: + if not session_yolo_enabled(session): + return True, None + estimated_next = _coerce_cost(pending.get("estimated_next_usd")) + remaining = session_remaining_usd(session) + if estimated_next is None: + if remaining is None or remaining > 0: + return True, None + return ( + False, + str( + pending.get("reason") + or "YOLO cap is reached. Raise or disable the cap to continue." + ), + ) + if remaining is not None and estimated_next > remaining: + return ( + False, + f"Estimated cost ${estimated_next:.2f} exceeds remaining YOLO cap ${remaining:.2f}.", + ) + return True, None diff --git a/agent/main.py b/agent/main.py index 4ecbefc50..6f29c06b8 100644 --- a/agent/main.py +++ b/agent/main.py @@ -9,8 +9,10 @@ import argparse import asyncio import json +import logging import os import signal +import subprocess import sys import time from dataclasses import dataclass @@ -21,10 +23,16 @@ from prompt_toolkit import PromptSession from agent.config import load_config +from agent.core.approval_policy import is_scheduled_operation from agent.core.agent_loop import submission_loop from agent.core import model_switcher +from agent.core.hf_access import fetch_whoami_v2, normalize_hf_user_plan +from agent.core.hf_tokens import resolve_hf_token +from agent.core.local_models import is_local_model_id +from agent.core.model_ids import strip_huggingface_model_prefix from agent.core.session import OpType from agent.core.tools import ToolRouter +from agent.messaging.gateway import NotificationGateway from agent.utils.reliability_checks import check_training_script_save_pattern from agent.utils.terminal_display import ( get_console, @@ -50,6 +58,75 @@ # on every error — users don't need it, and our friendly errors cover the case. litellm.suppress_debug_info = True +CLI_CONFIG_PATH = Path(__file__).parent.parent / "configs" / "cli_agent_config.json" +logger = logging.getLogger(__name__) + + +def _apply_tool_runtime_override(config: Any, *, sandbox_tools: bool) -> str: + if sandbox_tools: + config.tool_runtime = "sandbox" + return getattr(config, "tool_runtime", "local") + + +def _is_local_tool_runtime(config: Any) -> bool: + return getattr(config, "tool_runtime", "local") == "local" + + +def _tool_runtime_label(local_mode: bool) -> str: + return "local filesystem" if local_mode else "HF sandbox" + + +def _normalize_config_model(config: Any) -> None: + normalized = strip_huggingface_model_prefix(getattr(config, "model_name", None)) + if normalized: + config.model_name = normalized + + +def _validate_cli_model_override(model: str) -> str: + if not model_switcher.is_valid_model_id(model): + raise ValueError( + "Invalid model id. Use an HF Router id like " + "'zai-org/GLM-5.2:novita' or a supported local prefix." + ) + return model.removeprefix("huggingface/") + + +async def _wait_for_initial_sandbox_preload(session_holder: list | None) -> None: + session = session_holder[0] if session_holder else None + task = getattr(session, "sandbox_preload_task", None) + if not task: + return + try: + await asyncio.shield(task) + except asyncio.CancelledError: + raise + except Exception: + # The sandbox tool will surface the stored preload error on first use. + return + + +def _is_scheduled_hf_job_tool(tool_info: dict[str, Any]) -> bool: + if tool_info.get("tool") != "hf_jobs": + return False + arguments = tool_info.get("arguments") or {} + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + return False + if not isinstance(arguments, dict): + return False + return is_scheduled_operation(arguments.get("operation")) + + +def _configure_runtime_logging() -> None: + """Keep third-party warning spam from punching through the interactive UI.""" + import logging + + logging.getLogger("LiteLLM").setLevel(logging.ERROR) + logging.getLogger("litellm").setLevel(logging.ERROR) + + def _safe_get_args(arguments: dict) -> dict: """Safely extract args dict from arguments, handling cases where LLM passes string.""" args = arguments.get("args", {}) @@ -59,28 +136,37 @@ def _safe_get_args(arguments: dict) -> dict: return args if isinstance(args, dict) else {} -def _get_hf_token() -> str | None: - """Get HF token from environment, huggingface_hub API, or cached token file.""" - token = os.environ.get("HF_TOKEN") - if token: - return token +def _get_hf_user(token: str | None) -> str | None: + """Resolve the HF username for a token, if available.""" + if not token: + return None try: from huggingface_hub import HfApi - api = HfApi() - token = api.token - if token: - return token + + return HfApi(token=token).whoami().get("name") except Exception: - pass - # Fallback: read the cached token file directly - token_path = Path.home() / ".cache" / "huggingface" / "token" - if token_path.exists(): - token = token_path.read_text().strip() - if token: - return token + return None + + +def _get_hf_user_from_whoami(whoami: dict[str, Any] | None) -> str | None: + if not isinstance(whoami, dict): + return None + for key in ("name", "user", "preferred_username"): + value = whoami.get(key) + if isinstance(value, str) and value: + return value return None +async def _get_hf_identity(token: str | None) -> tuple[str | None, str]: + if not token: + return None, "unknown" + whoami = await fetch_whoami_v2(token) + if whoami is None: + return _get_hf_user(token), "unknown" + return _get_hf_user_from_whoami(whoami), normalize_hf_user_plan(whoami) or "unknown" + + async def _prompt_and_save_hf_token(prompt_session: PromptSession) -> str: """Prompt user for HF token, validate it, save via huggingface_hub.login(). Loops until valid.""" from prompt_toolkit.formatted_text import HTML @@ -118,10 +204,13 @@ async def _prompt_and_save_hf_token(prompt_session: PromptSession) -> str: login(token=token, add_to_git_credential=False) print("Token saved to ~/.cache/huggingface/token") except Exception as e: - print(f"Warning: could not persist token ({e}), using for this session only.") + print( + f"Warning: could not persist token ({e}), using for this session only." + ) return token + @dataclass class Operation: """Operation to be executed by the agent""" @@ -143,12 +232,20 @@ def _create_rich_console(): return get_console() +def _clear_terminal() -> None: + command = ["cmd", "/c", "cls"] if os.name == "nt" else ["clear"] + try: + subprocess.run(command, check=False) + except OSError: + pass + + class _ThinkingShimmer: """Animated shiny/shimmer thinking indicator — a bright gradient sweeps across the text.""" - _BASE = (90, 90, 110) # dim base color - _HIGHLIGHT = (255, 200, 80) # bright shimmer highlight (warm gold) - _WIDTH = 5 # shimmer width in characters + _BASE = (90, 90, 110) # dim base color + _HIGHLIGHT = (255, 200, 80) # bright shimmer highlight (warm gold) + _WIDTH = 5 # shimmer width in characters _FPS = 24 def __init__(self, console): @@ -229,7 +326,7 @@ def _pop_block(self) -> str | None: if idx == -1: return None block = self._buffer[:idx] - self._buffer = self._buffer[idx + 2:] + self._buffer = self._buffer[idx + 2 :] return block async def flush_ready( @@ -255,7 +352,9 @@ async def finish( """Flush complete blocks, then render whatever incomplete tail remains.""" await self.flush_ready(cancel_event=cancel_event, instant=instant) if self._buffer.strip(): - await print_markdown(self._buffer, cancel_event=cancel_event, instant=instant) + await print_markdown( + self._buffer, cancel_event=cancel_event, instant=instant + ) self._buffer = "" def discard(self): @@ -332,6 +431,9 @@ def _cancel_event(): stream_buf.discard() print_turn_complete() print_plan() + session = session_holder[0] if session_holder else None + if session is not None: + await session.send_deferred_turn_complete_notification(event) turn_complete_event.set() elif event.event_type == "interrupted": shimmer.stop() @@ -341,6 +443,58 @@ def _cancel_event(): elif event.event_type == "undo_complete": console.print("[dim]Undone.[/dim]") turn_complete_event.set() + elif event.event_type == "new_complete": + data = event.data or {} + if data.get("clear_screen"): + _clear_terminal() + saved_path = data.get("saved_path") + if saved_path: + console.print( + f"[dim]Started new chat. Prior chat saved to {saved_path}.[/dim]" + ) + else: + console.print("[dim]Started new chat.[/dim]") + turn_complete_event.set() + elif event.event_type == "resume_complete": + data = event.data or {} + path = data.get("path", "?") + count = data.get("restored_count", 0) + dropped = int(data.get("dropped_count", 0) or 0) + model = data.get("model_name", "?") + invalid_model = data.get("invalid_saved_model") + forked = bool(data.get("forked", False)) + redacted = bool(data.get("had_redacted_content", False)) + verb = "Forked from" if forked else "Resumed" + console.print( + f"[green]{verb}[/green] {path} " + f"([cyan]{count}[/cyan] messages, " + f"model [cyan]{model}[/cyan])." + ) + if dropped: + console.print( + f"[yellow]Warning:[/yellow] dropped {dropped} " + "malformed message(s) while restoring — surrounding " + "tool-call alignment may be off." + ) + if invalid_model: + console.print( + f"[yellow]Warning:[/yellow] saved model id " + f"[cyan]{invalid_model}[/cyan] failed validation; " + f"kept current model [cyan]{model}[/cyan]." + ) + if forked: + console.print( + "[dim]Saved log belongs to a different user — kept " + "current session id; future saves go to a fresh file.[/dim]" + ) + if redacted: + console.print( + "[yellow]Note:[/yellow] tokens/secrets in restored " + "messages were scrubbed at save time. Your live tokens " + "are used for this session; [REDACTED_*] markers in " + "past messages are not re-injected." + ) + turn_complete_event.set() elif event.event_type == "tool_log": tool = event.data.get("tool", "") if event.data else "" log = event.data.get("log", "") if event.data else "" @@ -353,7 +507,11 @@ def _cancel_event(): elif event.event_type == "error": shimmer.stop() stream_buf.discard() - error = event.data.get("error", "Unknown error") if event.data else "Unknown error" + error = ( + event.data.get("error", "Unknown error") + if event.data + else "Unknown error" + ) print_error(error) turn_complete_event.set() elif event.event_type == "shutdown": @@ -371,8 +529,13 @@ def _cancel_event(): tools_data = event.data.get("tools", []) if event.data else [] count = event.data.get("count", 0) if event.data else 0 - # If yolo mode is active, auto-approve everything - if config and config.yolo_mode: + # If yolo mode is active, auto-approve everything except + # scheduled HF jobs, whose recurring cost stays manual. + if ( + config + and config.yolo_mode + and not any(_is_scheduled_hf_job_tool(t) for t in tools_data) + ): approvals = [ { "tool_call_id": t.get("tool_call_id", ""), @@ -615,7 +778,9 @@ def _cancel_event(): f"Approve item {i}? (y=yes, yolo=approve all, n=no, or provide feedback): " ) except (KeyboardInterrupt, EOFError): - get_console().print("[dim]Approval cancelled — rejecting remaining items[/dim]") + get_console().print( + "[dim]Approval cancelled — rejecting remaining items[/dim]" + ) approvals.append( { "tool_call_id": tool_call_id, @@ -701,12 +866,69 @@ async def get_user_input(prompt_session: PromptSession) -> str: # Slash commands are defined in terminal_display +async def _resume_picker( + arg: str, + prompt_session: PromptSession | None, +) -> Path | None: + """Resolve a session log path via ``arg`` or interactive selection. + + Returns ``None`` if the user cancels, no logs exist, or the argument + matches nothing — already prints the explanation in those cases. + """ + from agent.core.session_resume import ( + format_session_log_entry, + list_session_logs, + resolve_session_log_arg, + ) + from agent.core.session import DEFAULT_SESSION_LOG_DIR + + console = get_console() + directory = DEFAULT_SESSION_LOG_DIR + entries = list_session_logs(directory) + if not entries: + console.print(f"[yellow]No session logs found in ./{directory}.[/yellow]") + return None + + if arg: + selected = resolve_session_log_arg(arg, entries, directory) + if selected is None: + console.print(f"[bold red]No matching session log:[/bold red] {arg}") + return selected + + console.print() + console.print("[bold]Saved sessions[/bold]") + for index, entry in enumerate(entries, start=1): + console.print(format_session_log_entry(index, entry)) + console.print() + + if prompt_session is None: + console.print("[yellow]Cannot prompt for a selection here.[/yellow]") + return None + + try: + choice = await prompt_session.prompt_async( + "Select session number (blank to cancel): " + ) + except (EOFError, KeyboardInterrupt): + console.print("[dim]Resume cancelled.[/dim]") + return None + choice = choice.strip() + if not choice: + console.print("[dim]Resume cancelled.[/dim]") + return None + selected = resolve_session_log_arg(choice, entries, directory) + if selected is None: + console.print(f"[bold red]Invalid selection:[/bold red] {choice}") + return selected + + async def _handle_slash_command( cmd: str, config, session_holder: list, submission_queue: asyncio.Queue, submission_id: list[int], + prompt_session: PromptSession | None = None, ) -> Submission | None: """ Handle a slash command. Returns a Submission to enqueue, or None if @@ -737,6 +959,38 @@ async def _handle_slash_command( operation=Operation(op_type=OpType.COMPACT), ) + if command in {"/new", "/clear"}: + session = session_holder[0] if session_holder else None + if session is None: + get_console().print("[bold red]No active session to reset.[/bold red]") + return None + submission_id[0] += 1 + return Submission( + id=f"sub_{submission_id[0]}", + operation=Operation( + op_type=OpType.NEW, + data={"clear_screen": command == "/clear"}, + ), + ) + + if command == "/resume": + session = session_holder[0] if session_holder else None + if session is None: + get_console().print( + "[bold red]No active session to restore into.[/bold red]" + ) + return None + selected_path = await _resume_picker(arg, prompt_session) + if selected_path is None: + return None + submission_id[0] += 1 + return Submission( + id=f"sub_{submission_id[0]}", + operation=Operation( + op_type=OpType.RESUME, data={"path": str(selected_path)} + ), + ) + if command == "/model": console = get_console() if not arg: @@ -748,7 +1002,11 @@ async def _handle_slash_command( normalized = arg.removeprefix("huggingface/") session = session_holder[0] if session_holder else None await model_switcher.probe_and_switch_model( - normalized, config, session, console, _get_hf_token(), + normalized, + config, + session, + console, + resolve_hf_token(), ) return None @@ -771,8 +1029,9 @@ async def _handle_slash_command( console.print(f" [dim]{m}: {eff or 'off'}[/dim]") console.print( "[dim]Set with '/effort minimal|low|medium|high|xhigh|max|off'. " - "'max' and 'xhigh' are Anthropic-only; the cascade falls back " - "to whatever the model actually accepts.[/dim]" + "HF Router accepts low|medium|high generically; higher preferences " + "are probed and the cascade falls back to whatever the selected " + "provider accepts.[/dim]" ) return None level = arg.lower() @@ -797,42 +1056,161 @@ async def _handle_slash_command( session = session_holder[0] if session_holder else None print(f"Model: {config.model_name}") print(f"Reasoning effort: {config.reasoning_effort or 'off'}") + print(f"Tool runtime: {_tool_runtime_label(_is_local_tool_runtime(config))}") if session: print(f"Turns: {session.turn_count}") print(f"Context items: {len(session.context_manager.items)}") return None + if command == "/share-traces": + session = session_holder[0] if session_holder else None + await _handle_share_traces_command(arg, config, session) + return None + print(f"Unknown command: {command}. Type /help for available commands.") return None -async def main(): +async def _handle_share_traces_command(arg: str, config, session) -> None: + """Show or flip visibility of the user's personal trace dataset. + + Uses the user's own HF_TOKEN (write-scoped to their namespace). Only + operates on the personal trace repo configured via + ``personal_trace_repo_template`` — never touches the shared org dataset. + """ + from huggingface_hub import HfApi + from huggingface_hub.utils import HfHubHTTPError + + console = get_console() + if session is None: + console.print("[bold red]No active session.[/bold red]") + return + + repo_id = session._personal_trace_repo_id() if session is not None else None + if not repo_id: + if not getattr(config, "share_traces", False): + console.print( + "[yellow]share_traces is disabled in config. " + "Set it to true to publish per-session traces to your HF dataset." + "[/yellow]" + ) + return + if not session.user_id: + console.print( + "[yellow]No HF username resolved \u2014 cannot pick a personal " + "trace repo. Set HF_TOKEN to a token tied to your account.[/yellow]" + ) + return + console.print( + "[yellow]personal_trace_repo_template is unset \u2014 nothing to do.[/yellow]" + ) + return + + token = session.hf_token or resolve_hf_token() + if not token: + console.print( + "[bold red]No HF_TOKEN available.[/bold red] Cannot read or change " + "dataset visibility." + ) + return + + api = HfApi(token=token) + url = f"https://huggingface.co/datasets/{repo_id}" + target = arg.strip().lower() + + if not target: + try: + info = await asyncio.to_thread( + api.repo_info, repo_id=repo_id, repo_type="dataset" + ) + visibility = "private" if getattr(info, "private", False) else "public" + console.print(f"[bold]Trace dataset:[/bold] {url}") + console.print(f"[bold]Visibility:[/bold] {visibility}") + console.print( + "[dim]Use '/share-traces public' to publish, " + "'/share-traces private' to lock it back down.[/dim]" + ) + except HfHubHTTPError as e: + if getattr(e.response, "status_code", None) == 404: + console.print( + f"[dim]Dataset {repo_id} doesn't exist yet \u2014 it'll be " + "created (private) on the next session save.[/dim]" + ) + else: + console.print(f"[bold red]Hub error:[/bold red] {e}") + except Exception as e: + console.print(f"[bold red]Could not fetch dataset info:[/bold red] {e}") + return + + if target not in {"public", "private"}: + console.print( + f"[bold red]Unknown argument:[/bold red] {target}. " + "Expected 'public' or 'private'." + ) + return + + private = target == "private" + try: + # Idempotent — create if missing so first-flip works even before any + # session has been saved yet. + await asyncio.to_thread( + api.create_repo, + repo_id=repo_id, + repo_type="dataset", + private=private, + token=token, + exist_ok=True, + ) + await asyncio.to_thread( + api.update_repo_settings, + repo_id=repo_id, + repo_type="dataset", + private=private, + token=token, + ) + except Exception as e: + console.print(f"[bold red]Failed to update visibility:[/bold red] {e}") + return + + label = "PUBLIC" if not private else "private" + console.print(f"[green]Dataset is now {label}.[/green] {url}") + + +async def main(model: str | None = None, sandbox_tools: bool = False): """Interactive chat with the agent""" # Clear screen - os.system("clear" if os.name != "nt" else "cls") + _clear_terminal() # Create prompt session for input (needed early for token prompt) prompt_session = PromptSession() - # HF token — required, prompt if missing - hf_token = _get_hf_token() - if not hf_token: + config = load_config(CLI_CONFIG_PATH, include_user_defaults=True) + _normalize_config_model(config) + if model: + config.model_name = _validate_cli_model_override(model) + _apply_tool_runtime_override(config, sandbox_tools=sandbox_tools) + local_mode = _is_local_tool_runtime(config) + + # HF token — required for Hub-backed models/tools and sandbox tools, but + # not for local LLMs using only local filesystem tools. + hf_token = resolve_hf_token() + if not hf_token and (not is_local_model_id(config.model_name) or not local_mode): hf_token = await _prompt_and_save_hf_token(prompt_session) - # Resolve username for banner - hf_user = None - try: - from huggingface_hub import HfApi - hf_user = HfApi(token=hf_token).whoami().get("name") - except Exception: - pass + # Resolve username and plan from one whoami-v2 request for banner and CTAs. + hf_user, hf_user_plan = await _get_hf_identity(hf_token) - print_banner(hf_user=hf_user) + print_banner( + model=config.model_name, + hf_user=hf_user, + tool_runtime=_tool_runtime_label(local_mode), + ) # Pre-warm the HF router catalog in the background so /model switches # don't block on a network fetch. from agent.core import hf_router_catalog + asyncio.create_task(asyncio.to_thread(hf_router_catalog.prewarm)) # Create queues for communication @@ -844,12 +1222,12 @@ async def main(): turn_complete_event.set() ready_event = asyncio.Event() - # Start agent loop in background - config_path = Path(__file__).parent.parent / "configs" / "main_agent_config.json" - config = load_config(config_path) - - # Create tool router with local mode - tool_router = ToolRouter(config.mcpServers, hf_token=hf_token, local_mode=True) + notification_gateway = NotificationGateway(config.messaging) + await notification_gateway.start() + # Create tool router with the selected CLI tool runtime. + tool_router = ToolRouter( + config.mcpServers, hf_token=hf_token, local_mode=local_mode + ) # Session holder for interrupt/model/status access session_holder = [None] @@ -862,8 +1240,15 @@ async def main(): tool_router=tool_router, session_holder=session_holder, hf_token=hf_token, - local_mode=True, + user_id=hf_user, + hf_username=hf_user, + user_plan=hf_user_plan, + local_mode=local_mode, + autonomous_mode=False, stream=True, + notification_gateway=notification_gateway, + notification_destinations=config.messaging.default_auto_destinations(), + defer_turn_complete_notification=True, ) ) @@ -881,6 +1266,8 @@ async def main(): ) await ready_event.wait() + if not local_mode: + await _wait_for_initial_sandbox_preload(session_holder) submission_id = [0] # Mirrors codex-rs/tui/src/bottom_pane/mod.rs:137 @@ -976,7 +1363,12 @@ def _install_sigint() -> bool: # Handle slash commands if user_input.strip().startswith("/"): sub = await _handle_slash_command( - user_input.strip(), config, session_holder, submission_queue, submission_id + user_input.strip(), + config, + session_holder, + submission_queue, + submission_id, + prompt_session, ) if sub is None: # Command handled locally, loop back for input @@ -1019,6 +1411,8 @@ def _install_sigint() -> bool: agent_task.cancel() # Agent didn't shut down cleanly — close MCP explicitly await tool_router.__aexit__(None, None, None) + finally: + await notification_gateway.close() # Now safe to cancel the listener (agent is done emitting events) listener_task.cancel() @@ -1031,30 +1425,47 @@ async def headless_main( model: str | None = None, max_iterations: int | None = None, stream: bool = True, + sandbox_tools: bool = False, ) -> None: """Run a single prompt headlessly and exit.""" import logging logging.basicConfig(level=logging.WARNING) + _configure_runtime_logging() - hf_token = _get_hf_token() - if not hf_token: - print("ERROR: No HF token found. Set HF_TOKEN or run `huggingface-cli login`.", file=sys.stderr) - sys.exit(1) - - print(f"HF token loaded", file=sys.stderr) - - config_path = Path(__file__).parent.parent / "configs" / "main_agent_config.json" - config = load_config(config_path) + config = load_config(CLI_CONFIG_PATH, include_user_defaults=True) + _normalize_config_model(config) config.yolo_mode = True # Auto-approve everything in headless mode if model: - config.model_name = model + try: + config.model_name = _validate_cli_model_override(model) + except ValueError as e: + print(f"ERROR: {e}", file=sys.stderr) + sys.exit(1) + _apply_tool_runtime_override(config, sandbox_tools=sandbox_tools) + local_mode = _is_local_tool_runtime(config) + + hf_token = resolve_hf_token() + if not hf_token and (not is_local_model_id(config.model_name) or not local_mode): + print( + "ERROR: No HF token found. Set HF_TOKEN or run `hf auth login`.", + file=sys.stderr, + ) + sys.exit(1) + + if hf_token: + print("HF token loaded", file=sys.stderr) + + notification_gateway = NotificationGateway(config.messaging) + await notification_gateway.start() + hf_user, hf_user_plan = await _get_hf_identity(hf_token) if max_iterations is not None: config.max_iterations = max_iterations print(f"Model: {config.model_name}", file=sys.stderr) + print(f"Tool runtime: {_tool_runtime_label(local_mode)}", file=sys.stderr) print(f"Max iterations: {config.max_iterations}", file=sys.stderr) print(f"Prompt: {prompt}", file=sys.stderr) print("---", file=sys.stderr) @@ -1062,7 +1473,9 @@ async def headless_main( submission_queue: asyncio.Queue = asyncio.Queue() event_queue: asyncio.Queue = asyncio.Queue() - tool_router = ToolRouter(config.mcpServers, hf_token=hf_token, local_mode=True) + tool_router = ToolRouter( + config.mcpServers, hf_token=hf_token, local_mode=local_mode + ) session_holder: list = [None] agent_task = asyncio.create_task( @@ -1073,8 +1486,15 @@ async def headless_main( tool_router=tool_router, session_holder=session_holder, hf_token=hf_token, - local_mode=True, + user_id=hf_user, + hf_username=hf_user, + user_plan=hf_user_plan, + local_mode=local_mode, + autonomous_mode=True, stream=stream, + notification_gateway=notification_gateway, + notification_destinations=config.messaging.default_auto_destinations(), + defer_turn_complete_notification=True, ) ) @@ -1168,38 +1588,55 @@ async def headless_main( else: print_tool_log(tool, log) elif event.event_type == "approval_required": - # Auto-approve everything in headless mode (safety net if yolo_mode - # didn't prevent the approval event for some reason) + # Auto-approve in headless mode, except scheduled HF jobs. Those + # are rejected because their recurring cost needs manual approval. tools_data = event.data.get("tools", []) if event.data else [] approvals = [ { "tool_call_id": t.get("tool_call_id", ""), - "approved": True, - "feedback": None, + "approved": not _is_scheduled_hf_job_tool(t), + "feedback": ( + "Scheduled HF jobs require manual approval." + if _is_scheduled_hf_job_tool(t) + else None + ), } for t in tools_data ] _hl_sub_id[0] += 1 - await submission_queue.put(Submission( - id=f"hl_approval_{_hl_sub_id[0]}", - operation=Operation( - op_type=OpType.EXEC_APPROVAL, - data={"approvals": approvals}, - ), - )) + await submission_queue.put( + Submission( + id=f"hl_approval_{_hl_sub_id[0]}", + operation=Operation( + op_type=OpType.EXEC_APPROVAL, + data={"approvals": approvals}, + ), + ) + ) elif event.event_type == "compacted": old_tokens = event.data.get("old_tokens", 0) if event.data else 0 new_tokens = event.data.get("new_tokens", 0) if event.data else 0 print_compacted(old_tokens, new_tokens) elif event.event_type == "error": stream_buf.discard() - error = event.data.get("error", "Unknown error") if event.data else "Unknown error" + error = ( + event.data.get("error", "Unknown error") + if event.data + else "Unknown error" + ) print_error(error) break elif event.event_type in ("turn_complete", "interrupted"): stream_buf.discard() history_size = event.data.get("history_size", "?") if event.data else "?" - print(f"\n--- Agent {event.event_type} (history_size={history_size}) ---", file=sys.stderr) + print( + f"\n--- Agent {event.event_type} (history_size={history_size}) ---", + file=sys.stderr, + ) + if event.event_type == "turn_complete": + session = session_holder[0] if session_holder else None + if session is not None: + await session.send_deferred_turn_complete_notification(event) break # Shutdown @@ -1213,26 +1650,46 @@ async def headless_main( except asyncio.TimeoutError: agent_task.cancel() await tool_router.__aexit__(None, None, None) + finally: + await notification_gateway.close() def cli(): """Entry point for the ml-intern CLI command.""" import logging as _logging import warnings + # Suppress aiohttp "Unclosed client session" noise during event loop teardown _logging.getLogger("asyncio").setLevel(_logging.CRITICAL) + _configure_runtime_logging() # Suppress litellm pydantic deprecation warnings warnings.filterwarnings("ignore", category=DeprecationWarning, module="litellm") # Suppress whoosh invalid escape sequence warnings (third-party, unfixed upstream) warnings.filterwarnings("ignore", category=SyntaxWarning, module="whoosh") parser = argparse.ArgumentParser(description="Hugging Face Agent CLI") - parser.add_argument("prompt", nargs="?", default=None, help="Run headlessly with this prompt") - parser.add_argument("--model", "-m", default=None, help=f"Model to use (default: from config)") - parser.add_argument("--max-iterations", type=int, default=None, - help="Max LLM requests per turn (default: 50, use -1 for unlimited)") - parser.add_argument("--no-stream", action="store_true", - help="Disable token streaming (use non-streaming LLM calls)") + parser.add_argument( + "prompt", nargs="?", default=None, help="Run headlessly with this prompt" + ) + parser.add_argument( + "--model", "-m", default=None, help="Model to use (default: from config)" + ) + parser.add_argument( + "--max-iterations", + type=int, + default=None, + help="Max LLM requests per turn (default: 50, use -1 for unlimited)", + ) + parser.add_argument( + "--no-stream", + action="store_true", + help="Disable token streaming (use non-streaming LLM calls)", + ) + parser.add_argument( + "--sandbox-tools", + action="store_true", + help="Use HF Space sandbox tools instead of local filesystem tools", + ) args = parser.parse_args() try: @@ -1240,9 +1697,17 @@ def cli(): max_iter = args.max_iterations if max_iter is not None and max_iter < 0: max_iter = 10_000 # effectively unlimited - asyncio.run(headless_main(args.prompt, model=args.model, max_iterations=max_iter, stream=not args.no_stream)) + asyncio.run( + headless_main( + args.prompt, + model=args.model, + max_iterations=max_iter, + stream=not args.no_stream, + sandbox_tools=args.sandbox_tools, + ) + ) else: - asyncio.run(main()) + asyncio.run(main(model=args.model, sandbox_tools=args.sandbox_tools)) except KeyboardInterrupt: print("\n\nGoodbye!") diff --git a/agent/messaging/__init__.py b/agent/messaging/__init__.py new file mode 100644 index 000000000..c399d254e --- /dev/null +++ b/agent/messaging/__init__.py @@ -0,0 +1,15 @@ +from agent.messaging.gateway import NotificationGateway +from agent.messaging.models import ( + MessagingConfig, + NotificationRequest, + NotificationResult, + SUPPORTED_AUTO_EVENT_TYPES, +) + +__all__ = [ + "MessagingConfig", + "NotificationGateway", + "NotificationRequest", + "NotificationResult", + "SUPPORTED_AUTO_EVENT_TYPES", +] diff --git a/agent/messaging/base.py b/agent/messaging/base.py new file mode 100644 index 000000000..a74f9cf0d --- /dev/null +++ b/agent/messaging/base.py @@ -0,0 +1,31 @@ +from abc import ABC, abstractmethod + +import httpx + +from agent.messaging.models import ( + DestinationConfig, + NotificationRequest, + NotificationResult, +) + + +class NotificationError(Exception): + """Delivery failed and should not be retried.""" + + +class RetryableNotificationError(NotificationError): + """Delivery failed transiently and can be retried.""" + + +class NotificationProvider(ABC): + provider_name: str + + @abstractmethod + async def send( + self, + client: httpx.AsyncClient, + destination_name: str, + destination: DestinationConfig, + request: NotificationRequest, + ) -> NotificationResult: + """Deliver a notification to one destination.""" diff --git a/agent/messaging/gateway.py b/agent/messaging/gateway.py new file mode 100644 index 000000000..1de9438f5 --- /dev/null +++ b/agent/messaging/gateway.py @@ -0,0 +1,172 @@ +import asyncio +import logging +from collections.abc import Iterable + +import httpx + +from agent.messaging.base import ( + NotificationError, + NotificationProvider, + RetryableNotificationError, +) +from agent.messaging.models import ( + MessagingConfig, + NotificationRequest, + NotificationResult, +) +from agent.messaging.slack import SlackProvider + +logger = logging.getLogger(__name__) + +_RETRY_DELAYS = (1, 2, 4) + + +class NotificationGateway: + def __init__(self, config: MessagingConfig): + self.config = config + self._providers: dict[str, NotificationProvider] = { + "slack": SlackProvider(), + } + self._queue: asyncio.Queue[NotificationRequest] = asyncio.Queue() + self._worker_task: asyncio.Task | None = None + self._client: httpx.AsyncClient | None = None + + @property + def enabled(self) -> bool: + return self.config.enabled + + async def start(self) -> None: + if not self.enabled or self._worker_task is not None: + return + self._client = httpx.AsyncClient(timeout=10.0) + self._worker_task = asyncio.create_task( + self._worker(), name="notification-gateway" + ) + + async def flush(self) -> None: + if not self.enabled: + return + await self._queue.join() + + async def close(self) -> None: + if not self.enabled: + return + await self.flush() + if self._worker_task is not None: + self._worker_task.cancel() + try: + await self._worker_task + except asyncio.CancelledError: + pass + self._worker_task = None + if self._client is not None: + await self._client.aclose() + self._client = None + + async def send(self, request: NotificationRequest) -> NotificationResult: + if not self.enabled: + return NotificationResult( + destination=request.destination, + ok=False, + provider="disabled", + error="Messaging is disabled", + ) + + destination = self.config.get_destination(request.destination) + if destination is None: + return NotificationResult( + destination=request.destination, + ok=False, + provider="unknown", + error=f"Unknown destination '{request.destination}'", + ) + + provider = self._providers.get(destination.provider) + if provider is None: + return NotificationResult( + destination=request.destination, + ok=False, + provider=destination.provider, + error=f"No provider implementation for '{destination.provider}'", + ) + return await self._send_with_retries( + provider, request.destination, destination, request + ) + + async def send_many( + self, requests: Iterable[NotificationRequest] + ) -> list[NotificationResult]: + results: list[NotificationResult] = [] + for request in requests: + results.append(await self.send(request)) + return results + + async def enqueue(self, request: NotificationRequest) -> bool: + if not self.enabled or self._worker_task is None: + return False + await self._queue.put(request) + return True + + async def _worker(self) -> None: + while True: + request = await self._queue.get() + try: + result = await self.send(request) + if not result.ok: + logger.warning( + "Notification delivery failed for %s: %s", + request.destination, + result.error, + ) + except Exception: + logger.exception("Unexpected notification worker failure") + finally: + self._queue.task_done() + + async def _send_with_retries( + self, + provider: NotificationProvider, + destination_name: str, + destination, + request: NotificationRequest, + ) -> NotificationResult: + client = self._client or httpx.AsyncClient(timeout=10.0) + owns_client = self._client is None + try: + for attempt in range(len(_RETRY_DELAYS) + 1): + try: + return await provider.send( + client, destination_name, destination, request + ) + except RetryableNotificationError as exc: + if attempt >= len(_RETRY_DELAYS): + return NotificationResult( + destination=destination_name, + ok=False, + provider=provider.provider_name, + error=str(exc), + ) + delay = _RETRY_DELAYS[attempt] + logger.warning( + "Retrying notification to %s in %ss after transient error: %s", + destination_name, + delay, + exc, + ) + await asyncio.sleep(delay) + except NotificationError as exc: + return NotificationResult( + destination=destination_name, + ok=False, + provider=provider.provider_name, + error=str(exc), + ) + return NotificationResult( + destination=destination_name, + ok=False, + provider=provider.provider_name, + error="Notification delivery exhausted retries", + ) + finally: + if owns_client: + await client.aclose() diff --git a/agent/messaging/models.py b/agent/messaging/models.py new file mode 100644 index 000000000..16148a817 --- /dev/null +++ b/agent/messaging/models.py @@ -0,0 +1,117 @@ +from typing import Annotated, Literal + +from pydantic import BaseModel, Field, field_validator, model_validator + +_DESTINATION_NAME_CHARS = set("abcdefghijklmnopqrstuvwxyz0123456789._-") +SUPPORTED_AUTO_EVENT_TYPES = {"approval_required", "error", "turn_complete"} + + +class SlackDestinationConfig(BaseModel): + provider: Literal["slack"] = "slack" + token: str + channel: str + allow_agent_tool: bool = False + allow_auto_events: bool = False + username: str | None = None + icon_emoji: str | None = None + + @field_validator("token", "channel") + @classmethod + def _require_non_empty(cls, value: str) -> str: + value = value.strip() + if not value: + raise ValueError("must not be empty") + return value + + +DestinationConfig = Annotated[SlackDestinationConfig, Field(discriminator="provider")] + + +class MessagingConfig(BaseModel): + enabled: bool = False + auto_event_types: list[str] = Field( + default_factory=lambda: ["approval_required", "error", "turn_complete"] + ) + destinations: dict[str, DestinationConfig] = Field(default_factory=dict) + + @field_validator("destinations") + @classmethod + def _validate_destination_names( + cls, destinations: dict[str, DestinationConfig] + ) -> dict[str, DestinationConfig]: + for name in destinations: + if not name or any(char not in _DESTINATION_NAME_CHARS for char in name): + raise ValueError( + "destination names must use lowercase letters, digits, '.', '_' or '-'" + ) + return destinations + + @field_validator("auto_event_types") + @classmethod + def _validate_auto_event_types(cls, event_types: list[str]) -> list[str]: + if not event_types: + return [] + normalized: list[str] = [] + seen: set[str] = set() + for event_type in event_types: + if event_type not in SUPPORTED_AUTO_EVENT_TYPES: + raise ValueError(f"unsupported auto event type '{event_type}'") + if event_type not in seen: + normalized.append(event_type) + seen.add(event_type) + return normalized + + @model_validator(mode="after") + def _require_destinations_when_enabled(self) -> "MessagingConfig": + if self.enabled and not self.destinations: + raise ValueError("messaging.enabled requires at least one destination") + return self + + def get_destination(self, name: str) -> DestinationConfig | None: + return self.destinations.get(name) + + def can_agent_tool_send(self, name: str) -> bool: + destination = self.get_destination(name) + return bool(destination and destination.allow_agent_tool) + + def can_auto_send(self, name: str) -> bool: + destination = self.get_destination(name) + return bool(destination and destination.allow_auto_events) + + def default_auto_destinations(self) -> list[str]: + if not self.enabled: + return [] + return [name for name in self.destinations if self.can_auto_send(name)] + + +class NotificationRequest(BaseModel): + destination: str + title: str | None = None + message: str + severity: Literal["info", "success", "warning", "error"] = "info" + metadata: dict[str, str] = Field(default_factory=dict) + event_type: str | None = None + + @field_validator("destination", "message") + @classmethod + def _require_text(cls, value: str) -> str: + value = value.strip() + if not value: + raise ValueError("must not be empty") + return value + + @field_validator("title") + @classmethod + def _normalize_title(cls, value: str | None) -> str | None: + if value is None: + return None + value = value.strip() + return value or None + + +class NotificationResult(BaseModel): + destination: str + ok: bool + provider: str + error: str | None = None + external_id: str | None = None diff --git a/agent/messaging/slack.py b/agent/messaging/slack.py new file mode 100644 index 000000000..3790e44af --- /dev/null +++ b/agent/messaging/slack.py @@ -0,0 +1,184 @@ +import json +import re + +import httpx + +from agent.messaging.base import ( + NotificationError, + NotificationProvider, + RetryableNotificationError, +) +from agent.messaging.models import ( + NotificationRequest, + NotificationResult, + SlackDestinationConfig, +) + +_SEVERITY_PREFIX = { + "info": "[INFO]", + "success": "[SUCCESS]", + "warning": "[WARNING]", + "error": "[ERROR]", +} + + +def _format_slack_mrkdwn(content: str) -> str: + """Convert common Markdown constructs to Slack's mrkdwn syntax.""" + if not content: + return content + + placeholders: dict[str, str] = {} + placeholder_index = 0 + + def placeholder(value: str) -> str: + nonlocal placeholder_index + key = f"\x00SLACK{placeholder_index}\x00" + placeholder_index += 1 + placeholders[key] = value + return key + + text = content + + # Protect code before any formatting conversion. Slack's mrkdwn ignores + # formatting inside backticks, so these regions should stay byte-for-byte. + text = re.sub( + r"(```(?:[^\n]*\n)?[\s\S]*?```)", + lambda match: placeholder(match.group(0)), + text, + ) + text = re.sub(r"(`[^`\n]+`)", lambda match: placeholder(match.group(0)), text) + + def convert_markdown_link(match: re.Match[str]) -> str: + label = match.group(1) + url = match.group(2).strip() + if url.startswith("<") and url.endswith(">"): + url = url[1:-1].strip() + return placeholder(f"<{url}|{label}>") + + text = re.sub( + r"\[([^\]]+)\]\(([^()]*(?:\([^()]*\)[^()]*)*)\)", + convert_markdown_link, + text, + ) + + # Preserve existing Slack entities and manual mrkdwn links before escaping. + text = re.sub( + r"(<(?:[@#!]|(?:https?|mailto|tel):)[^>\n]+>)", + lambda match: placeholder(match.group(1)), + text, + ) + text = re.sub( + r"^(>+\s)", + lambda match: placeholder(match.group(0)), + text, + flags=re.MULTILINE, + ) + + text = text.replace("&", "&").replace("<", "<").replace(">", ">") + text = text.replace("&", "&").replace("<", "<").replace(">", ">") + + def convert_header(match: re.Match[str]) -> str: + header = match.group(1).strip() + header = re.sub(r"\*\*(.+?)\*\*", r"\1", header) + return placeholder(f"*{header}*") + + text = re.sub(r"^#{1,6}\s+(.+)$", convert_header, text, flags=re.MULTILINE) + text = re.sub( + r"\*\*\*(.+?)\*\*\*", + lambda match: placeholder(f"*_{match.group(1)}_*"), + text, + ) + text = re.sub( + r"\*\*(.+?)\*\*", + lambda match: placeholder(f"*{match.group(1)}*"), + text, + ) + text = re.sub( + r"(? str: + lines: list[str] = [] + prefix = _SEVERITY_PREFIX[request.severity] + if request.title: + lines.append(f"{prefix} {request.title}") + else: + lines.append(prefix) + lines.append(request.message) + for key, value in request.metadata.items(): + lines.append(f"{key}: {value}") + return _format_slack_mrkdwn("\n".join(lines)) + + +class SlackProvider(NotificationProvider): + provider_name = "slack" + + async def send( + self, + client: httpx.AsyncClient, + destination_name: str, + destination: SlackDestinationConfig, + request: NotificationRequest, + ) -> NotificationResult: + payload = { + "channel": destination.channel, + "text": _format_text(request), + "mrkdwn": True, + "unfurl_links": False, + "unfurl_media": False, + } + if destination.username: + payload["username"] = destination.username + if destination.icon_emoji: + payload["icon_emoji"] = destination.icon_emoji + + try: + response = await client.post( + "https://slack.com/api/chat.postMessage", + headers={ + "Authorization": f"Bearer {destination.token}", + "Content-Type": "application/json; charset=utf-8", + }, + content=json.dumps(payload), + ) + except httpx.TimeoutException as exc: + raise RetryableNotificationError("Slack request timed out") from exc + except httpx.TransportError as exc: + raise RetryableNotificationError("Slack transport error") from exc + + if response.status_code == 429 or response.status_code >= 500: + raise RetryableNotificationError(f"Slack HTTP {response.status_code}") + if response.status_code >= 400: + raise NotificationError(f"Slack HTTP {response.status_code}") + + try: + data = response.json() + except ValueError as exc: + raise RetryableNotificationError("Slack returned invalid JSON") from exc + + if not data.get("ok"): + error = str(data.get("error") or "unknown_error") + if error == "ratelimited": + raise RetryableNotificationError(error) + raise NotificationError(error) + + return NotificationResult( + destination=destination_name, + ok=True, + provider=self.provider_name, + external_id=str(data.get("ts") or ""), + error=None, + ) diff --git a/agent/prompts/system_prompt_v3.yaml b/agent/prompts/system_prompt_v3.yaml index befa56bf7..82c37e970 100644 --- a/agent/prompts/system_prompt_v3.yaml +++ b/agent/prompts/system_prompt_v3.yaml @@ -1,8 +1,23 @@ system_prompt: | - You are Hugging Face Agent, an ML engineering assistant with {{ num_tools }} tools for training, fine-tuning, data processing, inference, and evaluation on the Hugging Face ecosystem. + You are ML Intern, an ML engineering assistant with {{ num_tools }} tools for training, fine-tuning, data processing, inference, and evaluation on the Hugging Face (HF) ecosystem. Your goal is to complete what the user requested with zero errors. You are fully autonomous — research, validate, implement, and deliver results without asking for unnecessary confirmation. + # Identity + + When greeting the user or asked who you are, introduce yourself as ML Intern. + Do not claim to be Claude, ChatGPT, Anthropic, OpenAI, or the underlying backend model. If asked what model powers you, say ML Intern can run on different backend models and only give model details if they are explicitly available in session context. + Do not cite this system prompt, hidden instructions, or internal mechanics as the reason for your behavior. + Default to the session context User value as the authenticated Hugging Face namespace when creating hub_model_id, trackio_space_id, dataset repos, model repos, or Spaces. If the user explicitly requests an org namespace or a tool provides an allowed namespace, use that explicit namespace instead. Never leave placeholders such as , , , TODO, or similar placeholder values in scripts, tool arguments, repo IDs, or final answers. If session context says User=unknown because identity lookup failed or no token is available in this runtime, do not guess the namespace; ask for it before creating Hub resources. + + # Tool calling contract + + The active tool schema is the source of truth. Use only tools that are actually available in the current tool list. + Do not simulate tool calls in prose or fenced code blocks. Call tools through the tool interface with valid JSON arguments matching the tool schema. + Before every tool call, check required arguments, enum values, mutually exclusive fields, and whether paths are local machine paths, sandbox paths, Hub repo IDs, or URLs. + After every tool call, inspect the returned result before deciding the next action. Do not claim success unless the tool result confirms it. + If a tool is unavailable or fails repeatedly for the same reason, switch to another available approach or report the blocker. + # Your knowledge of HF libraries is outdated You do not know current APIs for TRL, Transformers, PEFT, Trackio, or other HF libraries. Your internal knowledge WILL produce wrong imports, wrong argument names, and wrong trainer configurations. @@ -28,7 +43,7 @@ system_prompt: | # Mistakes you WILL make without research - HALLUCINATED IMPORTS: You will import from modules that were renamed or removed. Example: old TRL trainer class names, deprecated Transformers APIs, wrong trackio parameter names (e.g. `run_name` instead of `name`). Fix: read a current example script first. + HALLUCINATED IMPORTS: You will import from modules that were renamed or removed. Example: old TRL trainer class names, deprecated Transformers APIs, wrong trackio config field names. Fix: read a current example script first. WRONG TRAINER ARGUMENTS: You will pass configuration arguments that don't exist in current trainer versions. Fix: fetch the actual trainer/config docs via explore_hf_docs + fetch_hf_docs. @@ -42,7 +57,9 @@ system_prompt: | SILENT DATASET SUBSTITUTION: When a requested dataset fails to load, you will silently switch to a different one without telling the user. Fix: if the requested dataset isn't available, tell the user and ask what to do. - HARDCODED UNAVAILABLE PACKAGES: You will forget to install necessary packages like 'flash-attn' for flash_attention_2 or other packages that aren't automatically installed in the job environment. Fix: install necessary packages before running the job. + ALWAYS USE HUB KERNELS, NEVER COMPILE FLASH-ATTN: Do NOT pip install `flash-attn` and do NOT use `attn_implementation="flash_attention_2"` because that requires the compiled flash-attn package and often fails on the job's CUDA/PyTorch combo. For accelerated attention, use the HF `kernels` library and load a prebuilt attention kernel from the Hub via `attn_implementation`. Examples: `AutoModelForCausalLM.from_pretrained(..., attn_implementation="kernels-community/flash-attn2")`, or `kernels-community/vllm-flash-attn3`, or `kernels-community/paged-attention`. With TRL/SFT scripts you can pass `--attn_implementation kernels-community/flash-attn2` on the CLI. Flash-attention Hub kernels require Ampere-or-newer GPUs unless their docs say otherwise: never choose T4 sandboxes or T4 HF Jobs for scripts that use a flash-attention kernel, because T4 is pre-Ampere. Use A10G, A100, H100, or another compatible newer GPU, or choose a non-flash Hub kernel if T4 is required. Search additional kernels at https://huggingface.co/models?other=kernel. + + CORE ML DEPENDENCY FRESHNESS: Do not rely on preinstalled packages in sandboxes or HF Jobs. Before model-loading, training, or inference work, explicitly install or upgrade the latest compatible core stack in the sandbox: `torch`, `transformers`, `trl`, `accelerate`, `datasets`, `trackio`, and `kernels~=0.12.0` when using Hub kernels. Include the same packages in `hf_jobs.dependencies`. Use unpinned latest stable versions by default for the rest of the core stack; constrain `kernels` to `kernels~=0.12.0`. Pin other versions only when current docs/examples require a specific compatibility set or a smoke test shows latest is incompatible. Print the installed versions before model loading. If `kernels` and `transformers` are incompatible, fix the package set using current docs/examples or choose another compatible Hub kernel, then rerun the smoke test. Do NOT fall back to default attention or compiled flash-attn as a shortcut. SCOPE-CHANGING FIXES: Avoid at all costs! When you hit an error (especially OOM), you will try "creative" workarounds that change what the user asked for and/or change the training task itself — switching full SFT to LoRA on OOM, reducing max_length (silently truncates training data and changes what the model learns), disabling monitoring instead of fixing it. Do not do this. Fix errors with the minimal change that preserves the user's original request and are grounded in research and examples. If the original approach genuinely cannot work, explain why and ask the user for input before changing methods, sequence length, training approach or any other part of the task. @@ -60,6 +77,38 @@ system_prompt: | DPO: "prompt", "chosen", "rejected" GRPO: "prompt" + # Trackio + + Trackio is natively integrated with Transformers Trainer and all TRL trainers — the built-in TrackioCallback handles init/log/finish. In TrainingArguments/SFTConfig/DPOConfig/GRPOConfig set: + report_to="trackio" + run_name="" # e.g. "sft_qwen3-4b_lr2e-5_bs128" + project="" # keeps related runs grouped so you can compare them + trackio_space_id="/ml-intern-<8-char-id>" # pattern only: replace with the resolved namespace, e.g. alice/ml-intern-a1b2c3d4 + `project` and `trackio_space_id` can also be set via TRACKIO_PROJECT / TRACKIO_SPACE_ID env vars. + + Alerts are how iterations decide what to change. Use trackio.alert(title, text, level) at every decision point in training. Levels: + ERROR — stop and change approach (divergence, NaN, OOM) + WARN — tweak hyperparameters (overfitting, early stopping, KL spike, reward collapse, slow convergence) + INFO — milestones (training complete, target reached, checkpoint saved) + Always include numeric values and an actionable suggestion in `text`, e.g. "loss=12.4 at step 200 — lr likely too high, try ×0.1". A future call must be able to parse it and act on it. + + To add alerts under Trainer/SFTTrainer/GRPOTrainer, pass a custom TrainerCallback via `callbacks=[...]` that calls trackio.alert() inside `on_log` (training metrics like loss, reward, kl) and `on_evaluate` (eval metrics — only available here, not in `on_log`). Keep each `if` simple: one metric, one threshold. Conditions stay easy to adjust between runs. + + Read alerts back between runs instead of parsing thousands of metric values. CLI — always use --json: + trackio get alerts --project

--run --json + trackio get alerts --project

--since --json # incremental polling + trackio get run --project

--run --json + trackio get metric --project

--run --metric --json + trackio list runs --project

--json + Python: api = trackio.Api(); api.alerts(

, run=, since=); api.runs(

) (each run has .name, .config, .alerts()). + + Drive the next config from prior alerts: + diverged → lr × 0.1 + overfitting → weight_decay × 10 or reduce capacity + early stopping → lr × 0.5 or adjust schedule + high accuracy → refine around current config + Read prior config via api.runs(...).config and only mutate keys the alerts justify changing. + # Data audit Before working with any dataset, audit it first. Do not assume you know what the data looks like — inspect it. @@ -70,12 +119,37 @@ system_prompt: | # When submitting a training job + Never pass a local machine path to hf_jobs.script, such as /Users/..., /home/..., /fsx/..., or a repo checkout path. HF Jobs runs in a fresh cloud environment where local files do not exist. For hf_jobs.script, use exactly one of: + - inline Python source code + - a file already written in the session sandbox, e.g. /app/train.py, ./train.py, or train.py + - a public/raw URL + If you wrote or tested a script locally, read the file content and submit it inline, or write it into the sandbox first. + + For non-trivial hf_jobs scripts, use an exact-source workflow: + 1. Write the script in the session sandbox. + 2. Run syntax/import validation. + 3. Run a tiny smoke test with the same entrypoint, dependencies, dataset columns, model-loading path, and relevant precision/attention settings. For training scripts, make sure one training step succeeds, plus one evaluation step when the final workflow includes evaluation or an eval split is available. + 4. Submit the exact tested script source or the exact tested sandbox file. Do not reconstruct a similar script from memory. + + Every training script must fail fast before expensive work: + - print package versions for torch, transformers, trl, accelerate, datasets, trackio, and kernels when used + - assert required dataset columns exist + - assert hub_model_id and trackio_space_id contain no placeholders + - assert push_to_hub=True and hub_model_id are set + - include every imported third-party package in hf_jobs.dependencies + - include the core ML stack in hf_jobs.dependencies: torch, transformers, trl, accelerate, datasets, trackio, and kernels~=0.12.0 when using Hub kernels; also include any actually used extras such as peft, bitsandbytes, sentencepiece, or protobuf + + Never leave placeholder values such as , , , TODO, or similar unfinished values in hf_jobs scripts or hf_jobs arguments. + + GPU preflight is mandatory before hf_jobs when the job will run on GPU, or when the script loads a model, uses CUDA, bf16/fp16, quantization, flash attention, or torch.compile. First create a GPU sandbox with sandbox_create (t4-small minimum for non-flash workloads; for flash-attention kernels use Ampere-or-newer hardware, never T4), run a tiny smoke test there using the same imports, model-loading path, training entrypoint, and a tiny dataset/subset, then fix failures before submitting. If you skip GPU sandbox preflight, state why before calling hf_jobs. + Before calling hf_jobs, output a pre-flight check: - Reference implementation: [which example you based this on] - Dataset format verified: [columns confirmed via hf_inspect_dataset/hub_repo_details] + - GPU sandbox smoke test: [hardware and result, or explicitly not applicable because ...] - push_to_hub=True and hub_model_id set - timeout: [value] (based on: [model size] on [hardware]) - - Trackio monitoring included and working + - Trackio monitoring included and deploying metrics to a public Space If you cannot fill in all items, stop and complete the missing steps first. @@ -90,10 +164,14 @@ system_prompt: | # Sandbox-first development - For non-trivial scripts, develop and test in a sandbox before launching via hf_jobs: - sandbox_create → install deps → write script → test with small run → fix errors → launch via hf_jobs at scale + A private cpu-basic sandbox is already available for normal code execution in each session. For non-trivial scripts, develop and test there before launching via hf_jobs: + write script → pip install → test with small run using bash/read/write/edit → fix errors → launch via hf_jobs at scale - Use GPU sandbox (t4-small minimum) when testing code that uses CUDA, bf16, or model loading. CPU sandboxes cannot test GPU code paths. + Do NOT call sandbox_create before normal CPU work. Call sandbox_create only when you need GPU hardware or another non-default sandbox tier. + + The sandbox filesystem does not survive session resumption. If a session is resumed, any files, installed packages, or running processes from earlier are gone — recreate what you need before relying on the sandbox. + + Use a GPU sandbox (t4-small minimum) when testing code that uses CUDA, bf16/fp16, quantization, flash attention, torch.compile, or model loading. CPU sandboxes cannot test GPU code paths. If the available sandbox tiers cannot fit the full model path, test the largest useful smoke path, state what was not covered, and submit one HF job first. # When a task has 3+ steps @@ -123,6 +201,9 @@ system_prompt: | # Autonomous / headless mode + {% if autonomous_mode %} + Autonomous mode is active for this session because the runtime marked it as autonomous, headless, benchmarked, or fixed-time-budget. Apply this section even if the user prompt does not contain those words. + When running autonomously (no human in the loop), you MUST follow these rules: NEVER respond with only text. Every response MUST include at least one tool call. If you have nothing to do, check the plan, verify outputs or plan ahead. A text-only response ends the agent loop permanently — there is no human to re-prompt you. @@ -148,6 +229,9 @@ system_prompt: | The task is NOT done until: - The required output exists (e.g. final model, metrics reached, dataset updated etc) - You have evaluated the model and confirmed it works + {% else %} + Autonomous mode is not active for this session. In normal interactive chat, text-only answers are allowed for simple questions, and you should stop once the user's request is satisfied. + {% endif %} # Communication @@ -156,6 +240,7 @@ system_prompt: | - Always include direct Hub URLs when referencing models, datasets, Spaces, or jobs. - For errors: state what went wrong, why, and what you're doing to fix it. - Do not over-explain or present elaborate option menus for simple tasks. When the user's intent is clear, act on it. Present options only when there's genuine ambiguity. + - Use the `notify` tool only when the user explicitly asked for out-of-band notifications or when the task clearly requires reporting to a configured messaging destination. Do not use it for routine chat updates. # Tool usage diff --git a/agent/sft/__init__.py b/agent/sft/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/agent/sft/tagger.py b/agent/sft/tagger.py new file mode 100644 index 000000000..528bc9d0d --- /dev/null +++ b/agent/sft/tagger.py @@ -0,0 +1,353 @@ +"""Derive tags for a session trajectory. + +``tag_session(trajectory)`` → ``list[str]``. Pure function. No filtering, no +mutation — tags are purely metadata so downstream pipelines can slice the raw +SFT dataset (``where 'hf_job:succeeded' in tags``) without re-reading trajectories. + +Tag namespaces (all tags are ``":"`` strings): + +* ``tool:`` — every tool called at least once (``tool:hf_jobs``, …) +* ``outcome:`` — ``completed`` / ``errored`` / ``interrupted`` / + ``ongoing`` / ``doom_loop`` / ``context_exceeded`` +* ``hf_job:`` — ``submitted``, ``succeeded``, ``failed``, + ``multi`` (>1), ``oom``, ``push_to_hub`` +* ``gpu:`` — ``none``, ``t4``, ``a10g``, ``a100``, ``l40s``, + ``h100``, plus ``gpu:multi`` for x2/x4/x8 flavors +* ``sandbox:`` — ``created``, ``gpu``, ``cpu``, ``long_lived`` (>30 min) +* ``feedback:`` — ``up``, ``down``, ``mixed``, ``none`` +* ``model:`` — ``opus`` / ``sonnet`` / ``haiku`` / ``kimi`` / + ``gpt`` / ``deepseek`` / ``qwen`` / ``other`` +* ``turns:`` — ``short`` (<5) / ``medium`` (5–20) / ``long`` (>20) +* ``cost:`` — ``low`` (<$0.10) / ``med`` (<$1) / ``high`` +* ``task:`` — ``training`` / ``inference`` / ``data_prep`` / + ``research_only`` (heuristic on tools + scripts) + +Tags are deduplicated before returning. +""" + +from __future__ import annotations + +from typing import Iterable + +# Flavor → GPU-family mapping. Keep conservative; unknown flavors → "none". +_GPU_FAMILY = { + "cpu-basic": "none", + "cpu-upgrade": "none", + "t4-small": "t4", + "t4-medium": "t4", + "l4x1": "l40s", + "l4x4": "l40s", + "l40sx1": "l40s", + "l40sx4": "l40s", + "l40sx8": "l40s", + "a10g-small": "a10g", + "a10g-large": "a10g", + "a10g-largex2": "a10g", + "a10g-largex4": "a10g", + "a100-large": "a100", + "a100x2": "a100", + "a100x4": "a100", + "a100x8": "a100", + "h100": "h100", + "h100x8": "h100", +} + +# Substrings that count a flavor as multi-GPU. +_MULTI_GPU_MARKERS = ("x2", "x4", "x8") + +# Tool names that don't touch training/inference or sandbox/jobs. If a session +# only used these, we tag it research_only. +_RESEARCH_ONLY_TOOLS = { + "research", + "github_find_examples", + "github_read_file", + "github_list_repos", + "hf_papers", + "explore_hf_docs", + "fetch_hf_docs", + "hub_repo_details", + "plan", + "hf_inspect_dataset", + "web_search", +} + +# Tool names that signal data manipulation workflows. +_DATA_PREP_TOOLS = {"hf_inspect_dataset", "dataset_tools", "hub_repo_details"} + + +def _model_family(model_name: str | None) -> str: + if not model_name: + return "other" + n = model_name.lower() + if "opus" in n: + return "opus" + if "sonnet" in n: + return "sonnet" + if "haiku" in n: + return "haiku" + if "kimi" in n: + return "kimi" + if "gpt" in n: + return "gpt" + if "deepseek" in n: + return "deepseek" + if "qwen" in n: + return "qwen" + if "llama" in n: + return "llama" + return "other" + + +def _turns_bucket(n: int) -> str: + if n < 5: + return "short" + if n <= 20: + return "medium" + return "long" + + +def _cost_bucket(cost_usd: float) -> str: + if cost_usd < 0.10: + return "low" + if cost_usd < 1.0: + return "med" + return "high" + + +def _flavor_to_gpu_tags(flavor: str) -> list[str]: + family = _GPU_FAMILY.get(flavor, "none") + tags = [f"gpu:{family}"] + if any(m in flavor for m in _MULTI_GPU_MARKERS): + tags.append("gpu:multi") + return tags + + +def _has_oom_signal(tool_outputs: Iterable[str]) -> bool: + for out in tool_outputs: + if not isinstance(out, str): + continue + low = out.lower() + if "outofmemoryerror" in low or "cuda out of memory" in low or "oom" in low: + return True + return False + + +def _infer_task_tag( + tool_names: set[str], + hf_job_submit_scripts: list[str], +) -> str | None: + """Return a ``task:*`` tag or None if we can't tell. + + Heuristic order: training > inference > data_prep > research_only. + """ + # training: any hf_jobs script with a Trainer/SFT/training keyword, OR uses + # hf_jobs at all and a script mentions training APIs. + for script in hf_job_submit_scripts: + low = script.lower() + if any( + k in low + for k in ( + "sftconfig", + "sfttrainer", + "trainer(", + "trainingarguments", + "grpo", + "dpo", + ".train(", + "transformers import", + "trainer import", + "fine-tune", + "finetune", + ) + ): + return "training" + + # inference: sessions that use inference tools but never hf_jobs/sandbox + uses_compute = bool(tool_names & {"hf_jobs", "sandbox_create", "sandbox_exec"}) + if not uses_compute and tool_names & {"inference", "generate", "run_inference"}: + return "inference" + + # data_prep: primarily dataset tools and no training/inference + if tool_names & _DATA_PREP_TOOLS and not uses_compute: + return "data_prep" + + # research_only: every tool used is in the research allow-list + if tool_names and tool_names <= _RESEARCH_ONLY_TOOLS: + return "research_only" + + return None + + +def tag_session(trajectory: dict) -> list[str]: + """Derive tags from a session trajectory. Pure function.""" + tags: set[str] = set() + + events: list[dict] = trajectory.get("events") or [] + messages: list[dict] = trajectory.get("messages") or [] + model_name: str | None = trajectory.get("model_name") + + # model + tags.add(f"model:{_model_family(model_name)}") + + # turns + user_turns = sum(1 for m in messages if m.get("role") == "user") + tags.add(f"turns:{_turns_bucket(user_turns)}") + + # cost + tool-name enumeration + outcome detection + cost_usd = 0.0 + tool_names: set[str] = set() + tool_outputs: list[str] = [] + hf_job_submit_count = 0 + hf_job_submit_scripts: list[str] = [] + hf_job_success_count = 0 + hf_job_fail_count = 0 + hf_job_push_to_hub = False + gpu_tags_seen: set[str] = set() + + # Outcome is the *last* terminal signal. Seed with "ongoing" — overridden + # if we see a terminal event. + outcome = "ongoing" + had_error = False + had_doom_loop = False + had_compact = False + + feedback_up = 0 + feedback_down = 0 + + sandbox_created = False + sandbox_hardware: str | None = None + sandbox_lifetime_s: int | None = None + + for ev in events: + et = ev.get("event_type") + data = ev.get("data") or {} + + if et == "llm_call": + cost_usd += float(data.get("cost_usd") or 0.0) + + elif et == "tool_call": + name = data.get("tool") + if name: + tool_names.add(name) + + elif et == "tool_output": + out = data.get("output") + if isinstance(out, str): + tool_outputs.append(out) + + elif et == "hf_job_submit": + hf_job_submit_count += 1 + if data.get("push_to_hub"): + hf_job_push_to_hub = True + flavor = data.get("flavor") or "cpu-basic" + for t in _flavor_to_gpu_tags(flavor): + gpu_tags_seen.add(t) + + elif et == "hf_job_complete": + final = (data.get("final_status") or "").lower() + if final in ("completed", "succeeded", "success"): + hf_job_success_count += 1 + elif final in ("failed", "error", "timeout", "cancelled"): + hf_job_fail_count += 1 + + elif et == "sandbox_create": + sandbox_created = True + sandbox_hardware = data.get("hardware") + + elif et == "sandbox_destroy": + lt = data.get("lifetime_s") + if isinstance(lt, (int, float)): + sandbox_lifetime_s = int(lt) + + elif et == "feedback": + rating = data.get("rating") + if rating == "up": + feedback_up += 1 + elif rating == "down": + feedback_down += 1 + + elif et == "error": + had_error = True + elif et == "turn_complete": + if not had_error: + outcome = "completed" + elif et == "interrupted": + outcome = "interrupted" + elif et == "compacted": + had_compact = True + elif et == "tool_log": + log_text = (data.get("log") or "").lower() + if "doom loop" in log_text: + had_doom_loop = True + + if had_error and outcome not in ("completed", "interrupted"): + outcome = "errored" + + tags.add(f"outcome:{outcome}") + if had_doom_loop: + tags.add("outcome:doom_loop") + if had_compact: + tags.add("outcome:context_exceeded") + + # tools + for name in tool_names: + tags.add(f"tool:{name}") + + # hf_jobs facets + if hf_job_submit_count >= 1: + tags.add("hf_job:submitted") + if hf_job_submit_count > 1: + tags.add("hf_job:multi") + if hf_job_success_count > 0: + tags.add("hf_job:succeeded") + if hf_job_fail_count > 0: + tags.add("hf_job:failed") + if hf_job_push_to_hub: + tags.add("hf_job:push_to_hub") + if _has_oom_signal(tool_outputs): + tags.add("hf_job:oom") + + # gpu tags (from all submitted jobs) + tags.update(gpu_tags_seen) + if "gpu:none" in tags and len(gpu_tags_seen) > 1: + # If any GPU flavor was used, drop the "none" tag for clarity. + tags.discard("gpu:none") + + # sandbox facets + if sandbox_created: + tags.add("sandbox:created") + if sandbox_hardware: + fam = _GPU_FAMILY.get(sandbox_hardware, "none") + tags.add("sandbox:cpu" if fam == "none" else "sandbox:gpu") + if sandbox_lifetime_s is not None and sandbox_lifetime_s > 1800: + tags.add("sandbox:long_lived") + + # feedback + if feedback_up and feedback_down: + tags.add("feedback:mixed") + elif feedback_up: + tags.add("feedback:up") + elif feedback_down: + tags.add("feedback:down") + else: + tags.add("feedback:none") + + # cost bucket + tags.add(f"cost:{_cost_bucket(cost_usd)}") + + # task heuristic (needs scripts — pull from the hf_job_submit events' + # matching tool_call arguments in the event list). + for ev in events: + if ev.get("event_type") == "tool_call": + data = ev.get("data") or {} + if data.get("tool") == "hf_jobs": + args = data.get("arguments") or {} + script = args.get("script") or args.get("command") or "" + if isinstance(script, str): + hf_job_submit_scripts.append(script) + + task_tag = _infer_task_tag(tool_names, hf_job_submit_scripts) + if task_tag: + tags.add(f"task:{task_tag}") + + return sorted(tags) diff --git a/agent/tools/__init__.py b/agent/tools/__init__.py index 14ef45669..65c793cba 100644 --- a/agent/tools/__init__.py +++ b/agent/tools/__init__.py @@ -20,6 +20,7 @@ ) from agent.tools.jobs_tool import HF_JOBS_TOOL_SPEC, HfJobsTool, hf_jobs_handler from agent.tools.types import ToolResult +from agent.tools.web_search_tool import WEB_SEARCH_TOOL_SPEC, web_search_handler __all__ = [ "ToolResult", @@ -36,4 +37,6 @@ "github_search_code_handler", "HF_INSPECT_DATASET_TOOL_SPEC", "hf_inspect_dataset_handler", + "WEB_SEARCH_TOOL_SPEC", + "web_search_handler", ] diff --git a/agent/tools/dataset_tools.py b/agent/tools/dataset_tools.py index ef3f3c81b..20add683d 100644 --- a/agent/tools/dataset_tools.py +++ b/agent/tools/dataset_tools.py @@ -423,7 +423,9 @@ def _format_parquet_files(data: dict, max_rows: int = 10) -> str | None: } -async def hf_inspect_dataset_handler(arguments: dict[str, Any], session=None) -> tuple[str, bool]: +async def hf_inspect_dataset_handler( + arguments: dict[str, Any], session=None +) -> tuple[str, bool]: """Handler for agent tool router""" try: hf_token = session.hf_token if session else None diff --git a/agent/tools/docs_tools.py b/agent/tools/docs_tools.py index a1782107e..ee40ef353 100644 --- a/agent/tools/docs_tools.py +++ b/agent/tools/docs_tools.py @@ -932,7 +932,7 @@ async def _get_api_search_tool_spec() -> dict[str, Any]: "• argilla — Data annotation, feedback, and human-in-the-loop workflows.\n" "• distilabel — Synthetic data generation and distillation pipelines.\n" "• microsoft-azure — Azure deployment and integration guides.\n" - "• kernels — Lightweight execution environments and notebook-style workflows.\n" + "• kernels — Load prebuilt compute kernels (E.g. flash-attn2) from the Hub via `attn_implementation`; avoids compiling flash-attn from source.\n" "• google-cloud — GCP deployment and serving workflows.\n" ), }, diff --git a/agent/tools/edit_utils.py b/agent/tools/edit_utils.py index 6a9a3295e..1c6b95819 100644 --- a/agent/tools/edit_utils.py +++ b/agent/tools/edit_utils.py @@ -10,18 +10,18 @@ # ── Unicode normalization map ──────────────────────────────────────────── UNICODE_MAP = { - "\u2013": "-", # en-dash - "\u2014": "-", # em-dash - "\u2212": "-", # minus sign - "\u2018": "'", # left single quote - "\u2019": "'", # right single quote - "\u201c": '"', # left double quote - "\u201d": '"', # right double quote - "\u00a0": " ", # non-breaking space - "\u2003": " ", # em space - "\u2002": " ", # en space - "\u200b": "", # zero-width space - "\ufeff": "", # BOM + "\u2013": "-", # en-dash + "\u2014": "-", # em-dash + "\u2212": "-", # minus sign + "\u2018": "'", # left single quote + "\u2019": "'", # right single quote + "\u201c": '"', # left double quote + "\u201d": '"', # right double quote + "\u00a0": " ", # non-breaking space + "\u2003": " ", # em space + "\u2002": " ", # en space + "\u200b": "", # zero-width space + "\ufeff": "", # BOM } @@ -59,12 +59,12 @@ def _build_stripped(text: str, strip_fn): line_start_map[i] = original byte offset of the start of line i. """ orig_lines = text.split("\n") - stripped_lines = [strip_fn(l) for l in orig_lines] + stripped_lines = [strip_fn(line) for line in orig_lines] return "\n".join(stripped_lines), orig_lines, stripped_lines # Pass 2 — right-trim c_rt, c_orig_lines, c_rt_lines = _build_stripped(content, str.rstrip) - p_rt = "\n".join(l.rstrip() for l in pattern.split("\n")) + p_rt = "\n".join(line.rstrip() for line in pattern.split("\n")) idx = c_rt.find(p_rt) if idx != -1: orig_idx = _map_back(idx, c_orig_lines, c_rt_lines) @@ -72,7 +72,7 @@ def _build_stripped(text: str, strip_fn): # Pass 3 — both-sides trim c_st, _, c_st_lines = _build_stripped(content, str.strip) - p_st = "\n".join(l.strip() for l in pattern.split("\n")) + p_st = "\n".join(line.strip() for line in pattern.split("\n")) idx = c_st.find(p_st) if idx != -1: orig_idx = _map_back(idx, c_orig_lines, c_st_lines) @@ -114,7 +114,9 @@ def _map_back( return 0 -def fuzzy_find_original_match(content: str, pattern: str) -> tuple[str | None, str | None]: +def fuzzy_find_original_match( + content: str, pattern: str +) -> tuple[str | None, str | None]: """Find the *original* text in content that matches pattern fuzzily. Returns (original_matched_text, match_note) or (None, None). @@ -224,7 +226,9 @@ def apply_edit( return new_content, 1, fuzzy_note else: - raise ValueError(f"Unknown edit mode: {mode}. Use replace, append_after, or prepend_before.") + raise ValueError( + f"Unknown edit mode: {mode}. Use replace, append_after, or prepend_before." + ) # ── Syntax validation (Python) ─────────────────────────────────────────── @@ -255,14 +259,15 @@ def validate_python(content: str, path: str = "") -> list[str]: return warnings # 2. Training script heuristics - if any(kw in content for kw in ("TrainingArguments", "SFTConfig", "DPOConfig", "GRPOConfig")): + if any( + kw in content + for kw in ("TrainingArguments", "SFTConfig", "DPOConfig", "GRPOConfig") + ): if "push_to_hub" not in content: warnings.append( "Training script warning: no 'push_to_hub' found — model may be lost when job ends" ) if "hub_model_id" not in content: - warnings.append( - "Training script warning: no 'hub_model_id' found" - ) + warnings.append("Training script warning: no 'hub_model_id' found") return warnings diff --git a/agent/tools/hf_repo_files_tool.py b/agent/tools/hf_repo_files_tool.py index fd39a488f..d2226ac17 100644 --- a/agent/tools/hf_repo_files_tool.py +++ b/agent/tools/hf_repo_files_tool.py @@ -5,15 +5,14 @@ """ import asyncio -from typing import Any, Dict, Literal, Optional +from typing import Any, Dict, Optional from huggingface_hub import HfApi, hf_hub_download from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError +from agent.core.hub_artifacts import is_known_hub_artifact, register_hub_artifact from agent.tools.types import ToolResult -OperationType = Literal["list", "read", "upload", "delete"] - async def _async_call(func, *args, **kwargs): """Wrap synchronous HfApi calls for async context.""" @@ -39,8 +38,9 @@ def _format_size(size_bytes: int) -> str: class HfRepoFilesTool: """Tool for file operations on HF repos.""" - def __init__(self, hf_token: Optional[str] = None): + def __init__(self, hf_token: Optional[str] = None, session: Any = None): self.api = HfApi(token=hf_token) + self.session = session async def execute(self, args: Dict[str, Any]) -> ToolResult: """Execute the specified operation.""" @@ -61,7 +61,9 @@ async def execute(self, args: Dict[str, Any]) -> ToolResult: if handler: return await handler(args) else: - return self._error(f"Unknown operation: {operation}. Valid: list, read, upload, delete") + return self._error( + f"Unknown operation: {operation}. Valid: list, read, upload, delete" + ) except RepositoryNotFoundError: return self._error(f"Repository not found: {args.get('repo_id')}") @@ -96,17 +98,23 @@ async def _list(self, args: Dict[str, Any]) -> ToolResult: revision = args.get("revision", "main") path = args.get("path", "") - items = list(await _async_call( - self.api.list_repo_tree, - repo_id=repo_id, - repo_type=repo_type, - revision=revision, - path_in_repo=path, - recursive=True, - )) + items = list( + await _async_call( + self.api.list_repo_tree, + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + path_in_repo=path, + recursive=True, + ) + ) if not items: - return {"formatted": f"No files in {repo_id}", "totalResults": 0, "resultsShared": 0} + return { + "formatted": f"No files in {repo_id}", + "totalResults": 0, + "resultsShared": 0, + } lines = [] total_size = 0 @@ -118,9 +126,16 @@ async def _list(self, args: Dict[str, Any]) -> ToolResult: lines.append(f"{item.path}/") url = _build_repo_url(repo_id, repo_type) - response = f"**{repo_id}** ({len(items)} files, {_format_size(total_size)})\n{url}/tree/{revision}\n\n" + "\n".join(lines) + response = ( + f"**{repo_id}** ({len(items)} files, {_format_size(total_size)})\n{url}/tree/{revision}\n\n" + + "\n".join(lines) + ) - return {"formatted": response, "totalResults": len(items), "resultsShared": len(items)} + return { + "formatted": response, + "totalResults": len(items), + "resultsShared": len(items), + } async def _read(self, args: Dict[str, Any]) -> ToolResult: """Read file content from a repository.""" @@ -160,8 +175,13 @@ async def _read(self, args: Dict[str, Any]) -> ToolResult: except UnicodeDecodeError: import os + size = os.path.getsize(file_path) - return {"formatted": f"Binary file ({_format_size(size)})", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"Binary file ({_format_size(size)})", + "totalResults": 1, + "resultsShared": 1, + } async def _upload(self, args: Dict[str, Any]) -> ToolResult: """Upload content to a repository.""" @@ -194,6 +214,16 @@ async def _upload(self, args: Dict[str, Any]) -> ToolResult: create_pr=create_pr, ) + if not create_pr and is_known_hub_artifact(self.session, repo_id, repo_type): + await _async_call( + register_hub_artifact, + self.api, + repo_id, + repo_type, + session=self.session, + force=path == "README.md", + ) + url = _build_repo_url(repo_id, repo_type) if create_pr and hasattr(result, "pr_url"): response = f"**Uploaded as PR**\n{result.pr_url}" @@ -235,7 +265,12 @@ async def _delete(self, args: Dict[str, Any]) -> ToolResult: def _error(self, message: str) -> ToolResult: """Return an error result.""" - return {"formatted": message, "totalResults": 0, "resultsShared": 0, "isError": True} + return { + "formatted": message, + "totalResults": 0, + "resultsShared": 0, + "isError": True, + } # Tool specification @@ -312,11 +347,13 @@ def _error(self, message: str) -> ToolResult: } -async def hf_repo_files_handler(arguments: Dict[str, Any], session=None) -> tuple[str, bool]: +async def hf_repo_files_handler( + arguments: Dict[str, Any], session=None +) -> tuple[str, bool]: """Handler for agent tool router.""" try: hf_token = session.hf_token if session else None - tool = HfRepoFilesTool(hf_token=hf_token) + tool = HfRepoFilesTool(hf_token=hf_token, session=session) result = await tool.execute(arguments) return result["formatted"], not result.get("isError", False) except Exception as e: diff --git a/agent/tools/hf_repo_git_tool.py b/agent/tools/hf_repo_git_tool.py index d7e2323a3..672186c6d 100644 --- a/agent/tools/hf_repo_git_tool.py +++ b/agent/tools/hf_repo_git_tool.py @@ -5,21 +5,14 @@ """ import asyncio -from typing import Any, Dict, Literal, Optional +from typing import Any, Dict, Optional from huggingface_hub import HfApi from huggingface_hub.utils import RepositoryNotFoundError +from agent.core.hub_artifacts import register_hub_artifact from agent.tools.types import ToolResult -OperationType = Literal[ - "create_branch", "delete_branch", - "create_tag", "delete_tag", - "list_refs", - "create_pr", "list_prs", "get_pr", "merge_pr", "close_pr", "comment_pr", "change_pr_status", - "create_repo", "update_repo", -] - async def _async_call(func, *args, **kwargs): """Wrap synchronous HfApi calls for async context.""" @@ -36,8 +29,9 @@ def _build_repo_url(repo_id: str, repo_type: str = "model") -> str: class HfRepoGitTool: """Tool for git-like operations on HF repos.""" - def __init__(self, hf_token: Optional[str] = None): + def __init__(self, hf_token: Optional[str] = None, session: Any = None): self.api = HfApi(token=hf_token) + self.session = session async def execute(self, args: Dict[str, Any]) -> ToolResult: """Execute the specified operation.""" @@ -131,7 +125,11 @@ async def _create_branch(self, args: Dict[str, Any]) -> ToolResult: ) url = f"{_build_repo_url(repo_id, repo_type)}/tree/{branch}" - return {"formatted": f"**Branch created:** {branch}\n{url}", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**Branch created:** {branch}\n{url}", + "totalResults": 1, + "resultsShared": 1, + } async def _delete_branch(self, args: Dict[str, Any]) -> ToolResult: """Delete a branch.""" @@ -152,7 +150,11 @@ async def _delete_branch(self, args: Dict[str, Any]) -> ToolResult: repo_type=repo_type, ) - return {"formatted": f"**Branch deleted:** {branch}", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**Branch deleted:** {branch}", + "totalResults": 1, + "resultsShared": 1, + } # ========================================================================= # TAG OPERATIONS @@ -183,7 +185,11 @@ async def _create_tag(self, args: Dict[str, Any]) -> ToolResult: ) url = f"{_build_repo_url(repo_id, repo_type)}/tree/{tag}" - return {"formatted": f"**Tag created:** {tag}\n{url}", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**Tag created:** {tag}\n{url}", + "totalResults": 1, + "resultsShared": 1, + } async def _delete_tag(self, args: Dict[str, Any]) -> ToolResult: """Delete a tag.""" @@ -204,7 +210,11 @@ async def _delete_tag(self, args: Dict[str, Any]) -> ToolResult: repo_type=repo_type, ) - return {"formatted": f"**Tag deleted:** {tag}", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**Tag deleted:** {tag}", + "totalResults": 1, + "resultsShared": 1, + } # ========================================================================= # LIST REFS @@ -226,7 +236,9 @@ async def _list_refs(self, args: Dict[str, Any]) -> ToolResult: ) branches = [b.name for b in refs.branches] if refs.branches else [] - tags = [t.name for t in refs.tags] if hasattr(refs, 'tags') and refs.tags else [] + tags = ( + [t.name for t in refs.tags] if hasattr(refs, "tags") and refs.tags else [] + ) url = _build_repo_url(repo_id, repo_type) lines = [f"**{repo_id}**", url, ""] @@ -241,7 +253,11 @@ async def _list_refs(self, args: Dict[str, Any]) -> ToolResult: else: lines.append("**Tags:** none") - return {"formatted": "\n".join(lines), "totalResults": len(branches) + len(tags), "resultsShared": len(branches) + len(tags)} + return { + "formatted": "\n".join(lines), + "totalResults": len(branches) + len(tags), + "resultsShared": len(branches) + len(tags), + } # ========================================================================= # PR OPERATIONS @@ -270,7 +286,7 @@ async def _create_pr(self, args: Dict[str, Any]) -> ToolResult: url = f"{_build_repo_url(repo_id, repo_type)}/discussions/{result.num}" return { - "formatted": f"**Draft PR #{result.num} created:** {title}\n{url}\n\nAdd commits via upload with revision=\"refs/pr/{result.num}\"", + "formatted": f'**Draft PR #{result.num} created:** {title}\n{url}\n\nAdd commits via upload with revision="refs/pr/{result.num}"', "totalResults": 1, "resultsShared": 1, } @@ -285,17 +301,27 @@ async def _list_prs(self, args: Dict[str, Any]) -> ToolResult: repo_type = args.get("repo_type", "model") status = args.get("status", "all") # open, closed, all - discussions = list(self.api.get_repo_discussions( - repo_id=repo_id, - repo_type=repo_type, - discussion_status=status if status != "all" else None, - )) + discussions = list( + self.api.get_repo_discussions( + repo_id=repo_id, + repo_type=repo_type, + discussion_status=status if status != "all" else None, + ) + ) if not discussions: - return {"formatted": f"No discussions in {repo_id}", "totalResults": 0, "resultsShared": 0} + return { + "formatted": f"No discussions in {repo_id}", + "totalResults": 0, + "resultsShared": 0, + } url = _build_repo_url(repo_id, repo_type) - lines = [f"**{repo_id}** - {len(discussions)} discussions", f"{url}/discussions", ""] + lines = [ + f"**{repo_id}** - {len(discussions)} discussions", + f"{url}/discussions", + "", + ] for d in discussions[:20]: if d.status == "draft": @@ -309,7 +335,11 @@ async def _list_prs(self, args: Dict[str, Any]) -> ToolResult: type_label = "PR" if d.is_pull_request else "D" lines.append(f"{status_label} #{d.num} [{type_label}] {d.title}") - return {"formatted": "\n".join(lines), "totalResults": len(discussions), "resultsShared": min(20, len(discussions))} + return { + "formatted": "\n".join(lines), + "totalResults": len(discussions), + "resultsShared": min(20, len(discussions)), + } async def _get_pr(self, args: Dict[str, Any]) -> ToolResult: """Get PR details.""" @@ -335,7 +365,7 @@ async def _get_pr(self, args: Dict[str, Any]) -> ToolResult: "draft": "Draft", "open": "Open", "merged": "Merged", - "closed": "Closed" + "closed": "Closed", } status = status_map.get(pr.status, pr.status.capitalize()) type_label = "Pull Request" if pr.is_pull_request else "Discussion" @@ -349,9 +379,13 @@ async def _get_pr(self, args: Dict[str, Any]) -> ToolResult: if pr.is_pull_request: if pr.status == "draft": - lines.append(f"\nTo add commits: upload with revision=\"refs/pr/{pr_num}\"") + lines.append( + f'\nTo add commits: upload with revision="refs/pr/{pr_num}"' + ) elif pr.status == "open": - lines.append(f"\nTo add commits: upload with revision=\"refs/pr/{pr_num}\"") + lines.append( + f'\nTo add commits: upload with revision="refs/pr/{pr_num}"' + ) return {"formatted": "\n".join(lines), "totalResults": 1, "resultsShared": 1} @@ -377,7 +411,11 @@ async def _merge_pr(self, args: Dict[str, Any]) -> ToolResult: ) url = f"{_build_repo_url(repo_id, repo_type)}/discussions/{pr_num}" - return {"formatted": f"**PR #{pr_num} merged**\n{url}", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**PR #{pr_num} merged**\n{url}", + "totalResults": 1, + "resultsShared": 1, + } async def _close_pr(self, args: Dict[str, Any]) -> ToolResult: """Close a PR/discussion.""" @@ -401,7 +439,11 @@ async def _close_pr(self, args: Dict[str, Any]) -> ToolResult: repo_type=repo_type, ) - return {"formatted": f"**Discussion #{pr_num} closed**", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**Discussion #{pr_num} closed**", + "totalResults": 1, + "resultsShared": 1, + } async def _comment_pr(self, args: Dict[str, Any]) -> ToolResult: """Add a comment to a PR/discussion.""" @@ -427,7 +469,11 @@ async def _comment_pr(self, args: Dict[str, Any]) -> ToolResult: ) url = f"{_build_repo_url(repo_id, repo_type)}/discussions/{pr_num}" - return {"formatted": f"**Comment added to #{pr_num}**\n{url}", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**Comment added to #{pr_num}**\n{url}", + "totalResults": 1, + "resultsShared": 1, + } async def _change_pr_status(self, args: Dict[str, Any]) -> ToolResult: """Change PR/discussion status (mainly to convert draft to open).""" @@ -455,7 +501,11 @@ async def _change_pr_status(self, args: Dict[str, Any]) -> ToolResult: ) url = f"{_build_repo_url(repo_id, repo_type)}/discussions/{pr_num}" - return {"formatted": f"**PR #{pr_num} status changed to {new_status}**\n{url}", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**PR #{pr_num} status changed to {new_status}**\n{url}", + "totalResults": 1, + "resultsShared": 1, + } # ========================================================================= # REPO MANAGEMENT @@ -473,7 +523,9 @@ async def _create_repo(self, args: Dict[str, Any]) -> ToolResult: space_sdk = args.get("space_sdk") if repo_type == "space" and not space_sdk: - return self._error("space_sdk required for spaces (gradio/streamlit/docker/static)") + return self._error( + "space_sdk required for spaces (gradio/streamlit/docker/static)" + ) kwargs = { "repo_id": repo_id, @@ -485,6 +537,17 @@ async def _create_repo(self, args: Dict[str, Any]) -> ToolResult: kwargs["space_sdk"] = space_sdk result = await _async_call(self.api.create_repo, **kwargs) + extra_metadata = None + if repo_type == "space" and space_sdk: + extra_metadata = {"sdk": space_sdk} + await _async_call( + register_hub_artifact, + self.api, + repo_id, + repo_type, + session=self.session, + extra_metadata=extra_metadata, + ) return { "formatted": f"**Repository created:** {repo_id}\n**Private:** {private}\n{result}", @@ -504,7 +567,9 @@ async def _update_repo(self, args: Dict[str, Any]) -> ToolResult: gated = args.get("gated") if private is None and gated is None: - return self._error("Specify private (bool) or gated ('auto'/'manual'/false)") + return self._error( + "Specify private (bool) or gated ('auto'/'manual'/false)" + ) kwargs = {"repo_id": repo_id, "repo_type": repo_type} if private is not None: @@ -521,11 +586,20 @@ async def _update_repo(self, args: Dict[str, Any]) -> ToolResult: changes.append(f"gated={gated}") url = f"{_build_repo_url(repo_id, repo_type)}/settings" - return {"formatted": f"**Settings updated:** {', '.join(changes)}\n{url}", "totalResults": 1, "resultsShared": 1} + return { + "formatted": f"**Settings updated:** {', '.join(changes)}\n{url}", + "totalResults": 1, + "resultsShared": 1, + } def _error(self, message: str) -> ToolResult: """Return an error result.""" - return {"formatted": message, "totalResults": 0, "resultsShared": 0, "isError": True} + return { + "formatted": message, + "totalResults": 0, + "resultsShared": 0, + "isError": True, + } # Tool specification @@ -571,10 +645,20 @@ def _error(self, message: str) -> ToolResult: "operation": { "type": "string", "enum": [ - "create_branch", "delete_branch", - "create_tag", "delete_tag", "list_refs", - "create_pr", "list_prs", "get_pr", "merge_pr", "close_pr", "comment_pr", "change_pr_status", - "create_repo", "update_repo", + "create_branch", + "delete_branch", + "create_tag", + "delete_tag", + "list_refs", + "create_pr", + "list_prs", + "get_pr", + "merge_pr", + "close_pr", + "comment_pr", + "change_pr_status", + "create_repo", + "update_repo", ], "description": "Operation to execute", }, @@ -653,11 +737,13 @@ def _error(self, message: str) -> ToolResult: } -async def hf_repo_git_handler(arguments: Dict[str, Any], session=None) -> tuple[str, bool]: +async def hf_repo_git_handler( + arguments: Dict[str, Any], session=None +) -> tuple[str, bool]: """Handler for agent tool router.""" try: hf_token = session.hf_token if session else None - tool = HfRepoGitTool(hf_token=hf_token) + tool = HfRepoGitTool(hf_token=hf_token, session=session) result = await tool.execute(arguments) return result["formatted"], not result.get("isError", False) except Exception as e: diff --git a/agent/tools/jobs_tool.py b/agent/tools/jobs_tool.py index 2c6ebf6c7..f9afe782a 100644 --- a/agent/tools/jobs_tool.py +++ b/agent/tools/jobs_tool.py @@ -7,20 +7,24 @@ import asyncio import base64 import http.client -import os -import re -from typing import Any, Dict, Literal, Optional, Callable, Awaitable - import logging +import re +import shlex +from typing import Any, Awaitable, Callable, Dict, Optional import httpx from huggingface_hub import HfApi from huggingface_hub.utils import HfHubHTTPError +from agent.core.hf_access import ( + JobsAccessError, + is_billing_error, + resolve_jobs_namespace, +) +from agent.core.hub_artifacts import build_hub_artifact_sitecustomize from agent.core.session import Event +from agent.tools.trackio_seed import ensure_trackio_dashboard from agent.tools.types import ToolResult - -logger = logging.getLogger(__name__) from agent.tools.utilities import ( format_job_details, format_jobs_table, @@ -28,6 +32,8 @@ format_scheduled_jobs_table, ) +logger = logging.getLogger(__name__) + # Hardware flavors CPU_FLAVORS = ["cpu-basic", "cpu-upgrade"] GPU_FLAVORS = [ @@ -57,24 +63,6 @@ "l4x1(8vCPU/30GB/GPU 24GB), l4x4(48vCPU/186GB/GPU 96GB), " "l40sx1(8vCPU/62GB/GPU 48GB), l40sx4(48vCPU/382GB/GPU 192GB), l40sx8(192vCPU/1534GB/GPU 384GB)" ) -SPECIALIZED_FLAVORS = ["inf2x6"] -ALL_FLAVORS = CPU_FLAVORS + GPU_FLAVORS + SPECIALIZED_FLAVORS - -# Operation names -OperationType = Literal[ - "run", - "ps", - "logs", - "inspect", - "cancel", - "scheduled run", - "scheduled ps", - "scheduled inspect", - "scheduled delete", - "scheduled suspend", - "scheduled resume", -] - # Constants UV_DEFAULT_IMAGE = "ghcr.io/astral-sh/uv:python3.12-bookworm" @@ -117,11 +105,11 @@ def _filter_uv_install_output(logs: list[str]) -> list[str]: return logs -_ANSI_RE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07') +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07") def _strip_ansi(text: str) -> str: - return _ANSI_RE.sub('', text) + return _ANSI_RE.sub("", text) _DEFAULT_ENV = { @@ -233,6 +221,26 @@ def _resolve_uv_command( return _build_uv_command(script, with_deps, python, script_args) +def _wrap_command_with_artifact_bootstrap( + command: list[str], session: Any = None +) -> list[str]: + """Install sitecustomize hooks before the user command runs in HF Jobs.""" + sitecustomize = build_hub_artifact_sitecustomize(session) + if not sitecustomize: + return command + + encoded = base64.b64encode(sitecustomize.encode("utf-8")).decode("ascii") + original_command = shlex.join(command) + shell = ( + 'set -e; _ml_intern_artifacts_dir="$(mktemp -d)"; ' + f"printf %s {shlex.quote(encoded)} | base64 -d " + '> "$_ml_intern_artifacts_dir/sitecustomize.py"; ' + 'export PYTHONPATH="$_ml_intern_artifacts_dir${PYTHONPATH:+:$PYTHONPATH}"; ' + f"exec {original_command}" + ) + return ["/bin/sh", "-lc", shell] + + async def _async_call(func, *args, **kwargs): """Wrap synchronous HfApi calls for async context""" return await asyncio.to_thread(func, *args, **kwargs) @@ -298,6 +306,7 @@ def __init__( self, hf_token: Optional[str] = None, namespace: Optional[str] = None, + jobs_access: Any = None, log_callback: Optional[Callable[[str], Awaitable[None]]] = None, session: Any = None, tool_call_id: Optional[str] = None, @@ -305,6 +314,7 @@ def __init__( self.hf_token = hf_token self.api = HfApi(token=hf_token) self.namespace = namespace + self.jobs_access = jobs_access self.log_callback = log_callback self.session = session self.tool_call_id = tool_call_id @@ -379,6 +389,31 @@ async def execute(self, params: Dict[str, Any]) -> ToolResult: "isError": True, } + async def _seed_trackio_dashboard(self, space_id: str) -> None: + """Idempotently install trackio dashboard files into *space_id* before + the job runs. Surfaces seed progress as tool_log events but never + raises — a seed failure should not block job submission, since trackio + often still works when the Space already has dashboard code from a + previous run. + """ + loop = asyncio.get_running_loop() + + def _log(msg: str) -> None: + if self.session is None: + return + loop.call_soon_threadsafe( + self.session.event_queue.put_nowait, + Event(event_type="tool_log", data={"tool": "hf_jobs", "log": msg}), + ) + + try: + await asyncio.to_thread( + ensure_trackio_dashboard, space_id, self.hf_token, _log + ) + except Exception as e: + logger.warning(f"trackio dashboard seed failed for {space_id}: {e}") + _log(f"trackio dashboard seed failed: {e}") + async def _wait_for_job_completion( self, job_id: str, namespace: Optional[str] = None ) -> tuple[str, list[str]]: @@ -403,7 +438,9 @@ async def _wait_for_job_completion( def log_producer(): try: # fetch_job_logs is a blocking sync generator - logs_gen = self.api.fetch_job_logs(job_id=job_id, namespace=namespace) + logs_gen = self.api.fetch_job_logs( + job_id=job_id, namespace=namespace + ) for line in logs_gen: # Push line to queue thread-safely loop.call_soon_threadsafe(queue.put_nowait, line) @@ -527,17 +564,66 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult: image = args.get("image", "python:3.12") job_type = "Docker" + command = _wrap_command_with_artifact_bootstrap(command, self.session) + # Run the job - job = await _async_call( - self.api.run_job, - image=image, - command=command, - env=_add_default_env(args.get("env")), - secrets=_add_environment_variables(args.get("secrets"), self.hf_token), - flavor=args.get("hardware_flavor", "cpu-basic"), - timeout=args.get("timeout", "30m"), - namespace=self.namespace, - ) + flavor = args.get("hardware_flavor", "cpu-basic") + timeout_str = args.get("timeout", "30m") + + # Trackio: agent-declared space + project become env vars on the job + # so trackio.init() picks them up automatically. We also surface them + # in tool_state_change so the frontend can embed the dashboard. + env_dict = _add_default_env(args.get("env")) + trackio_space_id = args.get("trackio_space_id") + trackio_project = args.get("trackio_project") + if trackio_space_id: + env_dict["TRACKIO_SPACE_ID"] = trackio_space_id + await self._seed_trackio_dashboard(trackio_space_id) + if trackio_project: + env_dict["TRACKIO_PROJECT"] = trackio_project + + try: + job = await _async_call( + self.api.run_job, + image=image, + command=command, + env=env_dict, + secrets=_add_environment_variables( + args.get("secrets"), self.hf_token + ), + flavor=flavor, + timeout=timeout_str, + namespace=self.namespace, + ) + except HfHubHTTPError as e: + if is_billing_error(str(e)): + if self.session and self.tool_call_id: + await self.session.send_event( + Event( + event_type="tool_state_change", + data={ + "tool_call_id": self.tool_call_id, + "tool": "hf_jobs", + "state": "billing_required", + "namespace": self.namespace, + }, + ) + ) + return { + "formatted": ( + f"Hugging Face Jobs rejected this run because the " + f"namespace `{self.namespace}` has no available credits. " + "HF Jobs are billed with namespace credits, which are " + "separate from HF Pro membership. Tell the user to add " + "credits at https://huggingface.co/settings/billing — " + "once topped up, re-run this same job. (Switching " + "namespaces is fine if another wallet has credits.)" + ), + "totalResults": 0, + "resultsShared": 0, + "isError": True, + } + raise # Track job ID for cancellation on interrupt if self.session: @@ -545,17 +631,55 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult: # Send job URL immediately after job creation (before waiting for completion) if self.session and self.tool_call_id: + state_data: Dict[str, Any] = { + "tool_call_id": self.tool_call_id, + "tool": "hf_jobs", + "state": "running", + "jobUrl": job.url, + } + if trackio_space_id: + state_data["trackioSpaceId"] = trackio_space_id + if trackio_project: + state_data["trackioProject"] = trackio_project await self.session.send_event( - Event( - event_type="tool_state_change", - data={ - "tool_call_id": self.tool_call_id, - "tool": "hf_jobs", - "state": "running", - "jobUrl": job.url, - }, - ) + Event(event_type="tool_state_change", data=state_data) + ) + + # Telemetry: job submission + completion (infra consumption signal). + submit_ts = None + if self.session: + from agent.core import telemetry + + submit_ts = await telemetry.record_hf_job_submit( + self.session, + job, + { + **args, + "hardware_flavor": flavor, + "timeout": timeout_str, + "namespace": self.namespace, + }, + image=image, + job_type=job_type, + ) + # Top-up signal: this submit succeeded after a prior billing + # block in the same session, and we haven't fired the event + # yet — the user came back from the HF billing flow. + events = self.session.logged_events + already_fired = any( + e.get("event_type") == "credits_topped_up" for e in events ) + if not already_fired: + blocked = any( + e.get("event_type") == "tool_state_change" + and (e.get("data") or {}).get("state") == "billing_required" + for e in events + ) + if blocked: + await telemetry.record_credits_topped_up( + self.session, + namespace=self.namespace, + ) # Wait for completion and stream logs logger.info(f"{job_type} job started: {job.url}") @@ -566,29 +690,55 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult: namespace=self.namespace, ) + if self.session and submit_ts is not None: + from agent.core import telemetry + + usage = await telemetry.record_hf_job_complete( + self.session, + job, + flavor=flavor, + final_status=final_status, + submit_ts=submit_ts, + ) + if self.tool_call_id: + from agent.core.yolo_budget import reconcile_budget_reservation + + reconcile_budget_reservation( + self.session, + self.tool_call_id, + usage.get("estimated_cost_usd") + if isinstance(usage, dict) + else None, + allow_zero_actual=True, + ) + # Untrack job ID (completed or failed, no longer needs cancellation) if self.session: self.session._running_job_ids.discard(job.id) # Notify frontend of final status if self.session and self.tool_call_id: + final_data: Dict[str, Any] = { + "tool_call_id": self.tool_call_id, + "tool": "hf_jobs", + "state": final_status.lower(), + "jobUrl": job.url, + } + if trackio_space_id: + final_data["trackioSpaceId"] = trackio_space_id + if trackio_project: + final_data["trackioProject"] = trackio_project await self.session.send_event( - Event( - event_type="tool_state_change", - data={ - "tool_call_id": self.tool_call_id, - "tool": "hf_jobs", - "state": final_status.lower(), - "jobUrl": job.url, - }, - ) + Event(event_type="tool_state_change", data=final_data) ) # Filter out UV package installation output filtered_logs = _filter_uv_install_output(all_logs) # Format all logs for the agent - log_text = _strip_ansi("\n".join(filtered_logs)) if filtered_logs else "(no logs)" + log_text = ( + _strip_ansi("\n".join(filtered_logs)) if filtered_logs else "(no logs)" + ) response = f"""{job_type} job completed! @@ -780,6 +930,8 @@ async def _scheduled_run(self, args: Dict[str, Any]) -> ToolResult: image = args.get("image", "python:3.12") job_type = "Docker" + command = _wrap_command_with_artifact_bootstrap(command, self.session) + # Create scheduled job scheduled_job = await _async_call( self.api.create_scheduled_job, @@ -953,9 +1105,42 @@ async def _resume_scheduled_job(self, args: Dict[str, Any]) -> ToolResult: "- You MUST have called github_find_examples + github_read_file to find a working reference implementation. " "Scripts based on your internal knowledge WILL use outdated APIs and fail.\n" "- You MUST have validated dataset format via hf_inspect_dataset or hub_repo_details.\n" + "- For non-trivial scripts, write the script in the session sandbox, run syntax/import validation, " + "run a tiny smoke test, and for training scripts make sure one training step succeeds, plus one " + "evaluation step when the final workflow includes evaluation or an eval split is available, then " + "submit the exact tested script source or exact tested sandbox file. " + "Do NOT reconstruct a similar script from memory.\n" + "- If the job runs on GPU, or the script loads a model, uses CUDA, bf16/fp16, quantization, flash attention, " + "or torch.compile, you MUST create a GPU sandbox with sandbox_create first, run a tiny smoke test there, " + "and fix failures before submitting. If skipped, state why before calling hf_jobs.\n" + "- Do NOT install compiled flash-attn or use attn_implementation='flash_attention_2'. " + "For accelerated attention, use the HF kernels package with a Hub kernel such as " + "kernels-community/flash-attn2, and smoke-test the exact same attn_implementation. " + "Flash-attention Hub kernels require Ampere-or-newer GPUs unless their docs say otherwise: " + "never choose T4 sandboxes or T4 HF Jobs for scripts that use a flash-attention kernel, " + "because T4 is pre-Ampere. Use A10G, A100, H100, or another compatible newer GPU.\n" + "- Do NOT rely on preinstalled ML packages. Install/upgrade the latest compatible core stack " + "in the sandbox and include the same packages in dependencies: torch, transformers, trl, " + "accelerate, datasets, trackio, and kernels~=0.12.0 when using Hub kernels. " + "Use unpinned latest stable versions by default for the rest of the core stack; constrain " + "kernels to kernels~=0.12.0. Pin other versions only when current docs/examples require a specific " + "compatibility set or a smoke test shows latest is incompatible. Print installed versions " + "before model loading. If kernels and transformers are incompatible, fix the package set " + "or choose another compatible Hub kernel, then rerun the smoke test.\n" "- Training config MUST include push_to_hub=True and hub_model_id. " "Job storage is EPHEMERAL — all files are deleted when the job ends. Without push_to_hub, trained models are lost permanently.\n" - "- Include trackio monitoring and provide the dashboard URL to the user.\n\n" + "- Training scripts MUST fail fast on missing dataset columns, placeholder repo IDs, placeholder Trackio IDs, " + "missing hub_model_id, or missing push_to_hub=True.\n" + "- Do NOT leave placeholders such as , , , TODO, " + "or similar unfinished values in scripts or job arguments.\n" + "- dependencies MUST include every imported third-party package, including the core ML stack " + "torch, transformers, trl, accelerate, datasets, trackio, kernels~=0.12.0 when using Hub kernels, " + "and extras such as peft, bitsandbytes, sentencepiece, or protobuf when used.\n" + "- Include trackio monitoring and provide the dashboard URL to the user. " + "When the script uses report_to='trackio', also pass `trackio_space_id` " + "(pattern only: '/ml-intern-<8char>'; replace , e.g. 'alice/ml-intern-a1b2c3d4') " + "and `trackio_project` as tool args — " + "they are injected as TRACKIO_SPACE_ID/TRACKIO_PROJECT env vars and let the UI embed the live dashboard.\n\n" "BATCH/ABLATION JOBS: Submit ONE job first. Check logs to confirm it starts training successfully. " "Only then submit the remaining jobs. Never submit all at once — if there's a bug, all jobs fail.\n\n" "Operations: run, ps, logs, inspect, cancel, scheduled run/ps/inspect/delete/suspend/resume.\n\n" @@ -968,8 +1153,8 @@ async def _resume_scheduled_job(self, args: Dict[str, Any]) -> ToolResult: "3. Upgrade to larger GPU (a10g→a100→h100)\n" "Do NOT switch training methods (e.g. full SFT to LoRA) or reduce max_length — those change what the user gets and require explicit approval.\n\n" "Examples:\n" - "Training: {'operation': 'run', 'script': '/app/train.py', 'dependencies': ['transformers', 'trl', 'torch', 'datasets', 'trackio'], 'hardware_flavor': 'a100-large', 'timeout': '8h'}\n" - "Monitor: {'operation': 'ps'}, {'operation': 'logs', 'job_id': 'xxx'}, {'operation': 'cancel', 'job_id': 'xxx'}" + "Training: {'operation': 'run', 'script': '/app/train.py', 'dependencies': ['torch', 'transformers', 'trl', 'accelerate', 'datasets', 'trackio', 'kernels~=0.12.0'], 'hardware_flavor': 'a100-large', 'timeout': '8h'}\n" + "Monitor: {'operation': 'ps'}, {'operation': 'logs', 'job_id': 'xxx'}, {'operation': 'cancel', 'job_id': 'xxx'}\n" "Docker: {'operation': 'run', 'command': ['duckdb', '-c', 'select 1 + 2'], 'image': 'duckdb/duckdb', 'hardware_flavor': 'cpu-basic', 'timeout': '1h'}\n" ), "parameters": { @@ -995,8 +1180,11 @@ async def _resume_scheduled_job(self, args: Dict[str, Any]) -> ToolResult: "script": { "type": "string", "description": ( - "Python code or sandbox file path (e.g. '/app/train.py') or URL. " + "Python code, sandbox file path (e.g. '/app/train.py', './train.py', or bare 'train.py'), or URL. " "Triggers Python mode. For ML training: base this on a working example found via github_find_examples, not on internal knowledge. " + "For non-trivial scripts, submit the exact tested script source or exact tested sandbox file. " + "For GPU/model-loading training scripts, smoke-test in a GPU sandbox before submission. " + "Do not leave placeholders such as , , , or TODO. " "Mutually exclusive with 'command'." ), }, @@ -1005,7 +1193,10 @@ async def _resume_scheduled_job(self, args: Dict[str, Any]) -> ToolResult: "items": {"type": "string"}, "description": ( "Pip packages to install. Include ALL required packages. " - "Common training set: ['transformers', 'trl', 'torch', 'datasets', 'trackio', 'accelerate']. " + "Common training set: ['torch', 'transformers', 'trl', 'accelerate', 'datasets', 'trackio', 'kernels~=0.12.0']. " + "Use unpinned latest stable versions by default for the rest of the core stack; constrain kernels to kernels~=0.12.0. " + "Pin other versions only when current docs/examples require a compatibility set or a smoke test shows latest is incompatible. " + "Must include every imported third-party package and any used extras such as peft, bitsandbytes, sentencepiece, or protobuf. " "Only used with 'script'." ), }, @@ -1038,6 +1229,34 @@ async def _resume_scheduled_job(self, args: Dict[str, Any]) -> ToolResult: "type": "object", "description": "Environment variables {'KEY': 'VALUE'}. HF_TOKEN is auto-included.", }, + "trackio_space_id": { + "type": "string", + "description": ( + "Optional. The HF Space hosting the trackio dashboard for this run " + "(pattern only: '/ml-intern-<8char>'; replace , e.g. 'alice/ml-intern-a1b2c3d4'). " + "Injected as TRACKIO_SPACE_ID env var and used by the UI to embed " + "the live dashboard. Set this whenever the script uses " + "report_to='trackio'. The Space is auto-created and seeded with the " + "trackio dashboard before the job starts — DO NOT pre-create it via " + "hf_repo_git, that produces an empty Space that breaks the embed." + ), + }, + "trackio_project": { + "type": "string", + "description": ( + "Optional. The trackio project name to log this run under. " + "Injected as TRACKIO_PROJECT env var and used by the UI to filter " + "the embedded dashboard to this project." + ), + }, + "namespace": { + "type": "string", + "description": ( + "Optional namespace to run the job under. Must be the caller's own " + "account or an org they belong to. If omitted, defaults to the " + "caller's personal account. Credits are billed against this namespace." + ), + }, "job_id": { "type": "string", "description": "Job ID. Required for: logs, inspect, cancel.", @@ -1073,6 +1292,7 @@ async def log_callback(log: str): sandbox = getattr(session, "sandbox", None) if session else None if sandbox and script: from agent.tools.sandbox_tool import resolve_sandbox_script + content, error = await resolve_sandbox_script(sandbox, script) if error: return error, False @@ -1080,11 +1300,18 @@ async def log_callback(log: str): arguments = {**arguments, "script": content} hf_token = session.hf_token if session else None - namespace = os.environ.get("HF_NAMESPACE") or (HfApi(token=hf_token).whoami().get("name") if hf_token else None) + try: + namespace, jobs_access = await resolve_jobs_namespace( + hf_token or "", + arguments.get("namespace"), + ) + except JobsAccessError as e: + return str(e), False tool = HfJobsTool( namespace=namespace, hf_token=hf_token, + jobs_access=jobs_access, log_callback=log_callback if session else None, session=session, tool_call_id=tool_call_id, diff --git a/agent/tools/local_tools.py b/agent/tools/local_tools.py index fc456f682..50cd5bd65 100644 --- a/agent/tools/local_tools.py +++ b/agent/tools/local_tools.py @@ -15,6 +15,8 @@ from pathlib import Path from typing import Any +from agent.core.hub_artifacts import wrap_shell_command_with_hub_artifact_bootstrap + MAX_OUTPUT_CHARS = 25_000 MAX_LINE_LENGTH = 4000 @@ -22,7 +24,7 @@ DEFAULT_TIMEOUT = 120 MAX_TIMEOUT = 36000 # 10 hours — needed for long training runs (e.g. PostTrainBench) -_ANSI_RE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07') +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07") # Track files that have been read this session (enforces read-before-write/edit) _files_read: set[str] = set() @@ -63,17 +65,21 @@ def _atomic_write(path: Path, content: str) -> None: def _strip_ansi(text: str) -> str: - return _ANSI_RE.sub('', text) + return _ANSI_RE.sub("", text) -def _truncate_output(output: str, max_chars: int = MAX_OUTPUT_CHARS, head_ratio: float = 0.25) -> str: +def _truncate_output( + output: str, max_chars: int = MAX_OUTPUT_CHARS, head_ratio: float = 0.25 +) -> str: """Tail-biased truncation with temp file spillover for full output access.""" if len(output) <= max_chars: return output # Write full output to temp file so LLM can read specific sections spill_path = None try: - with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', prefix='bash_output_', delete=False) as f: + with tempfile.NamedTemporaryFile( + mode="w", suffix=".txt", prefix="bash_output_", delete=False + ) as f: f.write(output) spill_path = f.name except Exception: @@ -93,10 +99,14 @@ def _truncate_output(output: str, max_chars: int = MAX_OUTPUT_CHARS, head_ratio: # ── Handlers ──────────────────────────────────────────────────────────── -async def _bash_handler(args: dict[str, Any], **_kw) -> tuple[str, bool]: + +async def _bash_handler( + args: dict[str, Any], session: Any = None, **_kw +) -> tuple[str, bool]: command = args.get("command", "") if not command: return "No command provided.", False + command = wrap_shell_command_with_hub_artifact_bootstrap(command, session) work_dir = args.get("work_dir", ".") timeout = min(args.get("timeout") or DEFAULT_TIMEOUT, MAX_TIMEOUT) try: @@ -174,9 +184,12 @@ async def _write_handler(args: dict[str, Any], **_kw) -> tuple[str, bool]: # Syntax validation for Python files if p.suffix == ".py": from agent.tools.edit_utils import validate_python + warnings = validate_python(content, file_path) if warnings: - msg += "\n\nValidation warnings:\n" + "\n".join(f" ⚠ {w}" for w in warnings) + msg += "\n\nValidation warnings:\n" + "\n".join( + f" ⚠ {w}" for w in warnings + ) return msg, True except Exception as e: return f"write error: {e}", False @@ -229,7 +242,9 @@ async def _edit_handler(args: dict[str, Any], **_kw) -> tuple[str, bool]: if p.suffix == ".py": warnings = validate_python(new_text, file_path) if warnings: - msg += "\n\nValidation warnings:\n" + "\n".join(f" ⚠ {w}" for w in warnings) + msg += "\n\nValidation warnings:\n" + "\n".join( + f" ⚠ {w}" for w in warnings + ) return msg, True diff --git a/agent/tools/notify_tool.py b/agent/tools/notify_tool.py new file mode 100644 index 000000000..f926d5a58 --- /dev/null +++ b/agent/tools/notify_tool.py @@ -0,0 +1,108 @@ +from typing import Any + +from agent.messaging.models import NotificationRequest + +NOTIFY_TOOL_SPEC = { + "name": "notify", + "description": ( + "Send an out-of-band notification to configured messaging destinations. " + "Use this only when the user explicitly asked for proactive notifications " + "or when the task requires reporting progress outside the chat. " + "Destinations must be named server-side configs such as 'slack.ops'." + ), + "parameters": { + "type": "object", + "properties": { + "destinations": { + "type": "array", + "description": "Named messaging destinations to notify.", + "items": {"type": "string"}, + "minItems": 1, + }, + "message": { + "type": "string", + "description": "Main notification body.", + }, + "title": { + "type": "string", + "description": "Optional short title line.", + }, + "severity": { + "type": "string", + "enum": ["info", "success", "warning", "error"], + "description": "Notification severity label.", + }, + }, + "required": ["destinations", "message"], + }, +} + + +async def notify_handler( + arguments: dict[str, Any], session=None, **_kwargs +) -> tuple[str, bool]: + if session is None or session.notification_gateway is None: + return "Messaging is not configured for this session.", False + + raw_destinations = arguments.get("destinations", []) + if not isinstance(raw_destinations, list) or not raw_destinations: + return "destinations must be a non-empty array of destination names.", False + + destinations: list[str] = [] + seen: set[str] = set() + for raw_name in raw_destinations: + if not isinstance(raw_name, str): + return "Each destination must be a string.", False + name = raw_name.strip() + if not name: + return "Destination names must not be empty.", False + if name not in seen: + destinations.append(name) + seen.add(name) + + disallowed = [ + name + for name in destinations + if not session.config.messaging.can_agent_tool_send(name) + ] + if disallowed: + return ( + "These destinations are unavailable for the notify tool: " + + ", ".join(disallowed) + ), False + + message = arguments.get("message", "") + if not isinstance(message, str) or not message.strip(): + return "message must be a non-empty string.", False + + title = arguments.get("title") + severity = arguments.get("severity", "info") + if title is not None and not isinstance(title, str): + return "title must be a string when provided.", False + if severity not in {"info", "success", "warning", "error"}: + return "severity must be one of: info, success, warning, error.", False + + requests = [ + NotificationRequest( + destination=name, + title=title, + message=message, + severity=severity, + metadata={ + "session_id": session.session_id, + "model": session.config.model_name, + }, + ) + for name in destinations + ] + results = await session.notification_gateway.send_many(requests) + + lines = [] + all_ok = True + for result in results: + if result.ok: + lines.append(f"{result.destination}: sent") + else: + all_ok = False + lines.append(f"{result.destination}: failed ({result.error})") + return "\n".join(lines), all_ok diff --git a/agent/tools/papers_tool.py b/agent/tools/papers_tool.py index 4032a7703..ff2cf51f0 100644 --- a/agent/tools/papers_tool.py +++ b/agent/tools/papers_tool.py @@ -102,7 +102,9 @@ async def _s2_request( async def _s2_get_json( - client: httpx.AsyncClient, path: str, params: dict | None = None, + client: httpx.AsyncClient, + path: str, + params: dict | None = None, ) -> dict | None: """Cached S2 GET returning parsed JSON or None.""" key = _s2_cache_key(path, params) @@ -118,17 +120,6 @@ async def _s2_get_json( return None -async def _s2_get_paper( - client: httpx.AsyncClient, arxiv_id: str, fields: str, -) -> dict | None: - """Fetch a single paper from S2 by arxiv ID. Returns None on failure.""" - return await _s2_get_json( - client, - f"/graph/v1/paper/{_s2_paper_id(arxiv_id)}", - {"fields": fields}, - ) - - # --------------------------------------------------------------------------- # HTML paper parsing # --------------------------------------------------------------------------- @@ -322,7 +313,9 @@ def _format_paper_detail(paper: dict, s2_data: dict | None = None) -> str: if keywords: lines.append(f"**Keywords:** {', '.join(keywords)}") if s2_data and s2_data.get("s2FieldsOfStudy"): - fields = [f["category"] for f in s2_data["s2FieldsOfStudy"] if f.get("category")] + fields = [ + f["category"] for f in s2_data["s2FieldsOfStudy"] if f.get("category") + ] if fields: lines.append(f"**Fields:** {', '.join(fields)}") if s2_data and s2_data.get("venue"): @@ -393,7 +386,9 @@ def _format_datasets(datasets: list, arxiv_id: str, sort: str) -> str: ds_id = ds.get("id", "unknown") downloads = ds.get("downloads", 0) likes = ds.get("likes", 0) - desc = _truncate(_clean_description(ds.get("description") or ""), MAX_SUMMARY_LEN) + desc = _truncate( + _clean_description(ds.get("description") or ""), MAX_SUMMARY_LEN + ) tags = ds.get("tags") or [] interesting = [t for t in tags if not t.startswith(("arxiv:", "region:"))][:5] @@ -582,11 +577,15 @@ def _format_s2_paper_list(papers: list[dict], title: str) -> str: lines.append(f"**TL;DR:** {tldr}") lines.append("") - lines.append("Use paper_details with arxiv_id for full info, or read_paper to read sections.") + lines.append( + "Use paper_details with arxiv_id for full info, or read_paper to read sections." + ) return "\n".join(lines) -async def _s2_bulk_search(query: str, args: dict[str, Any], limit: int) -> ToolResult | None: +async def _s2_bulk_search( + query: str, args: dict[str, Any], limit: int +) -> ToolResult | None: """Search via S2 bulk endpoint with filters. Returns None on failure.""" params: dict[str, Any] = { "query": query, @@ -616,7 +615,9 @@ async def _s2_bulk_search(query: str, args: dict[str, Any], limit: int) -> ToolR params["sort"] = f"{sort_by}:desc" async with httpx.AsyncClient(timeout=15) as client: - resp = await _s2_request(client, "GET", "/graph/v1/paper/search/bulk", params=params) + resp = await _s2_request( + client, "GET", "/graph/v1/paper/search/bulk", params=params + ) if not resp or resp.status_code != 200: return None data = resp.json() @@ -629,7 +630,9 @@ async def _s2_bulk_search(query: str, args: dict[str, Any], limit: int) -> ToolR "resultsShared": 0, } - formatted = _format_s2_paper_list(papers[:limit], f"Papers matching '{query}' (Semantic Scholar)") + formatted = _format_s2_paper_list( + papers[:limit], f"Papers matching '{query}' (Semantic Scholar)" + ) return { "formatted": formatted, "totalResults": data.get("total", len(papers)), @@ -643,7 +646,10 @@ async def _op_search(args: dict[str, Any], limit: int) -> ToolResult: return _error("'query' is required for search operation.") # Route to S2 when filters are present - use_s2 = any(args.get(k) for k in ("date_from", "date_to", "categories", "min_citations", "sort_by")) + use_s2 = any( + args.get(k) + for k in ("date_from", "date_to", "categories", "min_citations", "sort_by") + ) if use_s2: result = await _s2_bulk_search(query, args, limit) if result is not None: @@ -806,7 +812,9 @@ def _format_citation_graph( lines.append("No citations found.") lines.append("") - lines.append("**Tip:** Use paper_details with an arxiv_id from above to explore further.") + lines.append( + "**Tip:** Use paper_details with an arxiv_id from above to explore further." + ) return "\n".join(lines) @@ -824,9 +832,13 @@ async def _op_citation_graph(args: dict[str, Any], limit: int) -> ToolResult: refs, cites = None, None coros = [] if direction in ("references", "both"): - coros.append(_s2_get_json(client, f"/graph/v1/paper/{s2_id}/references", params)) + coros.append( + _s2_get_json(client, f"/graph/v1/paper/{s2_id}/references", params) + ) if direction in ("citations", "both"): - coros.append(_s2_get_json(client, f"/graph/v1/paper/{s2_id}/citations", params)) + coros.append( + _s2_get_json(client, f"/graph/v1/paper/{s2_id}/citations", params) + ) results = await asyncio.gather(*coros, return_exceptions=True) idx = 0 @@ -841,7 +853,9 @@ async def _op_citation_graph(args: dict[str, Any], limit: int) -> ToolResult: cites = r.get("data", []) if refs is None and cites is None: - return _error(f"Could not fetch citation data for {arxiv_id}. Paper may not be indexed by Semantic Scholar.") + return _error( + f"Could not fetch citation data for {arxiv_id}. Paper may not be indexed by Semantic Scholar." + ) total = (len(refs) if refs else 0) + (len(cites) if cites else 0) return { @@ -1039,7 +1053,9 @@ def _format_snippets(snippets: list[dict], query: str) -> str: lines.append(f"> {_truncate(text, 400)}") lines.append("") - lines.append("Use paper_details or read_paper with arxiv_id to explore a paper further.") + lines.append( + "Use paper_details or read_paper with arxiv_id to explore a paper further." + ) return "\n".join(lines) @@ -1065,7 +1081,9 @@ async def _op_snippet_search(args: dict[str, Any], limit: int) -> ToolResult: params["minCitationCount"] = str(args["min_citations"]) async with httpx.AsyncClient(timeout=15) as client: - resp = await _s2_request(client, "GET", "/graph/v1/snippet/search", params=params) + resp = await _s2_request( + client, "GET", "/graph/v1/snippet/search", params=params + ) if not resp or resp.status_code != 200: return _error("Snippet search failed. Semantic Scholar may be unavailable.") data = resp.json() @@ -1102,16 +1120,28 @@ async def _op_recommend(args: dict[str, Any], limit: int) -> ToolResult: async with httpx.AsyncClient(timeout=15) as client: if positive_ids and not arxiv_id: # Multi-paper recommendations (POST, not cached) - pos = [_s2_paper_id(pid.strip()) for pid in positive_ids.split(",") if pid.strip()] + pos = [ + _s2_paper_id(pid.strip()) + for pid in positive_ids.split(",") + if pid.strip() + ] neg_raw = args.get("negative_ids", "") - neg = [_s2_paper_id(pid.strip()) for pid in neg_raw.split(",") if pid.strip()] if neg_raw else [] + neg = ( + [_s2_paper_id(pid.strip()) for pid in neg_raw.split(",") if pid.strip()] + if neg_raw + else [] + ) resp = await _s2_request( - client, "POST", "/recommendations/v1/papers/", + client, + "POST", + "/recommendations/v1/papers/", json={"positivePaperIds": pos, "negativePaperIds": neg}, params={"fields": fields, "limit": limit}, ) if not resp or resp.status_code != 200: - return _error("Recommendation request failed. Semantic Scholar may be unavailable.") + return _error( + "Recommendation request failed. Semantic Scholar may be unavailable." + ) data = resp.json() else: # Single-paper recommendations (cached) @@ -1121,7 +1151,9 @@ async def _op_recommend(args: dict[str, Any], limit: int) -> ToolResult: {"fields": fields, "limit": limit, "from": "recent"}, ) if not data: - return _error("Recommendation request failed. Semantic Scholar may be unavailable.") + return _error( + "Recommendation request failed. Semantic Scholar may be unavailable." + ) papers = data.get("recommendedPapers") or [] if not papers: diff --git a/agent/tools/plan_tool.py b/agent/tools/plan_tool.py index a923d53c2..218ec9b68 100644 --- a/agent/tools/plan_tool.py +++ b/agent/tools/plan_tool.py @@ -9,6 +9,13 @@ _current_plan: List[Dict[str, str]] = [] +def reset_current_plan() -> None: + """Clear the CLI-visible in-memory plan.""" + global _current_plan + + _current_plan = [] + + class PlanTool: """Tool for managing a list of todos with status tracking.""" @@ -54,20 +61,24 @@ async def execute(self, params: Dict[str, Any]) -> ToolResult: "isError": True, } - # Store the raw todos structure in memory - _current_plan = todos + # Store a session-scoped copy so the runtime can tell whether a + # text-only model response is trying to stop while work remains. + stored_todos = [dict(todo) for todo in todos] + _current_plan = stored_todos + if self.session is not None: + self.session.current_plan = stored_todos # Emit plan update event if session is available if self.session: await self.session.send_event( Event( event_type="plan_update", - data={"plan": todos}, + data={"plan": stored_todos}, ) ) # Format only for display using terminal_display utility - formatted_output = format_plan_tool_output(todos) + formatted_output = format_plan_tool_output(stored_todos) return { "formatted": formatted_output, diff --git a/agent/tools/private_hf_repo_tools.py b/agent/tools/private_hf_repo_tools.py deleted file mode 100644 index 090f6882b..000000000 --- a/agent/tools/private_hf_repo_tools.py +++ /dev/null @@ -1,650 +0,0 @@ -""" -Private HF Repos Tool - Manage private Hugging Face repositories - -PRIMARY USE: Store job outputs, training scripts, and logs from HF Jobs. -Since job results are ephemeral, this tool provides persistent storage in private repos. - -SECONDARY USE: Read back stored files and list repo contents. -""" - -import asyncio -from typing import Any, Dict, Literal, Optional - -from huggingface_hub import HfApi, hf_hub_download -from huggingface_hub.utils import HfHubHTTPError - -from agent.tools.types import ToolResult - -# Operation names -OperationType = Literal[ - "upload_file", "create_repo", "check_repo", "list_files", "read_file" -] - - -async def _async_call(func, *args, **kwargs): - """Wrap synchronous HfApi calls for async context.""" - return await asyncio.to_thread(func, *args, **kwargs) - - -def _build_repo_url(repo_id: str, repo_type: str = "dataset") -> str: - """Build the Hub URL for a repository.""" - type_path = "" if repo_type == "model" else f"{repo_type}s" - return f"https://huggingface.co/{type_path}/{repo_id}".replace("//", "/") - - -def _content_to_bytes(content: str | bytes) -> bytes: - """Convert string or bytes content to bytes.""" - if isinstance(content, str): - return content.encode("utf-8") - return content - - -class PrivateHfRepoTool: - """Tool for managing private Hugging Face repositories.""" - - def __init__(self, hf_token: Optional[str] = None): - self.api = HfApi(token=hf_token) - - async def execute(self, params: Dict[str, Any]) -> ToolResult: - """Execute the specified upload operation.""" - operation = params.get("operation") - args = params.get("args", {}) - - # If no operation provided, return usage instructions - if not operation: - return self._show_help() - - # Normalize operation name - operation = operation.lower() - - # Check if help is requested - if args.get("help"): - return self._show_operation_help(operation) - - try: - # Route to appropriate handler - if operation == "upload_file": - return await self._upload_file(args) - elif operation == "create_repo": - return await self._create_repo(args) - elif operation == "check_repo": - return await self._check_repo(args) - elif operation == "list_files": - return await self._list_files(args) - elif operation == "read_file": - return await self._read_file(args) - else: - return { - "formatted": f'Unknown operation: "{operation}"\n\n' - "Available operations: upload_file, create_repo, check_repo, list_files, read_file\n\n" - "Call this tool with no operation for full usage instructions.", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - except HfHubHTTPError as e: - return { - "formatted": f"API Error: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - except Exception as e: - return { - "formatted": f"Error executing {operation}: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - def _show_help(self) -> ToolResult: - """Show usage instructions when tool is called with no arguments.""" - usage_text = """# Private HF Repos Tool - -**PRIMARY USE:** Store job outputs, scripts, and logs from HF Jobs to private repos. -Since job results are ephemeral, use this tool for persistent storage. - -**SECONDARY USE:** Read back stored files and list repo contents. - -## Available Commands - -### Write Operations -- **upload_file** - Upload file content to a repository -- **create_repo** - Create a new private repository - -### Read Operations -- **list_files** - List all files in a repository -- **read_file** - Read content of a specific file from a repository -- **check_repo** - Check if a repository exists - -## Examples - -### Upload a script to a dataset repo -Call this tool with: -```json -{ - "operation": "upload_file", - "args": { - "file_content": "import pandas as pd\\nprint('Hello from HF!')", - "path_in_repo": "scripts/hello.py", - "repo_id": "my-dataset", - "repo_type": "dataset", - "create_if_missing": true, - "commit_message": "Add hello script" - } -} -``` - -### Upload logs from a job -Call this tool with: -```json -{ - "operation": "upload_file", - "args": { - "file_content": "Job started...\\nJob completed successfully!", - "path_in_repo": "jobs/job-abc123/logs.txt", - "repo_id": "job-results", - "create_if_missing": true - } -} -``` - -### Create a repository -Call this tool with: -```json -{ - "operation": "create_repo", - "args": { - "repo_id": "my-results", - "repo_type": "dataset" - } -} -``` - -### Create a Space -Call this tool with: -```json -{ - "operation": "create_repo", - "args": { - "repo_id": "my-gradio-app", - "repo_type": "space", - "space_sdk": "gradio" - } -} -``` -Note: Repositories are always created as private. For spaces, `space_sdk` is required (gradio, streamlit, static, or docker). - -### Check if a repository exists -Call this tool with: -```json -{ - "operation": "check_repo", - "args": { - "repo_id": "my-dataset", - "repo_type": "dataset" - } -} -``` - -### List files in a repository -Call this tool with: -```json -{ - "operation": "list_files", - "args": { - "repo_id": "job-results", - "repo_type": "dataset" - } -} -``` - -### Read a file from a repository -Call this tool with: -```json -{ - "operation": "read_file", - "args": { - "repo_id": "job-results", - "path_in_repo": "jobs/job-abc123/script.py", - "repo_type": "dataset" - } -} -``` - -## Repository Types - -- **dataset** (default) - For storing data, results, logs, scripts -- **model** - For ML models and related artifacts -- **space** - For Spaces and applications - -## Tips - -- **Content-based**: Pass file content directly as strings or bytes, not file paths -- **Repo ID format**: Use just the repo name (e.g., "my-dataset"). Username is automatically inferred from HF_TOKEN -- **Automatic repo creation**: Set `create_if_missing: true` to auto-create repos (requires user approval) -- **Organization**: Use path_in_repo to organize files (e.g., "jobs/job-123/script.py") -- **After jobs**: Upload job scripts and logs after compute jobs complete for reproducibility -""" - return {"formatted": usage_text, "totalResults": 1, "resultsShared": 1} - - def _show_operation_help(self, operation: str) -> ToolResult: - """Show help for a specific operation.""" - help_text = f"Help for operation: {operation}\n\nCall with appropriate arguments. Use the main help for examples." - return {"formatted": help_text, "totalResults": 1, "resultsShared": 1} - - async def _upload_file(self, args: Dict[str, Any]) -> ToolResult: - """Upload file content to a Hub repository.""" - # Validate required arguments - file_content = args.get("file_content") - path_in_repo = args.get("path_in_repo") - repo_id = args.get("repo_id") - - if not file_content: - return { - "formatted": "file_content is required", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - if not path_in_repo: - return { - "formatted": "path_in_repo is required", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - if not repo_id: - return { - "formatted": "repo_id is required", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - repo_type = args.get("repo_type", "dataset") - create_if_missing = args.get("create_if_missing", False) - - # Check if repo exists - try: - repo_exists = await _async_call( - self.api.repo_exists, repo_id=repo_id, repo_type=repo_type - ) - - # Create repo if needed - if not repo_exists and create_if_missing: - create_args = { - "repo_id": repo_id, - "repo_type": repo_type, - "private": True, - } - # Pass through space_sdk if provided (required for spaces) - if "space_sdk" in args: - create_args["space_sdk"] = args["space_sdk"] - await self._create_repo(create_args) - elif not repo_exists: - return { - "formatted": f"Repository {repo_id} does not exist. Set create_if_missing: true to create it.", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - except Exception as e: - return { - "formatted": f"Failed to check repository: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - # Convert content to bytes - file_bytes = _content_to_bytes(file_content) - - # Upload file - try: - await _async_call( - self.api.upload_file, - path_or_fileobj=file_bytes, - path_in_repo=path_in_repo, - repo_id=repo_id, - repo_type=repo_type, - commit_message=args.get("commit_message", f"Upload {path_in_repo}"), - ) - - repo_url = _build_repo_url(repo_id, repo_type) - file_url = f"{repo_url}/blob/main/{path_in_repo}" - - response = f"""✓ File uploaded successfully! - -**Repository:** {repo_id} -**File:** {path_in_repo} -**View at:** {file_url} -**Browse repo:** {repo_url}""" - - return {"formatted": response, "totalResults": 1, "resultsShared": 1} - - except Exception as e: - return { - "formatted": f"Failed to upload file: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - async def _create_repo(self, args: Dict[str, Any]) -> ToolResult: - """Create a new Hub repository.""" - repo_id = args.get("repo_id") - - if not repo_id: - return { - "formatted": "repo_id is required", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - repo_type = args.get("repo_type", "dataset") - private = True # Always create private repos - space_sdk = args.get("space_sdk") # Required if repo_type is "space" - - try: - # Check if repo already exists - repo_exists = await _async_call( - self.api.repo_exists, repo_id=repo_id, repo_type=repo_type - ) - - if repo_exists: - repo_url = _build_repo_url(repo_id, repo_type) - return { - "formatted": f"Repository {repo_id} already exists.\n**View at:** {repo_url}", - "totalResults": 1, - "resultsShared": 1, - } - - # Validate space_sdk for spaces - if repo_type == "space" and not space_sdk: - return { - "formatted": "space_sdk is required when creating a space. Valid values: gradio, streamlit, static, docker", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - # Create repository - create_kwargs = { - "repo_id": repo_id, - "repo_type": repo_type, - "private": private, - "exist_ok": True, - } - # Add space_sdk only for spaces - if repo_type == "space" and space_sdk: - create_kwargs["space_sdk"] = space_sdk - - repo_url = await _async_call(self.api.create_repo, **create_kwargs) - - response = f"""✓ Repository created successfully! - -**Repository:** {repo_id} -**Type:** {repo_type} -**Private:** Yes -**View at:** {repo_url}""" - - return {"formatted": response, "totalResults": 1, "resultsShared": 1} - - except Exception as e: - return { - "formatted": f"Failed to create repository: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - async def _check_repo(self, args: Dict[str, Any]) -> ToolResult: - """Check if a Hub repository exists.""" - repo_id = args.get("repo_id") - - if not repo_id: - return { - "formatted": "repo_id is required", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - repo_type = args.get("repo_type", "dataset") - - try: - repo_exists = await _async_call( - self.api.repo_exists, repo_id=repo_id, repo_type=repo_type - ) - - if repo_exists: - repo_url = _build_repo_url(repo_id, repo_type) - response = f"""✓ Repository exists! - -**Repository:** {repo_id} -**Type:** {repo_type} -**View at:** {repo_url}""" - else: - response = f"""Repository does not exist: {repo_id} - -To create it, call this tool with: -```json -{{ - "operation": "create_repo", - "args": {{ - "repo_id": "{repo_id}", - "repo_type": "{repo_type}" - }} -}} -```""" - - return { - "formatted": response, - "totalResults": 1 if repo_exists else 0, - "resultsShared": 1 if repo_exists else 0, - } - - except Exception as e: - return { - "formatted": f"Failed to check repository: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - async def _list_files(self, args: Dict[str, Any]) -> ToolResult: - """List all files in a Hub repository.""" - repo_id = args.get("repo_id") - - if not repo_id: - return { - "formatted": "repo_id is required", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - repo_type = args.get("repo_type", "dataset") - - try: - # List all files in the repository - files = await _async_call( - self.api.list_repo_files, repo_id=repo_id, repo_type=repo_type - ) - - if not files: - return { - "formatted": f"No files found in repository: {repo_id}", - "totalResults": 0, - "resultsShared": 0, - } - - # Format file list - file_list = "\n".join(f"- {f}" for f in sorted(files)) - repo_url = _build_repo_url(repo_id, repo_type) - - response = f"""✓ Files in repository: {repo_id} - -**Total files:** {len(files)} -**Repository URL:** {repo_url} - -**Files:** -{file_list}""" - - return { - "formatted": response, - "totalResults": len(files), - "resultsShared": len(files), - } - - except Exception as e: - return { - "formatted": f"Failed to list files: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - async def _read_file(self, args: Dict[str, Any]) -> ToolResult: - """Read content of a specific file from a Hub repository.""" - repo_id = args.get("repo_id") - path_in_repo = args.get("path_in_repo") - - if not repo_id: - return { - "formatted": "repo_id is required", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - if not path_in_repo: - return { - "formatted": "path_in_repo is required", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - repo_type = args.get("repo_type", "dataset") - - try: - # Download file to cache and read it - file_path = await _async_call( - hf_hub_download, - repo_id=repo_id, - filename=path_in_repo, - repo_type=repo_type, - token=self.api.token, - ) - - # Read file content - with open(file_path, "r", encoding="utf-8") as f: - content = f.read() - - repo_url = _build_repo_url(repo_id, repo_type) - file_url = f"{repo_url}/blob/main/{path_in_repo}" - - response = f"""✓ File read successfully! - -**Repository:** {repo_id} -**File:** {path_in_repo} -**Size:** {len(content)} characters -**View at:** {file_url} - -**Content:** -``` -{content} -```""" - - return {"formatted": response, "totalResults": 1, "resultsShared": 1} - - except UnicodeDecodeError: - # If file is binary, return size info instead - try: - with open(file_path, "rb") as f: - binary_content = f.read() - - return { - "formatted": f"File is binary ({len(binary_content)} bytes). Cannot display as text.", - "totalResults": 1, - "resultsShared": 1, - } - except Exception as e: - return { - "formatted": f"Failed to read binary file: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - except Exception as e: - return { - "formatted": f"Failed to read file: {str(e)}", - "totalResults": 0, - "resultsShared": 0, - "isError": True, - } - - -# Tool specification for agent registration -PRIVATE_HF_REPO_TOOL_SPEC = { - "name": "hf_private_repos", - "description": ( - "Manage private HF repositories - create, upload, read, list files in models/datasets/spaces. " - "⚠️ PRIMARY USE: Store job outputs persistently (job storage is EPHEMERAL - everything deleted after completion). " - "**Use when:** (1) Job completes and need to store logs/scripts/results, (2) Creating repos for training outputs, " - "(3) Reading back stored files, (4) Managing Space files, (5) Organizing job artifacts by path. " - "**Pattern:** hf_jobs (ephemeral) → hf_private_repos upload_file (persistent) → can read_file later. " - "ALWAYS pass file_content as string/bytes (✓), never file paths (✗) - this is content-based, no filesystem access. " - "**Operations:** create_repo (new private repo), upload_file (store content), read_file (retrieve content), list_files (browse), check_repo (verify exists). " - "**Critical for reliability:** Jobs lose all files after completion - use this tool to preserve important outputs. " - "Repositories created are ALWAYS private by default (good for sensitive training data/models). " - "For Spaces: must provide space_sdk ('gradio', 'streamlit', 'static', 'docker') when creating. " - "**Then:** After uploading, provide user with repository URL for viewing/sharing." - ), - "parameters": { - "type": "object", - "properties": { - "operation": { - "type": "string", - "enum": [ - "upload_file", - "create_repo", - "check_repo", - "list_files", - "read_file", - ], - "description": ( - "Operation to execute. Valid values: [upload_file, create_repo, check_repo, list_files, read_file]" - ), - }, - "args": { - "type": "object", - "description": ( - "Operation-specific arguments as a JSON object. " - "Write ops: file_content (string/bytes), path_in_repo (string), repo_id (string), " - "repo_type (dataset/model/space), create_if_missing (boolean), commit_message (string), " - "space_sdk (gradio/streamlit/static/docker - required when repo_type=space). " - "Read ops: repo_id (string), path_in_repo (for read_file), repo_type (optional)." - ), - "additionalProperties": True, - }, - }, - }, -} - - -async def private_hf_repo_handler(arguments: Dict[str, Any]) -> tuple[str, bool]: - """Handler for agent tool router.""" - try: - tool = PrivateHfRepoTool() - result = await tool.execute(arguments) - return result["formatted"], not result.get("isError", False) - except Exception as e: - return f"Error executing Private HF Repo tool: {str(e)}", False diff --git a/agent/tools/research_tool.py b/agent/tools/research_tool.py index 79692383a..3cc19b2b1 100644 --- a/agent/tools/research_tool.py +++ b/agent/tools/research_tool.py @@ -9,14 +9,22 @@ import json import logging +import time from typing import Any from litellm import Message, acompletion +from agent.core import telemetry from agent.core.doom_loop import check_for_doom_loop from agent.core.llm_params import _resolve_llm_params -from agent.core.prompt_caching import with_prompt_caching +from agent.core.model_ids import strip_huggingface_model_prefix +from agent.core.prompt_caching import ( + router_session_id_for, + with_prompt_cache_params, + with_prompt_caching, +) from agent.core.session import Event +from agent.core.yolo_budget import maybe_pause_yolo_after_spend logger = logging.getLogger(__name__) @@ -37,10 +45,56 @@ "github_find_examples", "github_list_repos", "github_read_file", + "web_search", "hf_inspect_dataset", "hf_repo_files", } + +async def _research_acompletion( + *, + session: Any, + research_model: str, + messages: list[Any], + tools: Any, + llm_params: dict[str, Any], + timeout: int, + tool_choice: str | None = None, +): + kwargs: dict[str, Any] = { + "messages": messages, + "tools": tools, + "stream": False, + "timeout": timeout, + **llm_params, + } + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice + return await acompletion(**kwargs) + + +async def _record_research_llm_call( + session: Any, + *, + research_model: str, + response: Any, + started_at: float, +) -> bool: + usage = await telemetry.record_llm_call( + session, + model=research_model, + response=response, + latency_ms=int((time.monotonic() - started_at) * 1000), + finish_reason=response.choices[0].finish_reason if response.choices else None, + kind="research", + ) + return await maybe_pause_yolo_after_spend( + session, + spend_kind="research", + observed_cost_usd=usage.get("cost_usd") if isinstance(usage, dict) else None, + ) + + RESEARCH_SYSTEM_PROMPT = """\ You are a research sub-agent for an ML engineering assistant. Your primary job: mine the literature to find the best training recipes — @@ -102,6 +156,8 @@ - `explore_hf_docs(endpoint)`: Search docs for a library. Endpoints: trl, transformers, datasets, peft, accelerate, trackio, vllm, inference-endpoints, etc. - `fetch_hf_docs(url)`: Fetch full page content from explore results - `find_hf_api(query=..., tag=...)`: Find REST API endpoints +- `web_search(query=..., allowed_domains=[...], blocked_domains=[...])`: + Search the current web when papers/docs/GitHub are not enough. ## Hub repo inspection - `hf_repo_files`: List/read files in any HF repo (model, dataset, space) @@ -215,11 +271,8 @@ def _get_research_model(main_model: str) -> str: - """Pick a cheaper model for research based on the main model.""" - if "anthropic" in main_model: - return "bedrock/us.anthropic.claude-sonnet-4-6" - # For non-Anthropic models (HF router etc.), use the same model - return main_model + """Normalize the main model id for the research sub-call.""" + return strip_huggingface_model_prefix(main_model) or main_model async def research_handler( @@ -244,13 +297,11 @@ async def research_handler( user_content = f"Context: {context}\n\n{user_content}" messages.append(Message(role="user", content=user_content)) - # Use a cheaper/faster model for research + # Use the normalized router model for research main_model = session.config.model_name research_model = _get_research_model(main_model) - # Research is a cheap sub-call — cap the main session's effort at "high" - # so a user preference of ``max``/``xhigh`` (valid for Opus 4.6/4.7) doesn't - # propagate to a Sonnet research model that may not accept those levels. - # We also haven't probed this sub-model so we don't know its ceiling. + # Research is a cheap sub-call — cap the main session's effort at "high". + # We also haven't probed this sub-call's model so we don't know its ceiling. _pref = getattr(session.config, "reasoning_effort", None) _capped = "high" if _pref in ("max", "xhigh") else _pref llm_params = _resolve_llm_params( @@ -258,6 +309,10 @@ async def research_handler( getattr(session, "hf_token", None), reasoning_effort=_capped, ) + llm_params = with_prompt_cache_params( + llm_params, + session_id=router_session_id_for(session), + ) # Get read-only tool specs from the session's tool router tool_specs = [ @@ -275,6 +330,7 @@ async def research_handler( _agent_id = tool_call_id else: import uuid + _agent_id = uuid.uuid4().hex[:8] _agent_label = "research: " + (task[:50] + "…" if len(task) > 50 else task) @@ -282,12 +338,15 @@ async def _log(text: str) -> None: """Send a progress event to the UI so it doesn't look frozen.""" try: await session.send_event( - Event(event_type="tool_log", data={ - "tool": "research", - "log": text, - "agent_id": _agent_id, - "label": _agent_label, - }) + Event( + event_type="tool_log", + data={ + "tool": "research", + "log": text, + "agent_id": _agent_id, + "label": _agent_label, + }, + ) ) except Exception: pass @@ -304,8 +363,10 @@ async def _log(text: str) -> None: # ── Doom-loop detection ── doom_prompt = check_for_doom_loop(messages) if doom_prompt: - logger.warning("Research sub-agent doom loop detected at iteration %d", _iteration) - await _log("Doom loop detected — injecting corrective prompt") + logger.warning( + "Research sub-agent repetition guard activated at iteration %d", + _iteration, + ) messages.append(Message(role="user", content=doom_prompt)) # ── Context budget: warn at 75%, hard-stop at 95% ── @@ -314,53 +375,92 @@ async def _log(text: str) -> None: "Research sub-agent hit context max (%d tokens) — forcing summary", _total_tokens, ) - await _log(f"Context limit reached ({_total_tokens} tokens) — forcing wrap-up") + await _log( + f"Context limit reached ({_total_tokens} tokens) — forcing wrap-up" + ) # Ask for a final summary with no tools - messages.append(Message( - role="user", - content=( - "[SYSTEM: CONTEXT LIMIT REACHED] You have used all available context. " - "Summarize your findings NOW. Do NOT call any more tools." - ), - )) + messages.append( + Message( + role="user", + content=( + "[SYSTEM: CONTEXT LIMIT REACHED] You have used all available context. " + "Summarize your findings NOW. Do NOT call any more tools." + ), + ) + ) try: - _msgs, _ = with_prompt_caching(messages, None, llm_params.get("model")) - response = await acompletion( - messages=_msgs, + _t0 = time.monotonic() + cached_messages, _ = with_prompt_caching(messages, None, llm_params) + response = await _research_acompletion( + session=session, + research_model=research_model, + messages=cached_messages, tools=None, # no tools — force text response - stream=False, + llm_params=llm_params, timeout=120, - **llm_params, ) + # Telemetry is best-effort; a logging blip must never mask a + # valid LLM response (the surrounding except would convert it + # to "summary call failed"). + try: + if await _record_research_llm_call( + session, + research_model=research_model, + response=response, + started_at=_t0, + ): + return ( + "Research paused because the YOLO cap was reached.", + False, + ) + except Exception as _telem_err: + logger.debug("research telemetry failed: %s", _telem_err) content = response.choices[0].message.content or "" - return content or "Research context exhausted — no summary produced.", bool(content) + return ( + content or "Research context exhausted — no summary produced.", + bool(content), + ) except Exception: return "Research context exhausted and summary call failed.", False if not _warned_context and _total_tokens >= _RESEARCH_CONTEXT_WARN: _warned_context = True await _log(f"Context at {_total_tokens} tokens — nudging to wrap up") - messages.append(Message( - role="user", - content=( - "[SYSTEM: You have used 75% of your context budget. " - "Start wrapping up: finish any critical lookups, then " - "produce your final summary within the next 1-2 iterations.]" - ), - )) + messages.append( + Message( + role="user", + content=( + "[SYSTEM: You have used 75% of your context budget. " + "Start wrapping up: finish any critical lookups, then " + "produce your final summary within the next 1-2 iterations.]" + ), + ) + ) try: - _msgs, _tools = with_prompt_caching( - messages, tool_specs if tool_specs else None, llm_params.get("model") + _t0 = time.monotonic() + cached_messages, cached_tools = with_prompt_caching( + messages, tool_specs if tool_specs else None, llm_params ) - response = await acompletion( - messages=_msgs, - tools=_tools, + response = await _research_acompletion( + session=session, + research_model=research_model, + messages=cached_messages, + tools=cached_tools, tool_choice="auto", - stream=False, + llm_params=llm_params, timeout=120, - **llm_params, ) + try: + if await _record_research_llm_call( + session, + research_model=research_model, + response=response, + started_at=_t0, + ): + return "Research paused because the YOLO cap was reached.", False + except Exception as _telem_err: + logger.debug("research telemetry failed: %s", _telem_err) except Exception as e: logger.error("Research sub-agent LLM error: %s", e) return f"Research agent LLM error: {e}", False @@ -384,11 +484,13 @@ async def _log(text: str) -> None: # LiteLLM's raw Message carries `provider_specific_fields` and # `reasoning_content`, which the HF router's OpenAI schema rejects # if we echo them back in the next request. - messages.append(Message( - role="assistant", - content=msg.content, - tool_calls=msg.tool_calls, - )) + messages.append( + Message( + role="assistant", + content=msg.content, + tool_calls=msg.tool_calls, + ) + ) for tc in msg.tool_calls: try: tool_args = json.loads(tc.function.arguments) @@ -422,7 +524,7 @@ async def _log(text: str) -> None: await _log(f"▸ {tool_name} {args_str}") output, _success = await session.tool_router.call_tool( - tool_name, tool_args, session=session + tool_name, tool_args, session=session, tool_call_id=tc.id ) _tool_uses += 1 await _log(f"tools:{_tool_uses}") @@ -443,22 +545,36 @@ async def _log(text: str) -> None: # ── Iteration limit: try to salvage findings ── await _log("Iteration limit reached — extracting summary") - messages.append(Message( - role="user", - content=( - "[SYSTEM: ITERATION LIMIT] You have reached the maximum number of research " - "iterations. Summarize ALL findings so far. Do NOT call any more tools." - ), - )) + messages.append( + Message( + role="user", + content=( + "[SYSTEM: ITERATION LIMIT] You have reached the maximum number of research " + "iterations. Summarize ALL findings so far. Do NOT call any more tools." + ), + ) + ) try: - _msgs, _ = with_prompt_caching(messages, None, llm_params.get("model")) - response = await acompletion( - messages=_msgs, + _t0 = time.monotonic() + cached_messages, _ = with_prompt_caching(messages, None, llm_params) + response = await _research_acompletion( + session=session, + research_model=research_model, + messages=cached_messages, tools=None, - stream=False, + llm_params=llm_params, timeout=120, - **llm_params, ) + try: + if await _record_research_llm_call( + session, + research_model=research_model, + response=response, + started_at=_t0, + ): + return "Research paused because the YOLO cap was reached.", False + except Exception as _telem_err: + logger.debug("research telemetry failed: %s", _telem_err) content = response.choices[0].message.content or "" if content: return content, True diff --git a/agent/tools/sandbox_client.py b/agent/tools/sandbox_client.py index 16982c76f..d031c914f 100644 --- a/agent/tools/sandbox_client.py +++ b/agent/tools/sandbox_client.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 # /// script # requires-python = ">=3.10" -# dependencies = ["huggingface_hub>=0.20.0", "httpx>=0.27.0"] +# dependencies = ["huggingface_hub>=1.12.0", "httpx>=0.27.0"] # /// """ Sandbox Tools — Agent-native primitives for HF Space dev-mode sandboxes. @@ -13,7 +13,7 @@ - Optionally deletes the Space when done Lifecycle: - sb = Sandbox.create(owner="burtenshaw") # duplicate, wait, connect + sb = Sandbox.create(owner="burtenshaw") # duplicate private Space, wait, connect sb = Sandbox.create(owner="burtenshaw", # with options hardware="t4-small", private=True, @@ -37,6 +37,7 @@ from __future__ import annotations import io +import secrets as secrets_lib import sys import time import uuid @@ -47,23 +48,23 @@ from huggingface_hub import CommitOperationAdd, HfApi TEMPLATE_SPACE = "burtenshaw/sandbox" -HARDWARE_OPTIONS = [ - "cpu-basic", - "cpu-upgrade", - "t4-small", - "t4-medium", - "a10g-small", - "a10g-large", - "a100-large", -] -OUTPUT_LIMIT = 25000 -LINE_LIMIT = 4000 DEFAULT_READ_LIMIT = 2000 DEFAULT_TIMEOUT = 240 MAX_TIMEOUT = 1200 WAIT_TIMEOUT = 600 WAIT_INTERVAL = 5 API_WAIT_TIMEOUT = 180 +CPU_BASIC_HARDWARE = "cpu-basic" + + +def _is_transient_space_visibility_error(error: Exception) -> bool: + """Return True when a newly duplicated Space is not queryable yet.""" + response = getattr(error, "response", None) + if getattr(response, "status_code", None) == 404: + return True + message = str(error) + return "Repository Not Found" in message or "404 Client Error" in message + _DOCKERFILE = """\ FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim @@ -99,8 +100,8 @@ _SANDBOX_SERVER = '''\ """Minimal FastAPI server for sandbox operations.""" -import os, subprocess, pathlib, signal, threading, re, tempfile -from fastapi import FastAPI +import hmac, os, subprocess, pathlib, signal, threading, re, tempfile +from fastapi import Depends, FastAPI, HTTPException, Request from pydantic import BaseModel from typing import Optional import uvicorn @@ -156,6 +157,24 @@ def _atomic_write(path: pathlib.Path, content: str): app = FastAPI() +def _bearer_token(header: str) -> str: + scheme, _, supplied = header.partition(" ") + if scheme.lower() != "bearer" or not supplied: + return "" + return supplied + +def _require_auth(request: Request) -> None: + sandbox_token = os.environ.get("SANDBOX_API_TOKEN") or "" + if not sandbox_token: + raise HTTPException(status_code=503, detail="Sandbox API token not configured") + supplied = _bearer_token(request.headers.get("x-sandbox-authorization", "")) + if not supplied: + raise HTTPException(status_code=401, detail="Missing bearer token") + if not hmac.compare_digest(supplied, sandbox_token): + raise HTTPException(status_code=401, detail="Invalid bearer token") + +_AUTH = [Depends(_require_auth)] + # Track active bash processes so they can be killed on cancel _active_procs = {} # pid -> subprocess.Popen _proc_lock = threading.Lock() @@ -344,7 +363,7 @@ def _validate_python(content, path=""): def health(): return {"status": "ok"} -@app.post("/api/bash") +@app.post("/api/bash", dependencies=_AUTH) def bash(req: BashReq): try: proc = subprocess.Popen( @@ -371,7 +390,7 @@ def bash(req: BashReq): except Exception as e: return {"success": False, "output": "", "error": str(e)} -@app.post("/api/kill") +@app.post("/api/kill", dependencies=_AUTH) def kill_all(): """Kill all active bash processes. Called when user cancels.""" with _proc_lock: @@ -389,7 +408,7 @@ def kill_all(): pass return {"success": True, "output": f"Killed {len(killed)} process(es): {killed}", "error": ""} -@app.post("/api/read") +@app.post("/api/read", dependencies=_AUTH) def read(req: ReadReq): try: p = pathlib.Path(req.path) @@ -406,7 +425,7 @@ def read(req: ReadReq): except Exception as e: return {"success": False, "output": "", "error": str(e)} -@app.post("/api/write") +@app.post("/api/write", dependencies=_AUTH) def write(req: WriteReq): try: p = pathlib.Path(req.path) @@ -420,7 +439,7 @@ def write(req: WriteReq): except Exception as e: return {"success": False, "output": "", "error": str(e)} -@app.post("/api/edit") +@app.post("/api/edit", dependencies=_AUTH) def edit(req: EditReq): try: p = pathlib.Path(req.path) @@ -447,7 +466,7 @@ def edit(req: EditReq): except Exception as e: return {"success": False, "output": "", "error": str(e)} -@app.post("/api/exists") +@app.post("/api/exists", dependencies=_AUTH) def exists(req: ExistsReq): return {"success": True, "output": str(pathlib.Path(req.path).exists()).lower(), "error": ""} @@ -467,9 +486,6 @@ def __str__(self): return self.output or "(no output)" return f"ERROR: {self.error}" - def to_dict(self) -> dict: - return {"success": self.success, "output": self.output, "error": self.error} - @dataclass class Sandbox: @@ -482,6 +498,7 @@ class Sandbox: space_id: str token: str | None = None + api_token: str | None = field(default=None, repr=False) work_dir: str = "/app" timeout: int = DEFAULT_TIMEOUT _owns_space: bool = field(default=False, repr=False) @@ -497,12 +514,26 @@ def __post_init__(self): self._base_url = f"https://{slug}.hf.space/api/" self._client = httpx.Client( base_url=self._base_url, - headers={"Authorization": f"Bearer {self.token}"} if self.token else {}, + headers=self._auth_headers(), timeout=httpx.Timeout(MAX_TIMEOUT, connect=30), follow_redirects=True, ) self._hf_api = HfApi(token=self.token) + def _auth_headers(self) -> dict[str, str]: + """Return headers for private HF Space access plus sandbox API auth. + + Private Spaces require the HF token in ``Authorization`` at the Hub + edge. The sandbox server requires its control-plane token in the + dedicated ``X-Sandbox-Authorization`` header. + """ + headers: dict[str, str] = {} + if self.token: + headers["Authorization"] = f"Bearer {self.token}" + if self.api_token: + headers["X-Sandbox-Authorization"] = f"Bearer {self.api_token}" + return headers + # ── Lifecycle ───────────────────────────────────────────────── class Cancelled(Exception): @@ -515,8 +546,8 @@ def create( *, name: str | None = None, template: str = TEMPLATE_SPACE, - hardware: str = "cpu-basic", - private: bool = False, + hardware: str = CPU_BASIC_HARDWARE, + private: bool = True, sleep_time: int | None = None, token: str | None = None, secrets: dict[str, str] | None = None, @@ -536,7 +567,7 @@ def create( A unique suffix is always appended. template: Source Space to duplicate (default: burtenshaw/sandbox). hardware: Hardware tier (cpu-basic, t4-small, etc.). - private: Whether the Space should be private. + private: Whether the Space should be private. Defaults to True. sleep_time: Auto-sleep after N seconds of inactivity. token: HF API token (from user's OAuth session). wait_timeout: Max seconds to wait for Space to start (default: 300). @@ -563,28 +594,47 @@ def _check_cancel(): base = name or "sandbox" suffix = uuid.uuid4().hex[:8] space_id = f"{owner}/{base}-{suffix}" + sandbox_api_token = secrets_lib.token_urlsafe(32) _log(f"Creating sandbox: {space_id} (from {template})...") kwargs = { "from_id": template, "to_id": space_id, + "repo_type": "space", "private": private, - "hardware": hardware, + "space_hardware": hardware, } if sleep_time is not None: - kwargs["sleep_time"] = sleep_time + kwargs["space_sleep_time"] = sleep_time - api.duplicate_space(**kwargs) + api.duplicate_repo(**kwargs) _log(f"Space created: https://huggingface.co/spaces/{space_id}") _check_cancel() + # ``duplicate_repo`` sends hardware and sleepTimeSeconds in the + # initial create request. Avoid a second /hardware call: deployed HF + # OAuth tokens can 401 on that endpoint for a just-created private + # Space even though duplication itself succeeded. We rely on the + # duplicate endpoint to honor sleepTimeSeconds for upgraded hardware; + # cpu-basic auto-sleep is fixed by the Hub. + _log(f"Using duplicated Space hardware: {hardware}") + if sleep_time is not None: + if hardware == CPU_BASIC_HARDWARE: + _log( + f"Requested duplicated Space sleep time: {sleep_time}s " + "(cpu-basic auto-sleep is fixed by the Hub)" + ) + else: + _log(f"Using duplicated Space sleep time: {sleep_time}s") + # Inject secrets BEFORE uploading server files (which triggers rebuild). # Secrets added after a Space is running aren't available until restart, # so they must be set before the build/start cycle. - if secrets: - for key, val in secrets.items(): + sandbox_secrets = {**(secrets or {}), "SANDBOX_API_TOKEN": sandbox_api_token} + if sandbox_secrets: + for key, val in sandbox_secrets.items(): api.add_space_secret(space_id, key, val) # Upload sandbox server and Dockerfile (triggers rebuild) @@ -597,8 +647,22 @@ def _check_cancel(): deadline = time.time() + wait_timeout while time.time() < deadline: _check_cancel() - runtime = api.get_space_runtime(space_id) + try: + runtime = api.get_space_runtime(space_id) + except Exception as e: + if _is_transient_space_visibility_error(e): + _log(" Space runtime not visible yet...") + time.sleep(WAIT_INTERVAL) + continue + raise if runtime.stage == "RUNNING": + current_hardware = runtime.hardware or getattr( + runtime, "requested_hardware", None + ) + if current_hardware != hardware: + _log(f" RUNNING on {current_hardware}; waiting for {hardware}...") + time.sleep(WAIT_INTERVAL) + continue _log(f"Space is running (hardware: {runtime.hardware})") break if runtime.stage in ("RUNTIME_ERROR", "BUILD_ERROR"): @@ -617,7 +681,12 @@ def _check_cancel(): _check_cancel() # Wait for the API server to be responsive (non-fatal) - sb = cls(space_id=space_id, token=token, _owns_space=True) + sb = cls( + space_id=space_id, + token=token, + api_token=sandbox_api_token, + _owns_space=True, + ) try: sb._wait_for_api(timeout=API_WAIT_TIMEOUT, log=_log) except TimeoutError as e: @@ -627,7 +696,9 @@ def _check_cancel(): return sb @staticmethod - def _setup_server(space_id: str, api: HfApi, *, log: Callable[[str], object] = print) -> None: + def _setup_server( + space_id: str, api: HfApi, *, log: Callable[[str], object] = print + ) -> None: """Upload embedded sandbox server + Dockerfile to the Space (single commit).""" log(f"Uploading sandbox server to {space_id}...") api.create_commit( @@ -648,17 +719,30 @@ def _setup_server(space_id: str, api: HfApi, *, log: Callable[[str], object] = p log("Server files uploaded, rebuild triggered.") @classmethod - def connect(cls, space_id: str, *, token: str | None = None) -> Sandbox: + def connect( + cls, + space_id: str, + *, + token: str | None = None, + api_token: str | None = None, + ) -> Sandbox: """ Connect to an existing running Space. Does a health check to verify the Space is reachable. """ - sb = cls(space_id=space_id, token=token, _owns_space=False) + sb = cls( + space_id=space_id, + token=token, + api_token=api_token, + _owns_space=False, + ) sb._wait_for_api(timeout=60) return sb - def _wait_for_api(self, timeout: int = API_WAIT_TIMEOUT, log: Callable[[str], object] = print): + def _wait_for_api( + self, timeout: int = API_WAIT_TIMEOUT, log: Callable[[str], object] = print + ): """Poll the health endpoint until the server responds.""" deadline = time.time() + timeout last_err = None @@ -678,26 +762,23 @@ def _wait_for_api(self, timeout: int = API_WAIT_TIMEOUT, log: Callable[[str], ob f"Last status: {last_status}, last error: {last_err}" ) - def delete(self): + def delete(self, log: Callable[[str], object] | None = None): """Delete the Space. Only works if this Sandbox created it.""" if not self._owns_space: raise RuntimeError( f"This Sandbox did not create {self.space_id}. " f"Use self._hf_api.delete_repo() directly if you're sure." ) - print(f"Deleting sandbox: {self.space_id}...") + if log: + log(f"Deleting sandbox: {self.space_id}...") self._hf_api.delete_repo(self.space_id, repo_type="space") + # Clear ownership so a second cleanup call (e.g. delete_session + + # _run_session.finally both fire) early-returns instead of retrying + # a 404 delete and emitting a spurious ERROR log. + self._owns_space = False self._client.close() - print("Deleted.") - - def pause(self): - """Pause the Space (stops billing, preserves state).""" - self._hf_api.pause_space(self.space_id) - - def restart(self): - """Restart the Space.""" - self._hf_api.restart_space(self.space_id) - self._wait_for_api() + if log: + log("Deleted.") @property def url(self) -> str: @@ -831,7 +912,12 @@ def write(self, path: str, content: str) -> ToolResult: return result def edit( - self, path: str, old_str: str, new_str: str, *, replace_all: bool = False, + self, + path: str, + old_str: str, + new_str: str, + *, + replace_all: bool = False, mode: str = "replace", ) -> ToolResult: if old_str == new_str: @@ -1022,10 +1108,6 @@ def kill_all(self) -> ToolResult: }, } - @classmethod - def tool_definitions(cls) -> list[dict]: - return [{"name": name, **spec} for name, spec in cls.TOOLS.items()] - def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: dispatch = { "bash": lambda a: self.bash( diff --git a/agent/tools/sandbox_tool.py b/agent/tools/sandbox_tool.py index 74c6a7885..a550c0180 100644 --- a/agent/tools/sandbox_tool.py +++ b/agent/tools/sandbox_tool.py @@ -2,36 +2,245 @@ Sandbox tools — expose the Sandbox client as agent tools. 5 tools total: - sandbox_create — explicit sandbox creation (requires approval) - bash, read, write, edit — operations on the sandbox + sandbox_create — create/replace sandbox for non-default hardware + bash, read, write, edit — operations on the active sandbox -If any operation tool is called without an active sandbox, -a cpu-basic sandbox is auto-created (no approval needed). +A cpu-basic sandbox is preloaded for each session. Operation tools wait for it +if startup is still in progress. """ from __future__ import annotations import asyncio +import logging +import re import threading +import uuid +import weakref +from collections.abc import Callable +from datetime import datetime, timezone from typing import Any from huggingface_hub import HfApi, SpaceHardware +from agent.core.cost_estimation import ( + DEFAULT_SANDBOX_RESERVATION_HOURS, + SPACE_PRICE_USD_PER_HOUR, + CostEstimate, +) +from agent.core.hub_artifacts import wrap_shell_command_with_hub_artifact_bootstrap from agent.core.session import Event from agent.tools.sandbox_client import Sandbox +from agent.tools.trackio_seed import ensure_trackio_dashboard + +logger = logging.getLogger(__name__) + +DEFAULT_CPU_SANDBOX_HARDWARE = "cpu-basic" + +# Match the exact suffix pattern Sandbox.create produces: "sandbox-<8 hex>". +# Used to identify orphan sandboxes from prior sessions safely (won't match +# user-renamed lookalikes). +SANDBOX_SPACE_NAME_RE = re.compile(r"^sandbox-[a-f0-9]{8}$") + +# HF Space duplication/build APIs can behave poorly when multiple private +# sandboxes are created concurrently for the same namespace. Keep session +# creation non-blocking, but serialize the actual Hub create path per owner. +_SANDBOX_CREATE_LOCKS: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, dict[str, asyncio.Lock] +] = weakref.WeakKeyDictionary() +_SANDBOX_YOLO_RENEWAL_FRACTION = 0.95 + + +def _get_sandbox_create_lock(owner: str) -> asyncio.Lock: + loop = asyncio.get_running_loop() + locks = _SANDBOX_CREATE_LOCKS.setdefault(loop, {}) + lock = locks.get(owner) + if lock is None: + lock = asyncio.Lock() + locks[owner] = lock + return lock + + +def _sandbox_window_cost_usd(hardware: str) -> float | None: + price = SPACE_PRICE_USD_PER_HOUR.get(str(hardware)) + if price is None: + return None + return round(float(price) * DEFAULT_SANDBOX_RESERVATION_HOURS, 4) + + +def _sandbox_window_estimate(hardware: str) -> CostEstimate: + cost = _sandbox_window_cost_usd(hardware) + if cost is None: + return CostEstimate( + estimated_cost_usd=None, + billable=True, + block_reason=f"No price is available for sandbox hardware '{hardware}'.", + label=hardware, + ) + return CostEstimate( + estimated_cost_usd=cost, + billable=cost > 0, + label=hardware, + ) + + +def _sandbox_yolo_renewal_delay_s() -> float: + return max( + 1.0, + DEFAULT_SANDBOX_RESERVATION_HOURS * 3600 * _SANDBOX_YOLO_RENEWAL_FRACTION, + ) + + +def _sandbox_yolo_finalized_cost_usd(session: Any) -> float: + return max(0.0, float(getattr(session, "_sandbox_yolo_finalized_cost_usd", 0.0))) + + +def _add_sandbox_yolo_finalized_cost(session: Any, amount_usd: float | None) -> None: + if amount_usd is None or amount_usd <= 0: + return + session._sandbox_yolo_finalized_cost_usd = round( + _sandbox_yolo_finalized_cost_usd(session) + float(amount_usd), + 4, + ) + + +def _cancel_sandbox_yolo_renewal(session: Any) -> None: + task = getattr(session, "_sandbox_yolo_renewal_task", None) + if task and not task.done() and task is not asyncio.current_task(): + task.cancel() + session._sandbox_yolo_renewal_task = None + + +def _start_sandbox_yolo_renewal( + session: Any, + *, + hardware: str, + reservation_id: str, +) -> None: + if hardware == DEFAULT_CPU_SANDBOX_HARDWARE: + return + _cancel_sandbox_yolo_renewal(session) + task = asyncio.create_task( + _sandbox_yolo_renewal_loop( + session, + hardware=hardware, + reservation_id=reservation_id, + ) + ) + session._sandbox_yolo_renewal_task = task + + def _log_task_error(done: asyncio.Task) -> None: + if done.cancelled(): + return + try: + done.result() + except Exception as e: + logger.warning("Sandbox YOLO renewal task failed: %s", e) + + task.add_done_callback(_log_task_error) + + +async def _sandbox_yolo_renewal_loop( + session: Any, + *, + hardware: str, + reservation_id: str, +) -> None: + from agent.core.yolo_budget import ( + reconcile_budget_reservation, + reserve_session_budget, + session_yolo_enabled, + ) + + active_reservation_id = reservation_id + while True: + await asyncio.sleep(_sandbox_yolo_renewal_delay_s()) + if ( + getattr(session, "_sandbox_yolo_reservation_id", None) + != active_reservation_id + ): + return + if not getattr(session, "sandbox", None): + return + if getattr(session, "sandbox_hardware", None) != hardware: + return + + window_cost = _sandbox_window_cost_usd(hardware) + reconcile_budget_reservation(session, active_reservation_id, window_cost) + _add_sandbox_yolo_finalized_cost(session, window_cost) + + if not session_yolo_enabled(session): + session._sandbox_yolo_reservation_id = None + return + + estimate = _sandbox_window_estimate(hardware) + next_reservation_id = f"sandbox-renew-{uuid.uuid4().hex[:10]}" + decision = reserve_session_budget( + session, + estimate, + spend_kind="sandbox", + reservation_id=next_reservation_id, + ) + if not decision.allowed: + session._sandbox_yolo_reservation_id = None + await session.send_event( + Event( + event_type="tool_log", + data={ + "tool": "sandbox", + "log": ( + "YOLO usage cap reached for the active sandbox; " + "tearing it down before the reserved budget expires." + ), + }, + ) + ) + await teardown_session_sandbox(session) + return + + active_reservation_id = ( + decision.reservation.reservation_id + if decision.reservation + else next_reservation_id + ) + session._sandbox_yolo_reservation_id = active_reservation_id + + +def _session_tool_logger( + session: Any, *, tool: str = "sandbox" +) -> Callable[[str], object] | None: + event_queue = getattr(session, "event_queue", None) + if event_queue is None: + return None + + loop = asyncio.get_running_loop() + + def _log(msg: str) -> None: + loop.call_soon_threadsafe( + event_queue.put_nowait, + Event(event_type="tool_log", data={"tool": tool, "log": msg}), + ) + + return _log def _looks_like_path(script: str) -> bool: """Return True if the script string looks like a file path (not inline code).""" - return ( + if not ( isinstance(script, str) and script.strip() == script and not any(c in script for c in "\r\n\0") - and ( - script.startswith("/") - or script.startswith("./") - or script.startswith("../") - ) + ): + return False + + if script.startswith("http://") or script.startswith("https://"): + return False + + return ( + script.startswith("/") + or script.startswith("./") + or script.startswith("../") + or (script.endswith(".py") and not any(c.isspace() for c in script)) ) @@ -62,11 +271,80 @@ async def resolve_sandbox_script( return None, f"Failed to read {script} from sandbox: {e}" +async def _seed_trackio_dashboard_safe(session: Any, space_id: str) -> None: + """Idempotently seed *space_id* with trackio dashboard files using the + session's HF token. Logs progress, swallows errors — a failed seed should + not block sandbox creation.""" + if not session or not getattr(session, "hf_token", None): + return + loop = asyncio.get_running_loop() + + def _log(msg: str) -> None: + loop.call_soon_threadsafe( + session.event_queue.put_nowait, + Event(event_type="tool_log", data={"tool": "sandbox_create", "log": msg}), + ) + + try: + await asyncio.to_thread( + ensure_trackio_dashboard, space_id, session.hf_token, _log + ) + except Exception as e: + _log(f"trackio dashboard seed failed: {e}") + + +async def _update_persisted_sandbox_fields(session: Any, **fields: Any) -> None: + """Best-effort update of sandbox metadata on the durable session record.""" + store = getattr(session, "persistence_store", None) + session_id = getattr(session, "session_id", None) + if not (store and session_id and hasattr(store, "update_session_fields")): + return + try: + await store.update_session_fields(session_id, **fields) + except Exception as e: + logger.warning("Failed to persist sandbox metadata for %s: %s", session_id, e) + + +async def _persist_active_sandbox( + session: Any, + sandbox: Sandbox, + *, + hardware: str, +) -> None: + space_id = getattr(sandbox, "space_id", None) + if not space_id: + return + owner = space_id.split("/", 1)[0] if "/" in space_id else None + await _update_persisted_sandbox_fields( + session, + sandbox_space_id=space_id, + sandbox_hardware=hardware, + sandbox_owner=owner, + sandbox_created_at=datetime.now(timezone.utc), + sandbox_status="active", + ) + + +async def _clear_persisted_sandbox(session: Any) -> None: + await _update_persisted_sandbox_fields( + session, + sandbox_space_id=None, + sandbox_hardware=None, + sandbox_owner=None, + sandbox_created_at=None, + sandbox_status="destroyed", + ) + + # ── Tool name mapping (short agent names → Sandbox client names) ────── async def _ensure_sandbox( - session: Any, hardware: str = "cpu-basic", **create_kwargs + session: Any, + hardware: str = DEFAULT_CPU_SANDBOX_HARDWARE, + extra_secrets: dict[str, str] | None = None, + cancel_event: threading.Event | None = None, + **create_kwargs, ) -> tuple[Sandbox | None, str | None]: """ Ensure a sandbox exists on the session. Auto-creates with given hardware if needed. @@ -90,6 +368,45 @@ async def _ensure_sandbox( if not owner: return None, "Could not determine HF username from token." + create_lock = _get_sandbox_create_lock(owner) + if create_lock.locked(): + await session.send_event( + Event( + event_type="tool_log", + data={ + "tool": "sandbox", + "log": "Waiting for sandbox creation slot...", + }, + ) + ) + + async with create_lock: + if getattr(session, "sandbox", None): + return session.sandbox, None + + return await _create_sandbox_locked( + session, + api=api, + owner=owner, + hardware=hardware, + extra_secrets=extra_secrets, + cancel_event=cancel_event, + **create_kwargs, + ) + + +async def _create_sandbox_locked( + session: Any, + *, + api: HfApi, + owner: str, + hardware: str, + extra_secrets: dict[str, str] | None = None, + cancel_event: threading.Event | None = None, + **create_kwargs, +) -> tuple[Sandbox | None, str | None]: + """Create the Space while the per-owner sandbox creation lock is held.""" + token = session.hf_token await session.send_event( Event( event_type="tool_log", @@ -100,19 +417,13 @@ async def _ensure_sandbox( ) ) - # Thread-safe log callback: posts tool_log events from the worker thread - loop = asyncio.get_running_loop() - - def _log(msg: str) -> None: - loop.call_soon_threadsafe( - session.event_queue.put_nowait, - Event(event_type="tool_log", data={"tool": "sandbox", "log": msg}), - ) + # Thread-safe log callback: posts tool_log events from worker threads. + _log = _session_tool_logger(session) or (lambda msg: None) # Bridge asyncio cancel event to a threading.Event for the blocking create call. # We poll session._cancelled from the main loop in a background task and set # a threading.Event that Sandbox.create checks during its polling loops. - cancel_flag = threading.Event() + cancel_flag = cancel_event or threading.Event() async def _watch_cancel(): await session._cancelled.wait() @@ -120,35 +431,55 @@ async def _watch_cancel(): watcher_task = asyncio.create_task(_watch_cancel()) + secrets: dict[str, str] = {"HF_TOKEN": token} + if extra_secrets: + secrets.update({k: v for k, v in extra_secrets.items() if v}) + + create_kwargs["private"] = True # enforce: overrides any caller-supplied value kwargs = { "owner": owner, "hardware": hardware, "token": token, - "secrets": {"HF_TOKEN": token}, + "secrets": secrets, "log": _log, "cancel_event": cancel_flag, **create_kwargs, } - if hardware != "cpu-basic": + if hardware != DEFAULT_CPU_SANDBOX_HARDWARE: kwargs["sleep_time"] = 2700 + import time as _t + + _t_start = _t.monotonic() try: sb = await asyncio.to_thread(Sandbox.create, **kwargs) except Sandbox.Cancelled: return None, "Sandbox creation cancelled by user." finally: watcher_task.cancel() - session.sandbox = sb - # Set a descriptive title (template title is inherited on duplicate) - from huggingface_hub import metadata_update + if cancel_flag.is_set(): + if getattr(sb, "_owns_space", False): + try: + await asyncio.to_thread(sb.delete, log=_log) + except Exception as e: + logger.warning( + "Failed to delete cancelled sandbox %s: %s", sb.space_id, e + ) + return None, "Sandbox creation cancelled by user." - await asyncio.to_thread( - metadata_update, - sb.space_id, - {"title": "ml-intern sandbox"}, - repo_type="space", - overwrite=True, - token=token, + session.sandbox = sb + session.sandbox_hardware = hardware + session.sandbox_preload_error = None + await _persist_active_sandbox(session, sb, hardware=hardware) + + # Telemetry: sandbox creation (infra consumption signal) + from agent.core import telemetry + + await telemetry.record_sandbox_create( + session, + sb, + hardware=hardware, + create_latency_s=int(_t.monotonic() - _t_start), ) await session.send_event( @@ -161,24 +492,194 @@ async def _watch_cancel(): return sb, None +def start_cpu_sandbox_preload(session: Any) -> asyncio.Task | None: + """Start a background ``cpu-basic`` sandbox for this session.""" + if not session or getattr(session, "sandbox", None): + return None + + existing_task = getattr(session, "sandbox_preload_task", None) + if existing_task and not existing_task.done(): + return existing_task + + cancel_event = threading.Event() + session.sandbox_preload_cancel_event = cancel_event + session.sandbox_preload_error = None + + async def _preload() -> Sandbox | None: + try: + sb, error = await _ensure_sandbox( + session, + hardware=DEFAULT_CPU_SANDBOX_HARDWARE, + cancel_event=cancel_event, + ) + if error: + session.sandbox_preload_error = error + return None + return sb + except asyncio.CancelledError: + cancel_event.set() + session.sandbox_preload_error = "Sandbox creation cancelled by user." + raise + except Exception as e: + session.sandbox_preload_error = f"Failed to create sandbox: {e}" + logger.warning("CPU sandbox preload failed: %s", e) + return None + + task = asyncio.create_task(_preload()) + session.sandbox_preload_task = task + return task + + +async def cancel_sandbox_preload(session: Any) -> None: + """Best-effort cancellation for an in-flight CPU sandbox preload.""" + cancel_event = getattr(session, "sandbox_preload_cancel_event", None) + if cancel_event is not None: + cancel_event.set() + + task = getattr(session, "sandbox_preload_task", None) + if not task or task.done(): + return + + current_task = asyncio.current_task() + if task is current_task: + return + + try: + await asyncio.wait_for(asyncio.shield(task), timeout=30) + except asyncio.TimeoutError: + logger.warning( + "Timed out waiting for CPU sandbox preload cancellation; " + "task is still live, cancelling asyncio wrapper" + ) + task.cancel() + except asyncio.CancelledError: + raise + except Exception: + pass + + +async def get_active_or_preloaded_sandbox( + session: Any, +) -> tuple[Sandbox | None, str | None]: + """Return the active sandbox, waiting for the startup preload if needed.""" + if not session: + return None, "No session available." + if getattr(session, "sandbox", None): + return session.sandbox, None + + task = getattr(session, "sandbox_preload_task", None) + if task: + try: + await asyncio.shield(task) + except asyncio.CancelledError: + raise + except Exception as e: + session.sandbox_preload_error = f"Failed to create sandbox: {e}" + + if getattr(session, "sandbox", None): + return session.sandbox, None + + preload_error = getattr(session, "sandbox_preload_error", None) + if preload_error: + return None, preload_error + + return None, "Sandbox is still starting. Please retry shortly." + + +async def teardown_session_sandbox(session: Any) -> None: + """Cancel sandbox preload and delete the active owned sandbox, if present.""" + if not session: + return + + await cancel_sandbox_preload(session) + _cancel_sandbox_yolo_renewal(session) + + sandbox = getattr(session, "sandbox", None) + session.sandbox = None + + if not sandbox: + session.sandbox_hardware = None + return + + try: + if not getattr(sandbox, "_owns_space", False): + return + + space_id = getattr(sandbox, "space_id", None) + delete_log = _session_tool_logger(session) + last_err: Exception | None = None + for attempt in range(3): + try: + logger.info( + "Deleting sandbox %s (attempt %s/3)...", + space_id, + attempt + 1, + ) + await asyncio.to_thread(sandbox.delete, log=delete_log) + from agent.core import telemetry + + usage = await telemetry.record_sandbox_destroy(session, sandbox) + from agent.core.yolo_budget import ( + adjust_session_spend, + reconcile_budget_reservation, + ) + + actual_total = ( + usage.get("estimated_cost_usd") if isinstance(usage, dict) else None + ) + finalized = _sandbox_yolo_finalized_cost_usd(session) + active_reservation_id = getattr( + session, "_sandbox_yolo_reservation_id", None + ) + actual_unfinalized = None + if actual_total is not None: + actual_unfinalized = max(0.0, float(actual_total) - finalized) + reconcile_budget_reservation( + session, + active_reservation_id, + actual_unfinalized, + allow_zero_actual=True, + ) + if active_reservation_id is None and actual_unfinalized: + adjust_session_spend(session, actual_unfinalized) + session._sandbox_yolo_reservation_id = None + session._sandbox_yolo_finalized_cost_usd = 0.0 + return + except Exception as e: + last_err = e + if attempt < 2: + await asyncio.sleep(2**attempt) + logger.error( + "Failed to delete sandbox %s after 3 attempts: %s. " + "Orphan — sweep script will pick it up.", + space_id, + last_err, + ) + finally: + session.sandbox_hardware = None + await _clear_persisted_sandbox(session) + + # ── sandbox_create tool ────────────────────────────────────────────── SANDBOX_CREATE_TOOL_SPEC = { "name": "sandbox_create", "description": ( - "Create a persistent remote Linux environment for developing and testing scripts.\n\n" - "Workflow: sandbox_create → write script → pip install → test with small run → fix errors → hf_jobs at scale.\n" - "The sandbox persists across tool calls within the session. pip install works out of the box.\n\n" - "Use this when: you need to develop, test, and iterate on scripts before launching via hf_jobs. " - "Especially for training scripts where you need to verify imports, test on a small subset, and fix errors interactively.\n\n" - "Skip this when: the task is a simple one-shot operation (status check, resource search, quick data query), " - "or the script is copied from a verified working example with minimal changes.\n\n" + "Create or replace the session sandbox when non-default hardware is needed.\n\n" + "A private cpu-basic sandbox is already started automatically for each session. " + "For normal CPU code execution, call bash/read/write/edit directly; do NOT call sandbox_create first.\n\n" + "Use sandbox_create when: you need GPU hardware, cpu-upgrade, or Trackio secrets before running code. " + "The active sandbox persists across tool calls within the session. pip install works out of the box. " + "Sandboxes are always created as private HF Spaces.\n\n" "For ML code that uses CUDA, bf16, or model loading: use GPU hardware (t4-small minimum). " "CPU sandboxes cannot run GPU code paths — your test will not catch GPU-related errors.\n\n" "Before choosing hardware, estimate your VRAM needs (models you run, training data size). Rule of thumb: bf16/fp16 ≈ 2 bytes/param, " "fp32 ≈ 4 bytes/param, plus ~20% overhead for optimizer states during training.\n" "Common picks: t4-small (16GB VRAM, fits ≤1-3B), a10g-small (24GB, ≤7B), a100-large (80GB, ≤30B). " "If the model won't fit, pick larger hardware upfront — OOM on a sandbox wastes time.\n\n" + "If you intend to run a training script in this sandbox that uses report_to='trackio', " + "pass `trackio_space_id` (e.g. '/ml-intern-<8char>') and `trackio_project` so they " + "are set as TRACKIO_SPACE_ID/TRACKIO_PROJECT secrets in the sandbox and the UI can embed the live dashboard.\n\n" "Hardware: " + ", ".join([e.value for e in SpaceHardware]) + ".\n" ), "parameters": { @@ -189,11 +690,27 @@ async def _watch_cancel(): "hardware": { "type": "string", "enum": [e.value for e in SpaceHardware], - "description": "Hardware tier for the sandbox (default: cpu-basic)", + "description": ( + "Hardware tier for the sandbox. Omit for the existing auto-started " + "cpu-basic sandbox; choose GPU/cpu-upgrade only when needed." + ), }, - "private": { - "type": "boolean", - "description": "If true, create a private Space", + "trackio_space_id": { + "type": "string", + "description": ( + "Optional. The HF Space hosting the trackio dashboard for runs in this sandbox " + "(e.g. '/ml-intern-<8char>', under YOUR HF namespace). Injected as " + "TRACKIO_SPACE_ID secret and surfaced to the UI. The Space is auto-created and " + "seeded with the trackio dashboard — DO NOT pre-create it via hf_repo_git, " + "that produces an empty Space that breaks the embed." + ), + }, + "trackio_project": { + "type": "string", + "description": ( + "Optional. The trackio project name. Injected as TRACKIO_PROJECT secret and " + "used by the UI to filter the embedded dashboard to this project." + ), }, }, }, @@ -201,35 +718,136 @@ async def _watch_cancel(): async def sandbox_create_handler( - args: dict[str, Any], session: Any = None + args: dict[str, Any], session: Any = None, tool_call_id: str | None = None ) -> tuple[str, bool]: """Handle sandbox_create tool calls.""" - # If sandbox already exists, return its info + hardware = args.get("hardware", DEFAULT_CPU_SANDBOX_HARDWARE) + trackio_space_id = args.get("trackio_space_id") or None + trackio_project = args.get("trackio_project") or None + + async def _emit_trackio_state(sb: Sandbox) -> None: + """Tell the frontend which trackio dashboard to embed for this sandbox.""" + if not (session and tool_call_id and trackio_space_id): + return + data: dict[str, Any] = { + "tool_call_id": tool_call_id, + "tool": "sandbox_create", + "state": "running", + "trackioSpaceId": trackio_space_id, + } + if trackio_project: + data["trackioProject"] = trackio_project + await session.send_event(Event(event_type="tool_state_change", data=data)) + + preload_task = getattr(session, "sandbox_preload_task", None) + if ( + session + and not getattr(session, "sandbox", None) + and preload_task + and not preload_task.done() + and hardware == DEFAULT_CPU_SANDBOX_HARDWARE + ): + sb, error = await get_active_or_preloaded_sandbox(session) + if error: + return error, False + if sb: + await _emit_trackio_state(sb) + return ( + f"Sandbox already active: {sb.space_id}\n" + f"URL: {sb.url}\n" + f"Hardware: {DEFAULT_CPU_SANDBOX_HARDWARE}\n" + f"Use bash/read/write/edit to interact with it." + ), True + + if ( + session + and not getattr(session, "sandbox", None) + and preload_task + and not preload_task.done() + and hardware != DEFAULT_CPU_SANDBOX_HARDWARE + ): + await cancel_sandbox_preload(session) + + # If sandbox already exists, return its info or replace the auto CPU sandbox if session and getattr(session, "sandbox", None): sb = session.sandbox - return ( - f"Sandbox already active: {sb.space_id}\n" - f"URL: {sb.url}\n" - f"Use bash/read/write/edit to interact with it." - ), True - - hardware = args.get("hardware", "cpu-basic") - create_kwargs = {} - if "private" in args: - create_kwargs["private"] = args["private"] + active_hardware = getattr(session, "sandbox_hardware", None) + if active_hardware == hardware: + await _emit_trackio_state(sb) + return ( + f"Sandbox already active: {sb.space_id}\n" + f"URL: {sb.url}\n" + f"Hardware: {active_hardware}\n" + f"Use bash/read/write/edit to interact with it." + ), True + + requested_hardware = args.get("hardware") + lockout_note = "" + if ( + active_hardware == DEFAULT_CPU_SANDBOX_HARDWARE + and hardware != DEFAULT_CPU_SANDBOX_HARDWARE + ): + await teardown_session_sandbox(session) + elif requested_hardware: + lockout_note = ( + f"\nRequested hardware: {requested_hardware}\n" + "Hardware cannot be changed by calling sandbox_create again. " + "Delete the existing sandbox first if you need a different tier." + ) + await _emit_trackio_state(sb) + return ( + f"Sandbox already active: {sb.space_id}\n" + f"URL: {sb.url}\n" + f"{lockout_note}\n" + f"Use bash/read/write/edit to interact with it." + ), True + else: + await _emit_trackio_state(sb) + return ( + f"Sandbox already active: {sb.space_id}\n" + f"URL: {sb.url}\n" + f"Hardware: {active_hardware or 'unknown'}\n" + f"Use bash/read/write/edit to interact with it." + ), True + + create_kwargs: dict[str, Any] = {} + + extra_secrets: dict[str, str] = {} + if trackio_space_id: + extra_secrets["TRACKIO_SPACE_ID"] = trackio_space_id + await _seed_trackio_dashboard_safe(session, trackio_space_id) + if trackio_project: + extra_secrets["TRACKIO_PROJECT"] = trackio_project try: - sb, error = await _ensure_sandbox(session, hardware=hardware, **create_kwargs) + sb, error = await _ensure_sandbox( + session, + hardware=hardware, + extra_secrets=extra_secrets or None, + **create_kwargs, + ) except Exception as e: return f"Failed to create sandbox: {e}", False if error: return error, False + if session and tool_call_id and hardware != DEFAULT_CPU_SANDBOX_HARDWARE: + session._sandbox_yolo_reservation_id = tool_call_id + session._sandbox_yolo_finalized_cost_usd = 0.0 + _start_sandbox_yolo_renewal( + session, + hardware=hardware, + reservation_id=tool_call_id, + ) + + await _emit_trackio_state(sb) + return ( f"Sandbox created: {sb.space_id}\n" f"URL: {sb.url}\n" f"Hardware: {hardware}\n" + "Visibility: private\n" f"Use bash/read/write/edit to interact with it." ), True @@ -238,13 +856,21 @@ def _make_tool_handler(sandbox_tool_name: str): """Factory: create a handler for a sandbox operation tool.""" async def handler(args: dict[str, Any], session: Any = None) -> tuple[str, bool]: - # Require sandbox to exist — user must approve sandbox_create first - if not session or not getattr(session, "sandbox", None): - return "No sandbox running. Call sandbox_create first to start one.", False - - sb = session.sandbox + sb, error = await get_active_or_preloaded_sandbox(session) + if error: + return error, False + if not sb: + return "Sandbox is still starting. Please retry shortly.", False try: + if sandbox_tool_name == "bash" and args.get("command"): + args = { + **args, + "command": wrap_shell_command_with_hub_artifact_bootstrap( + args["command"], + session, + ), + } result = await asyncio.to_thread(sb.call_tool, sandbox_tool_name, args) if result.success: output = result.output or "(no output)" @@ -267,7 +893,7 @@ def get_sandbox_tools(): tools = [] - # sandbox_create (explicit creation, requires approval) + # sandbox_create (for GPU or other non-default hardware) tools.append( ToolSpec( name=SANDBOX_CREATE_TOOL_SPEC["name"], @@ -280,10 +906,15 @@ def get_sandbox_tools(): # Operation tools (auto-execute, no approval needed) for name in Sandbox.TOOLS.keys(): spec = Sandbox.TOOLS[name] + description = ( + "Uses the session's active sandbox. A private cpu-basic sandbox is " + "started automatically for normal CPU work; call sandbox_create only " + "for GPU or other non-default hardware.\n\n" + spec["description"] + ) tools.append( ToolSpec( name=name, - description=spec["description"], + description=description, parameters=spec["parameters"], handler=_make_tool_handler(name), ) diff --git a/agent/tools/trackio_seed.py b/agent/tools/trackio_seed.py new file mode 100644 index 000000000..1062e1b5e --- /dev/null +++ b/agent/tools/trackio_seed.py @@ -0,0 +1,205 @@ +"""Seed an HF Space with the trackio dashboard. + +Background: when the agent creates a Space via `hf_repo_git create_repo` (or +the user pre-creates one), it ships with no app.py — so the iframe shows the +default Gradio "Get started" template instead of charts. Trackio's `init()` +detects the existing Space but does NOT auto-bootstrap dashboard files into it, +so the dashboard never materializes. + +This helper writes the three files trackio's runtime expects (README.md, +requirements.txt, app.py) into the Space, idempotently, BEFORE the job that +will call `trackio.init()` runs. We deliberately omit `hf_oauth: true` from +the README so the embedded iframe in ml-intern renders without a login click — +per-user privacy is enforced by namespace ownership instead. + +Beyond the dashboard files, the helper also creates the metrics bucket and +mounts it on the Space at `/data` (with `TRACKIO_DIR` / `TRACKIO_BUCKET_ID` +Space variables). Without this, the running job writes metrics into a bucket +that the dashboard Space can't read, and the iframe shows "No projects". +""" + +from __future__ import annotations + +import io +from typing import Callable, Optional + +from huggingface_hub import ( + HfApi, + Volume, + add_space_variable, + create_bucket, + create_repo, +) +from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError + + +_README = """--- +title: Trackio Dashboard +emoji: 📊 +colorFrom: pink +colorTo: gray +sdk: gradio +app_file: app.py +pinned: false +tags: + - trackio +--- + +Embedded trackio dashboard for ml-intern runs. +""" + +_REQUIREMENTS = "trackio\n" +_APP_PY = "import trackio\ntrackio.show()\n" + +# ml-intern brand mark surfaced inside the trackio dashboard. Trackio reads +# `TRACKIO_LOGO_LIGHT_URL` / `TRACKIO_LOGO_DARK_URL` from Space variables and +# renders them in place of its own logo. We point at the publicly-resolvable +# copy on the smolagents/ml-intern Space repo so any seeded dashboard inherits +# the ml-intern branding without each user having to host the asset. +_LOGO_URL = ( + "https://huggingface.co/spaces/smolagents/ml-intern/" + "resolve/main/frontend/public/smolagents.webp" +) + +_FILES = { + "README.md": _README, + "requirements.txt": _REQUIREMENTS, + "app.py": _APP_PY, +} + + +def _already_seeded(api: HfApi, space_id: str) -> bool: + """Cheap check: does the Space already have a trackio dashboard app.py? + + Avoids re-uploading the same three files on every job submission. We look + for the literal `trackio.show` call which is the load-bearing line — any + other app.py shape (the default gradio shell, a stale custom one) means + we should re-seed. + """ + try: + path = api.hf_hub_download( + repo_id=space_id, repo_type="space", filename="app.py" + ) + except (EntryNotFoundError, RepositoryNotFoundError, OSError): + return False + try: + with open(path, "r", encoding="utf-8") as f: + return "trackio.show" in f.read() + except OSError: + return False + + +def _get_space_volumes(api: HfApi, space_id: str) -> list: + """Return mounted volumes for a Space. + + `get_space_runtime()` doesn't always populate `volumes` even when the + mount exists; mirror trackio's fallback to `space_info().runtime.volumes`. + """ + runtime = api.get_space_runtime(space_id) + if getattr(runtime, "volumes", None): + return list(runtime.volumes) + info = api.space_info(space_id) + if info.runtime and getattr(info.runtime, "volumes", None): + return list(info.runtime.volumes) + return [] + + +def _ensure_bucket_mounted( + api: HfApi, + space_id: str, + bucket_id: str, + hf_token: str, + log: Optional[Callable[[str], None]] = None, +) -> None: + """Create the bucket if missing, mount it at `/data` on the Space, and + set the `TRACKIO_DIR` / `TRACKIO_BUCKET_ID` Space variables. Idempotent — + skips work that has already been done. + """ + create_bucket(bucket_id, private=True, exist_ok=True, token=hf_token) + + existing = _get_space_volumes(api, space_id) + already_mounted = any( + getattr(v, "type", None) == "bucket" + and getattr(v, "source", None) == bucket_id + and getattr(v, "mount_path", None) == "/data" + for v in existing + ) + if not already_mounted: + preserved = [ + v + for v in existing + if not ( + getattr(v, "type", None) == "bucket" + and ( + getattr(v, "source", None) == bucket_id + or getattr(v, "mount_path", None) == "/data" + ) + ) + ] + api.set_space_volumes( + space_id, + preserved + [Volume(type="bucket", source=bucket_id, mount_path="/data")], + ) + if log: + log(f"mounted bucket {bucket_id} at /data on {space_id}") + + variables = api.get_space_variables(space_id) + desired = { + "TRACKIO_DIR": "/data/trackio", + "TRACKIO_BUCKET_ID": bucket_id, + "TRACKIO_LOGO_LIGHT_URL": _LOGO_URL, + "TRACKIO_LOGO_DARK_URL": _LOGO_URL, + } + for key, value in desired.items(): + if getattr(variables.get(key), "value", None) != value: + add_space_variable(space_id, key, value, token=hf_token) + + +def ensure_trackio_dashboard( + space_id: str, + hf_token: str, + log: Optional[Callable[[str], None]] = None, +) -> bool: + """Make sure *space_id* is fully wired for trackio: + 1. Space exists with our dashboard files (README without `hf_oauth`, + `requirements.txt`, `app.py` calling `trackio.show`). + 2. Bucket `-bucket` exists, is mounted at `/data`, and the + Space has `TRACKIO_DIR` / `TRACKIO_BUCKET_ID` variables set. + + Idempotent — re-running is cheap. Returns True if any seeding happened + in step (1), False if the dashboard files were already in place. Bucket + mount is always re-checked. + """ + api = HfApi(token=hf_token) + + create_repo( + repo_id=space_id, + repo_type="space", + space_sdk="gradio", + exist_ok=True, + token=hf_token, + ) + + seeded_files = False + if _already_seeded(api, space_id): + if log: + log(f"trackio dashboard already seeded on {space_id}") + else: + if log: + log(f"seeding trackio dashboard files into {space_id}") + for path_in_repo, content in _FILES.items(): + api.upload_file( + path_or_fileobj=io.BytesIO(content.encode("utf-8")), + path_in_repo=path_in_repo, + repo_id=space_id, + repo_type="space", + commit_message=f"ml-intern: seed trackio dashboard ({path_in_repo})", + ) + seeded_files = True + + bucket_id = f"{space_id}-bucket" + _ensure_bucket_mounted(api, space_id, bucket_id, hf_token, log) + + if log: + log(f"trackio dashboard ready: https://huggingface.co/spaces/{space_id}") + return seeded_files diff --git a/agent/tools/web_search_tool.py b/agent/tools/web_search_tool.py new file mode 100644 index 000000000..5c1841085 --- /dev/null +++ b/agent/tools/web_search_tool.py @@ -0,0 +1,276 @@ +"""DuckDuckGo HTML web search tool. + +This mirrors Claw Code's Rust WebSearch behavior: fetch DuckDuckGo's HTML +endpoint, extract result links, optionally filter domains, and return a +JSON payload the model can cite. +""" + +from __future__ import annotations + +import asyncio +import html +import json +import os +import time +from dataclasses import dataclass +from html.parser import HTMLParser +from typing import Any +from urllib.parse import parse_qsl, parse_qs, urlencode, urlparse, urlunparse + +import requests + +DEFAULT_SEARCH_URL = "https://html.duckduckgo.com/html/" +WEB_SEARCH_BASE_URL_ENV = "CLAWD_WEB_SEARCH_BASE_URL" +USER_AGENT = "clawd-rust-tools/0.1" +REQUEST_TIMEOUT_SECONDS = 20 +MAX_RESULTS = 8 + + +@dataclass(frozen=True) +class SearchHit: + title: str + url: str + + def as_json(self) -> dict[str, str]: + return {"title": self.title, "url": self.url} + + +class _AnchorParser(HTMLParser): + def __init__(self, *, require_result_class: bool) -> None: + super().__init__(convert_charrefs=True) + self.require_result_class = require_result_class + self.hits: list[tuple[str, str]] = [] + self._active_href: str | None = None + self._active_text: list[str] = [] + + def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: + if tag.lower() != "a": + return + attr_map = {key.lower(): value or "" for key, value in attrs} + href = attr_map.get("href") + if not href: + return + if self.require_result_class and "result__a" not in attr_map.get("class", ""): + return + self._active_href = href + self._active_text = [] + + def handle_data(self, data: str) -> None: + if self._active_href is not None: + self._active_text.append(data) + + def handle_entityref(self, name: str) -> None: + if self._active_href is not None: + self._active_text.append(f"&{name};") + + def handle_charref(self, name: str) -> None: + if self._active_href is not None: + self._active_text.append(f"&#{name};") + + def handle_endtag(self, tag: str) -> None: + if tag.lower() != "a" or self._active_href is None: + return + title = collapse_whitespace(html.unescape("".join(self._active_text))).strip() + self.hits.append((self._active_href, title)) + self._active_href = None + self._active_text = [] + + +def build_search_url(query: str) -> str: + base = os.environ.get(WEB_SEARCH_BASE_URL_ENV, DEFAULT_SEARCH_URL) + parsed = urlparse(base) + if parsed.scheme not in {"http", "https"} or not parsed.netloc: + raise ValueError(f"invalid search base URL: {base}") + + query_pairs = parse_qsl(parsed.query, keep_blank_values=True) + query_pairs.append(("q", query)) + return urlunparse(parsed._replace(query=urlencode(query_pairs))) + + +def collapse_whitespace(value: str) -> str: + return " ".join(value.split()) + + +def decode_duckduckgo_redirect(url: str) -> str | None: + if url.startswith("http://") or url.startswith("https://"): + return html.unescape(url) + if url.startswith("//"): + joined = f"https:{url}" + elif url.startswith("/"): + joined = f"https://duckduckgo.com{url}" + else: + return None + + parsed = urlparse(joined) + if parsed.path in {"/l", "/l/"}: + uddg = parse_qs(parsed.query).get("uddg", []) + if uddg: + return html.unescape(uddg[0]) + return joined + + +def _extract_links(search_html: str, *, require_result_class: bool) -> list[SearchHit]: + parser = _AnchorParser(require_result_class=require_result_class) + parser.feed(search_html) + + hits: list[SearchHit] = [] + for raw_url, title in parser.hits: + if not title: + continue + decoded_url = decode_duckduckgo_redirect(raw_url) + if decoded_url and ( + decoded_url.startswith("http://") or decoded_url.startswith("https://") + ): + hits.append(SearchHit(title=title, url=decoded_url)) + return hits + + +def extract_search_hits(search_html: str) -> list[SearchHit]: + return _extract_links(search_html, require_result_class=True) + + +def extract_search_hits_from_generic_links(search_html: str) -> list[SearchHit]: + return _extract_links(search_html, require_result_class=False) + + +def normalize_domain_filter(domain: str) -> str: + trimmed = domain.strip() + parsed = urlparse(trimmed) + candidate = parsed.hostname if parsed.scheme and parsed.hostname else trimmed + return candidate.strip().lstrip(".").rstrip("/").lower() + + +def host_matches_list(url: str, domains: list[str]) -> bool: + host = urlparse(url).hostname + if not host: + return False + normalized_host = host.lower() + for domain in domains: + normalized = normalize_domain_filter(domain) + if normalized and ( + normalized_host == normalized or normalized_host.endswith(f".{normalized}") + ): + return True + return False + + +def dedupe_hits(hits: list[SearchHit]) -> list[SearchHit]: + seen: set[str] = set() + deduped: list[SearchHit] = [] + for hit in hits: + if hit.url in seen: + continue + seen.add(hit.url) + deduped.append(hit) + return deduped + + +def execute_web_search( + query: str, + allowed_domains: list[str] | None = None, + blocked_domains: list[str] | None = None, + tool_use_id: str = "web_search_1", +) -> dict[str, Any]: + started = time.monotonic() + search_url = build_search_url(query) + response = requests.get( + search_url, + headers={"User-Agent": USER_AGENT}, + timeout=REQUEST_TIMEOUT_SECONDS, + allow_redirects=True, + ) + + hits = extract_search_hits(response.text) + if not hits and urlparse(response.url or search_url).hostname: + hits = extract_search_hits_from_generic_links(response.text) + + if allowed_domains is not None: + hits = [hit for hit in hits if host_matches_list(hit.url, allowed_domains)] + if blocked_domains is not None: + hits = [hit for hit in hits if not host_matches_list(hit.url, blocked_domains)] + + hits = dedupe_hits(hits)[:MAX_RESULTS] + rendered_hits = "\n".join(f"- [{hit.title}]({hit.url})" for hit in hits) + if hits: + summary = ( + f"Search results for {query!r}. Include a Sources section in the final answer.\n" + f"{rendered_hits}" + ) + else: + summary = f"No web search results matched the query {query!r}." + + return { + "query": query, + "results": [ + summary, + { + "tool_use_id": tool_use_id, + "content": [hit.as_json() for hit in hits], + }, + ], + "durationSeconds": time.monotonic() - started, + } + + +WEB_SEARCH_TOOL_SPEC = { + "name": "web_search", + "description": "Search the web for current information and return cited results.", + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string", "minLength": 2}, + "allowed_domains": { + "type": "array", + "items": {"type": "string"}, + "description": "Optional allowlist of domains or URLs. Subdomains match.", + }, + "blocked_domains": { + "type": "array", + "items": {"type": "string"}, + "description": "Optional blocklist of domains or URLs. Subdomains match.", + }, + }, + "required": ["query"], + "additionalProperties": False, + }, +} + + +def _optional_string_list(arguments: dict[str, Any], key: str) -> list[str] | None: + value = arguments.get(key) + if value is None: + return None + if not isinstance(value, list) or not all(isinstance(item, str) for item in value): + raise ValueError(f"{key} must be an array of strings") + return value + + +async def web_search_handler( + arguments: dict[str, Any], + session: Any = None, + tool_call_id: str | None = None, + **_kw: Any, +) -> tuple[str, bool]: + query_value = arguments.get("query", "") + if not isinstance(query_value, str): + return ( + "Error: web_search requires a query string with at least 2 characters.", + False, + ) + + query = query_value.strip() + if len(query) < 2: + return "Error: web_search requires a query with at least 2 characters.", False + + try: + output = await asyncio.to_thread( + execute_web_search, + query=query, + allowed_domains=_optional_string_list(arguments, "allowed_domains"), + blocked_domains=_optional_string_list(arguments, "blocked_domains"), + tool_use_id=tool_call_id or "web_search_1", + ) + except Exception as exc: + return f"Error executing web search: {exc}", False + + return json.dumps(output, indent=2), True diff --git a/agent/utils/braille.py b/agent/utils/braille.py index 3b6ee4318..4621b735b 100644 --- a/agent/utils/braille.py +++ b/agent/utils/braille.py @@ -41,8 +41,7 @@ def render(self) -> list[str]: for row in range(self.term_height): offset = row * self.term_width line = "".join( - chr(0x2800 + self._buf[offset + col]) - for col in range(self.term_width) + chr(0x2800 + self._buf[offset + col]) for col in range(self.term_width) ) lines.append(line) return lines @@ -52,6 +51,7 @@ def render(self) -> list[str]: _FONT: dict[str, list[str]] = {} + def _define_font() -> None: """Define a simple 5×7 bitmap font for uppercase ASCII.""" glyphs = { @@ -113,8 +113,9 @@ def text_to_pixels(text: str, scale: int = 1) -> list[tuple[int, int]]: if cell == "#": for sy in range(scale): for sx in range(scale): - pixels.append((cursor_x + col_idx * scale + sx, - row_idx * scale + sy)) + pixels.append( + (cursor_x + col_idx * scale + sx, row_idx * scale + sy) + ) glyph_width = max(len(r) for r in glyph) cursor_x += (glyph_width + 1) * scale return pixels diff --git a/agent/utils/crt_boot.py b/agent/utils/crt_boot.py index f36ea50e8..da0867188 100644 --- a/agent/utils/crt_boot.py +++ b/agent/utils/crt_boot.py @@ -55,7 +55,10 @@ def run_boot_sequence(console: Console, boot_lines: list[tuple[str, str]]) -> No # Render previously completed lines for prev_text, prev_style in displayed_lines: if rng.random() < prev_glitch_chance: - result.append(_glitch_text(prev_text, prev_glitch_intensity, rng), style=prev_style) + result.append( + _glitch_text(prev_text, prev_glitch_intensity, rng), + style=prev_style, + ) else: result.append(prev_text, style=prev_style) result.append("\n") @@ -86,7 +89,7 @@ def run_boot_sequence(console: Console, boot_lines: list[tuple[str, str]]) -> No live.update(result) # Variable typing speed - if line_text[char_idx - 1:char_idx] in " .": + if line_text[char_idx - 1 : char_idx] in " .": time.sleep(0.025) else: time.sleep(0.010) diff --git a/agent/utils/particle_logo.py b/agent/utils/particle_logo.py index 7b2ff5a36..eb0da59bd 100644 --- a/agent/utils/particle_logo.py +++ b/agent/utils/particle_logo.py @@ -23,7 +23,9 @@ class Particle: __slots__ = ("x", "y", "target_x", "target_y", "vx", "vy", "phase", "delay") - def __init__(self, x: float, y: float, target_x: float, target_y: float, delay: float = 0): + def __init__( + self, x: float, y: float, target_x: float, target_y: float, delay: float = 0 + ): self.x = x self.y = y self.target_x = target_x @@ -57,10 +59,6 @@ def update_converge(self, t: float, strength: float = 0.08, damping: float = 0.9 self.x += self.vx self.y += self.vy - @property - def at_target(self) -> bool: - return abs(self.x - self.target_x) < 1.5 and abs(self.y - self.target_y) < 1.5 - def run_particle_logo(console: Console, hold_seconds: float = 1.5) -> None: """Run the particle coalesce effect.""" diff --git a/agent/utils/terminal_display.py b/agent/utils/terminal_display.py index 34d879108..45850e893 100644 --- a/agent/utils/terminal_display.py +++ b/agent/utils/terminal_display.py @@ -2,9 +2,11 @@ Terminal display utilities — rich-powered CLI formatting. """ +import asyncio import re from rich.console import Console +from rich.markup import escape from rich.markdown import Heading, Markdown from rich.panel import Panel from rich.theme import Theme @@ -57,23 +59,26 @@ def _clip_to_width(s: str, width: int) -> str: out.append("\033[0m…") return "".join(out) -_THEME = Theme({ - "tool.name": "bold rgb(255,200,80)", - "tool.args": "dim", - "tool.ok": "dim green", - "tool.fail": "dim red", - "info": "dim", - "muted": "dim", - # Markdown emphasis colors - "markdown.strong": "bold rgb(255,200,80)", - "markdown.emphasis": "italic rgb(180,140,40)", - "markdown.code": "rgb(120,220,255)", - "markdown.code_block": "rgb(120,220,255)", - "markdown.link": "underline rgb(90,180,255)", - "markdown.h1": "bold rgb(255,200,80)", - "markdown.h2": "bold rgb(240,180,95)", - "markdown.h3": "bold rgb(220,165,100)", -}) + +_THEME = Theme( + { + "tool.name": "bold rgb(255,200,80)", + "tool.args": "dim", + "tool.ok": "dim green", + "tool.fail": "dim red", + "info": "dim", + "muted": "dim", + # Markdown emphasis colors + "markdown.strong": "bold rgb(255,200,80)", + "markdown.emphasis": "italic rgb(180,140,40)", + "markdown.code": "rgb(120,220,255)", + "markdown.code_block": "rgb(120,220,255)", + "markdown.link": "underline rgb(90,180,255)", + "markdown.h1": "bold rgb(255,200,80)", + "markdown.h2": "bold rgb(240,180,95)", + "markdown.h3": "bold rgb(220,165,100)", + } +) _console = Console(theme=_THEME, highlight=False) @@ -87,7 +92,12 @@ def get_console() -> Console: # ── Banner ───────────────────────────────────────────────────────────── -def print_banner(model: str | None = None, hf_user: str | None = None) -> None: + +def print_banner( + model: str | None = None, + hf_user: str | None = None, + tool_runtime: str | None = None, +) -> None: """Print particle logo then CRT boot sequence with system info.""" from agent.utils.particle_logo import run_particle_logo from agent.utils.crt_boot import run_boot_sequence @@ -99,7 +109,7 @@ def print_banner(model: str | None = None, hf_user: str | None = None) -> None: _console.file.write("\033[2J\033[H") _console.file.flush() - model_label = model or "bedrock/us.anthropic.claude-opus-4-6-v1" + model_label = model or "unknown" user_label = hf_user or "not logged in" # Warm gold palette matching the shimmer highlight (255, 200, 80) @@ -110,6 +120,7 @@ def print_banner(model: str | None = None, hf_user: str | None = None) -> None: (f"{_I}Initializing agent runtime...", gold), (f"{_I} User: {user_label}", dim_gold), (f"{_I} Model: {model_label}", dim_gold), + (f"{_I} Tool runtime: {tool_runtime or 'local filesystem'}", dim_gold), (f"{_I} Tools: loading...", dim_gold), ("", ""), (f"{_I}/help for commands · /model to switch · /quit to exit", gold), @@ -120,12 +131,16 @@ def print_banner(model: str | None = None, hf_user: str | None = None) -> None: # ── Init progress ────────────────────────────────────────────────────── + def print_init_done(tool_count: int = 0) -> None: import time + f = _console.file # Overwrite the "Tools: loading..." line with actual count - f.write(f"\033[A\033[A\033[A\033[K") # Move up 3 lines (blank + help + blank) then up to tools line - f.write(f"\033[A\033[K") + f.write( + "\033[A\033[A\033[A\033[K" + ) # Move up 3 lines (blank + help + blank) then up to tools line + f.write("\033[A\033[K") gold = "\033[38;2;180;140;40m" reset = "\033[0m" tool_text = f"{_I} Tools: {tool_count} loaded" @@ -135,16 +150,22 @@ def print_init_done(tool_count: int = 0) -> None: time.sleep(0.012) f.write("\n\n") # Reprint the help line - f.write(f"{_I}\033[38;2;255;200;80m/help for commands · /model to switch · /quit to exit{reset}\n\n") + f.write( + f"{_I}\033[38;2;255;200;80m/help for commands · /model to switch · /quit to exit{reset}\n\n" + ) # Ready message — minimal padding - f.write(f"{_I}\033[38;2;255;200;80mReady. Let's build something impressive.{reset}\n") + f.write( + f"{_I}\033[38;2;255;200;80mReady. Let's build something impressive.{reset}\n" + ) f.flush() # ── Tool calls ───────────────────────────────────────────────────────── + def print_tool_call(tool_name: str, args_preview: str) -> None: import time + f = _console.file # CRT-style: type out tool name in HF yellow gold = "\033[38;2;255;200;80m" @@ -180,11 +201,10 @@ class SubAgentDisplayManager: def __init__(self): self._agents: dict[str, dict] = {} # agent_id -> state dict self._lines_on_screen = 0 - self._ticker_task = None def start(self, agent_id: str, label: str = "research") -> None: - import asyncio import time + self._agents[agent_id] = { "label": label, "calls": [], @@ -192,8 +212,6 @@ def start(self, agent_id: str, label: str = "research") -> None: "token_count": 0, "start_time": time.monotonic(), } - if not self._ticker_task: - self._ticker_task = asyncio.ensure_future(self._tick()) self._redraw() def set_tokens(self, agent_id: str, tokens: int) -> None: @@ -222,11 +240,7 @@ def clear(self, agent_id: str) -> None: _console.file.write(line + "\n") _console.file.flush() self._lines_on_screen = 0 - if not self._agents: - if self._ticker_task: - self._ticker_task.cancel() - self._ticker_task = None - else: + if self._agents: self._redraw() @staticmethod @@ -239,19 +253,10 @@ def _render_completion_line(agent: dict) -> str: line += f" \033[2m({stats})\033[0m" return line - async def _tick(self) -> None: - import asyncio - try: - while True: - await asyncio.sleep(1.0) - if self._agents: - self._redraw() - except asyncio.CancelledError: - pass - @staticmethod def _format_stats(agent: dict) -> str: import time + start = agent["start_time"] if start is None: return "" @@ -294,7 +299,7 @@ def _render_agent_lines(self, agent: dict, compact: bool = False) -> list[str]: header += f" \033[2m·\033[0m \033[2m{short}\033[0m" return [header] lines = [header] - visible = agent["calls"][-self._MAX_VISIBLE:] + visible = agent["calls"][-self._MAX_VISIBLE :] for desc in visible: lines.append(f"{_I} \033[2m{desc}\033[0m") return lines @@ -337,13 +342,14 @@ def print_tool_log(tool: str, log: str, agent_id: str = "", label: str = "") -> # ── Messages ─────────────────────────────────────────────────────────── + async def print_markdown( text: str, cancel_event: "asyncio.Event | None" = None, instant: bool = False, ) -> None: - import asyncio - import io, random + import io + import random from rich.padding import Padding _console.print() @@ -413,47 +419,113 @@ def print_interrupted() -> None: def print_compacted(old_tokens: int, new_tokens: int) -> None: - _console.print(f"{_I}[dim]context compacted: {old_tokens:,} → {new_tokens:,} tokens[/dim]") + _console.print( + f"{_I}[dim]context compacted: {old_tokens:,} → {new_tokens:,} tokens[/dim]" + ) # ── Approval ─────────────────────────────────────────────────────────── + def print_approval_header(count: int) -> None: label = f"Approval required — {count} item{'s' if count != 1 else ''}" _console.print() - _console.print(f"{_I}", Panel(f"[bold yellow]{label}[/bold yellow]", border_style="yellow", expand=False)) + _console.print( + f"{_I}", + Panel( + f"[bold yellow]{label}[/bold yellow]", border_style="yellow", expand=False + ), + ) def print_approval_item(index: int, total: int, tool_name: str, operation: str) -> None: - _console.print(f"\n{_I}[bold]\\[{index}/{total}][/bold] [tool.name]{tool_name}[/tool.name] {operation}") + _console.print( + f"\n{_I}[bold]\\[{index}/{total}][/bold] [tool.name]{tool_name}[/tool.name] {operation}" + ) def print_yolo_approve(count: int) -> None: - _console.print(f"{_I}[bold yellow]yolo →[/bold yellow] auto-approved {count} item(s)") + _console.print( + f"{_I}[bold yellow]yolo →[/bold yellow] auto-approved {count} item(s)" + ) # ── Help ─────────────────────────────────────────────────────────────── -HELP_TEXT = f"""\ -{_I}[bold]Commands[/bold] -{_I} [cyan]/help[/cyan] Show this help -{_I} [cyan]/undo[/cyan] Undo last turn -{_I} [cyan]/compact[/cyan] Compact context window -{_I} [cyan]/model[/cyan] [id] Show available models or switch -{_I} [cyan]/effort[/cyan] [level] Reasoning effort (minimal|low|medium|high|xhigh|max|off) -{_I} [cyan]/yolo[/cyan] Toggle auto-approve mode -{_I} [cyan]/status[/cyan] Current model & turn count -{_I} [cyan]/quit[/cyan] Exit""" +HELP_ROWS: tuple[tuple[str, str, str], ...] = ( + ("/help", "", "Show this help"), + ("/new", "", "Start a fresh chat"), + ("/clear", "", "Clear terminal and start fresh"), + ("/undo", "", "Undo last turn"), + ("/compact", "", "Compact context window"), + ("/resume", "[index|id|path]", "Pick up from ./session_logs"), + ("/model", "[id]", "Show available models or switch"), + ( + "/effort", + "[level]", + "Set reasoning effort preference", + ), + ("/yolo", "", "Toggle auto-approve mode"), + ("/status", "", "Current model & turn count"), + ( + "/share-traces", + "[public|private]", + "Show or change HF trace visibility", + ), + ("/quit", "", "Exit"), +) + + +def _help_column_widths( + rows: tuple[tuple[str, str, str], ...], +) -> tuple[int, int]: + return ( + max(len(command) for command, _, _ in rows), + max(len(args) for _, args, _ in rows), + ) + + +def _format_help_row( + command: str, + args: str, + description: str, + command_width: int, + args_width: int, +) -> str: + command_gap = " " * (command_width - len(command) + 2) + args_gap = " " * (args_width - len(args) + 2) + command_markup = f"[cyan]{escape(command)}[/cyan]" + args_markup = f"[muted]{escape(args)}[/muted]" if args else "" + return f"{_I} {command_markup}{command_gap}{args_markup}{args_gap}{description}" + + +def format_help_text(rows: tuple[tuple[str, str, str], ...] | None = None) -> str: + help_rows = HELP_ROWS if rows is None else rows + command_width, args_width = _help_column_widths(help_rows) + return "\n".join( + [f"{_I}[bold]Commands[/bold]"] + + [ + _format_help_row( + command, + args, + description, + command_width, + args_width, + ) + for command, args, description in help_rows + ] + ) def print_help() -> None: _console.print() - _console.print(HELP_TEXT) + _console.print(format_help_text()) _console.print() # ── Plan display ─────────────────────────────────────────────────────── + def format_plan_display() -> str: """Format the current plan for display.""" from agent.tools.plan_tool import get_current_plan @@ -487,6 +559,7 @@ def print_plan() -> None: # ── Formatting for plan_tool output (used by plan_tool handler) ──────── + def format_plan_tool_output(todos: list) -> str: if not todos: return "Plan is empty." @@ -509,6 +582,7 @@ def format_plan_tool_output(todos: list) -> str: # ── Internal helpers ─────────────────────────────────────────────────── + def _truncate(text: str, max_lines: int = 6) -> str: lines = text.split("\n") if len(lines) <= max_lines: diff --git a/backend/dataset_uploads.py b/backend/dataset_uploads.py new file mode 100644 index 000000000..94c0839e0 --- /dev/null +++ b/backend/dataset_uploads.py @@ -0,0 +1,305 @@ +"""Helpers for session-scoped dataset uploads to the Hugging Face Hub.""" + +import asyncio +import os +import re +import uuid +from dataclasses import dataclass +from urllib.parse import quote + +from fastapi import HTTPException, UploadFile +from huggingface_hub import HfApi + +MAX_DATASET_UPLOAD_BYTES = 100 * 1024 * 1024 +ALLOWED_DATASET_EXTENSIONS = {"csv", "json", "jsonl"} +_SAFE_FILENAME_RE = re.compile(r"[^A-Za-z0-9._-]+") +_SAFE_NAMESPACE_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,95}$") + + +@dataclass(frozen=True) +class DatasetUpload: + session_id: str + repo_id: str + repo_type: str + private: bool + upload_id: str + config_name: str + filename: str + original_filename: str + path_in_repo: str + size_bytes: int + format: str + hub_url: str + load_dataset_snippet: str + + def response_payload(self) -> dict[str, str | int | bool]: + return { + "session_id": self.session_id, + "repo_id": self.repo_id, + "repo_type": self.repo_type, + "private": self.private, + "upload_id": self.upload_id, + "config_name": self.config_name, + "filename": self.filename, + "path_in_repo": self.path_in_repo, + "size_bytes": self.size_bytes, + "format": self.format, + "hub_url": self.hub_url, + "load_dataset_snippet": self.load_dataset_snippet, + } + + +def sanitize_dataset_filename(filename: str | None) -> str: + """Return a Hub-safe basename while preserving the extension.""" + raw = os.path.basename(filename or "").strip() + if not raw: + raw = "dataset.csv" + + safe = _SAFE_FILENAME_RE.sub("-", raw).strip(".-_") + if not safe: + safe = "dataset.csv" + + stem, ext = os.path.splitext(safe) + if not stem: + stem = "dataset" + if not ext: + ext = ".csv" + + max_stem_len = 96 - len(ext) + stem = stem[:max_stem_len].strip(".-_") or "dataset" + return f"{stem}{ext.lower()}" + + +def display_filename(filename: str | None, fallback: str) -> str: + raw = os.path.basename(filename or "").strip() + if not raw: + return fallback + cleaned = "".join(char for char in raw if ord(char) >= 32) + return cleaned[:160] or fallback + + +def dataset_format_from_filename(filename: str) -> str: + ext = os.path.splitext(filename)[1].lower().lstrip(".") + if ext not in ALLOWED_DATASET_EXTENSIONS: + raise HTTPException( + status_code=400, + detail="Only .csv, .json, and .jsonl dataset files are supported.", + ) + return ext + + +def session_dataset_repo_id(hf_username: str | None, session_id: str) -> str: + namespace = (hf_username or "").strip() + if not namespace or not _SAFE_NAMESPACE_RE.fullmatch(namespace): + raise HTTPException( + status_code=400, + detail="Could not determine a valid Hugging Face namespace.", + ) + + safe_session_id = re.sub(r"[^A-Za-z0-9]+", "-", session_id).strip("-") + if not safe_session_id: + safe_session_id = uuid.uuid4().hex[:8] + return f"{namespace}/ml-intern-{safe_session_id[:8]}-datasets" + + +async def upload_size_bytes(upload: UploadFile) -> int: + await asyncio.to_thread(upload.file.seek, 0, os.SEEK_END) + size = await asyncio.to_thread(upload.file.tell) + await asyncio.to_thread(upload.file.seek, 0) + return int(size) + + +async def validate_dataset_upload(upload: UploadFile) -> tuple[str, str, int]: + dataset_format = dataset_format_from_filename(upload.filename or "") + safe_filename = sanitize_dataset_filename(upload.filename) + size = await upload_size_bytes(upload) + if size <= 0: + raise HTTPException(status_code=400, detail="Uploaded dataset file is empty.") + if size > MAX_DATASET_UPLOAD_BYTES: + raise HTTPException( + status_code=413, + detail="Dataset upload exceeds the 100 MB limit.", + ) + return safe_filename, dataset_format, size + + +def dataset_hub_url(repo_id: str, path_in_repo: str) -> str: + quoted_path = quote(path_in_repo, safe="/") + return f"https://huggingface.co/datasets/{repo_id}/blob/main/{quoted_path}" + + +def dataset_config_name(upload_id: str) -> str: + safe_upload_id = re.sub(r"[^A-Za-z0-9]+", "_", upload_id).strip("_").lower() + if not safe_upload_id: + safe_upload_id = "dataset" + return f"upload_{safe_upload_id[:32]}" + + +def dataset_config_name_from_path(path_in_repo: str) -> str: + parts = path_in_repo.split("/") + if len(parts) >= 3 and parts[0] == "uploads": + return dataset_config_name(parts[1]) + stem = os.path.splitext(os.path.basename(path_in_repo))[0] + return dataset_config_name(stem) + + +def is_dataset_upload_path(path_in_repo: str) -> bool: + parts = path_in_repo.split("/") + if len(parts) != 3 or parts[0] != "uploads" or not parts[1] or not parts[2]: + return False + extension = os.path.splitext(path_in_repo)[1].lower().lstrip(".") + return extension in ALLOWED_DATASET_EXTENSIONS + + +def unique_dataset_upload_paths(paths: list[str]) -> list[str]: + seen = set() + upload_paths = [] + for path in paths: + if not is_dataset_upload_path(path) or path in seen: + continue + seen.add(path) + upload_paths.append(path) + return upload_paths + + +def load_dataset_snippet(repo_id: str, config_name: str) -> str: + return ( + "from datasets import load_dataset\n\n" + f'dataset = load_dataset("{repo_id}", "{config_name}", ' + 'split="train", token=True)' + ) + + +def dataset_repo_card(repo_id: str, upload_paths: list[str]) -> bytes: + config_lines = [] + unique_upload_paths = unique_dataset_upload_paths(upload_paths) + if unique_upload_paths: + config_lines.append("configs:") + for path in unique_upload_paths: + config_lines.extend( + [ + f"- config_name: {dataset_config_name_from_path(path)}", + " data_files:", + " - split: train", + f' path: "{path}"', + ] + ) + + configs = "\n".join(config_lines) + if configs: + configs = f"{configs}\n" + + content = f"""--- +tags: +- ml-intern +- uploaded-dataset +{configs}--- + +# {repo_id} + +Private dataset files uploaded through ML Intern. + +Files are stored under `uploads//` and are attached to the +corresponding ML Intern session context by Hub reference, not by copying file +contents into the chat. + +Each uploaded file is exposed as its own dataset config so files with different +schemas can coexist in the same session repo. +""" + return content.encode("utf-8") + + +def dataset_context_note(upload: DatasetUpload) -> str: + return f"""[SYSTEM: The user uploaded a dataset file for this session. + +Use this Hugging Face Hub dataset reference when the task needs the uploaded data. +Do not look for the uploaded file on local disk and do not ask the user to +upload it again unless this Hub reference fails. + +- Repo ID: {upload.repo_id} +- Repo type: dataset +- Dataset config: {upload.config_name} +- File in repo: {upload.path_in_repo} +- Original filename: {upload.original_filename} +- Stored filename: {upload.filename} +- Format: {upload.format} +- Size: {upload.size_bytes} bytes +- Hub URL: {upload.hub_url} + +Load it with: +```python +{upload.load_dataset_snippet} +``` +]""" + + +async def push_dataset_upload_to_hub( + *, + upload: UploadFile, + session_id: str, + hf_username: str, + hf_token: str, +) -> DatasetUpload: + safe_filename, dataset_format, size = await validate_dataset_upload(upload) + original_filename = display_filename(upload.filename, safe_filename) + upload_id = uuid.uuid4().hex[:12] + config_name = dataset_config_name(upload_id) + repo_id = session_dataset_repo_id(hf_username, session_id) + path_in_repo = f"uploads/{upload_id}/{safe_filename}" + hub_url = dataset_hub_url(repo_id, path_in_repo) + snippet = load_dataset_snippet(repo_id, config_name) + api = HfApi(token=hf_token) + + await asyncio.to_thread( + api.create_repo, + repo_id=repo_id, + repo_type="dataset", + private=True, + exist_ok=True, + ) + await asyncio.to_thread( + api.update_repo_settings, + repo_id=repo_id, + repo_type="dataset", + private=True, + ) + repo_files = await asyncio.to_thread( + api.list_repo_files, + repo_id=repo_id, + repo_type="dataset", + ) + upload_paths = unique_dataset_upload_paths([*repo_files, path_in_repo]) + await asyncio.to_thread(upload.file.seek, 0) + file_bytes = await asyncio.to_thread(upload.file.read) + await asyncio.to_thread( + api.upload_file, + path_or_fileobj=file_bytes, + path_in_repo=path_in_repo, + repo_id=repo_id, + repo_type="dataset", + commit_message=f"Upload dataset file {safe_filename}", + ) + await asyncio.to_thread( + api.upload_file, + path_or_fileobj=dataset_repo_card(repo_id, upload_paths), + path_in_repo="README.md", + repo_id=repo_id, + repo_type="dataset", + commit_message="Update ML Intern dataset upload configs", + ) + + return DatasetUpload( + session_id=session_id, + repo_id=repo_id, + repo_type="dataset", + private=True, + upload_id=upload_id, + config_name=config_name, + filename=safe_filename, + original_filename=original_filename, + path_in_repo=path_in_repo, + size_bytes=size, + format=dataset_format, + hub_url=hub_url, + load_dataset_snippet=snippet, + ) diff --git a/backend/dependencies.py b/backend/dependencies.py index 97a4e2860..5a71123bc 100644 --- a/backend/dependencies.py +++ b/backend/dependencies.py @@ -7,36 +7,89 @@ import logging import os import time +from collections.abc import Iterable +from hashlib import sha256 from typing import Any import httpx from fastapi import HTTPException, Request, status +from agent.core.hf_tokens import bearer_token_from_header, clean_hf_token + +from agent.core.hf_access import fetch_whoami_v2, normalize_hf_user_plan + logger = logging.getLogger(__name__) OPENID_PROVIDER_URL = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co") AUTH_ENABLED = bool(os.environ.get("OAUTH_CLIENT_ID", "")) -HF_EMPLOYEE_ORG = os.environ.get("HF_EMPLOYEE_ORG", "huggingface") # Simple in-memory token cache: token -> (user_info, expiry_time) _token_cache: dict[str, tuple[dict[str, Any], float]] = {} TOKEN_CACHE_TTL = 300 # 5 minutes -# Org membership cache: key -> expiry_time (only caches positive results) -_org_member_cache: dict[str, float] = {} - DEV_USER: dict[str, Any] = { "user_id": "dev", "username": "dev", "authenticated": True, - "plan": "org", # Dev runs at the Pro/Org quota tier so local testing isn't capped. + "plan": "pro", } -# Plan field discovery — log the whoami-v2 shape once at DEBUG so we can -# confirm the actual key in production without hammering the HF API. +INTERNAL_HF_TOKEN_KEY = "_hf_token" +OAUTH_SCOPE_COOKIE = "hf_oauth_scope_hash" +REQUIRED_OAUTH_SCOPES: tuple[str, ...] = ( + "openid", + "profile", + "read-billing", + "read-repos", + "write-repos", + "contribute-repos", + "manage-repos", + "write-collections", + "inference-api", + "jobs", + "write-discussions", +) + +# Log the whoami-v2 shape once at DEBUG so we can confirm the production Pro +# signal without hammering the HF API. _WHOAMI_SHAPE_LOGGED = False +def normalize_oauth_scopes(scopes: Iterable[str]) -> tuple[str, ...]: + """Return stable, de-duplicated OAuth scopes preserving declaration order.""" + seen: set[str] = set() + normalized: list[str] = [] + for scope in scopes: + value = str(scope).strip() + if not value or value in seen: + continue + seen.add(value) + normalized.append(value) + return tuple(normalized) + + +def configured_oauth_scopes() -> tuple[str, ...]: + """Return the scopes this backend should request from HF OAuth. + + Spaces expose README ``hf_oauth_scopes`` through ``OAUTH_SCOPES``. Unioning + that value with the app-required scopes keeps the local request and Space + metadata in sync while ensuring new required scopes are never omitted. + """ + env_scopes = os.environ.get("OAUTH_SCOPES", "").split() + return normalize_oauth_scopes((*env_scopes, *REQUIRED_OAUTH_SCOPES)) + + +def oauth_scope_fingerprint(scopes: Iterable[str] | None = None) -> str: + """Return a non-secret fingerprint for the current OAuth scope contract.""" + scope_list = configured_oauth_scopes() if scopes is None else scopes + payload = " ".join(sorted(normalize_oauth_scopes(scope_list))) + return sha256(payload.encode("utf-8")).hexdigest()[:16] + + +def _cookie_has_current_oauth_scope_marker(request: Request) -> bool: + return request.cookies.get(OAUTH_SCOPE_COOKIE) == oauth_scope_fingerprint() + + async def _validate_token(token: str) -> dict[str, Any] | None: """Validate a token against HF OAuth userinfo endpoint. @@ -80,76 +133,34 @@ def _user_from_info(user_info: dict[str, Any]) -> dict[str, Any]: } -def _normalize_plan(whoami: dict[str, Any]) -> str: - """Map an HF /api/whoami-v2 payload to one of: 'free' | 'pro' | 'org'. - - The exact field shape in whoami-v2 isn't documented for our purposes, - so we try a handful of likely keys and fall back to 'free'. The first - call logs the raw shape at DEBUG (see `_fetch_user_plan`) so we can - pin the real key post-deploy. - """ - plan_str = "" - for key in ("plan", "type", "accountType"): - val = whoami.get(key) - if isinstance(val, str) and val: - plan_str = val.lower() - break - - if not plan_str: - if whoami.get("isPro") is True or whoami.get("is_pro") is True: - return "pro" - - if "pro" in plan_str or "enterprise" in plan_str or "team" in plan_str: - return "pro" - - # Org tier: anyone in a paid / enterprise org. We don't pay for this - # right now, but the "pro" cap applies identically. - orgs = whoami.get("orgs") or [] - if isinstance(orgs, list): - for org in orgs: - if isinstance(org, dict): - org_plan = str(org.get("plan") or org.get("type") or "").lower() - if "pro" in org_plan or "enterprise" in org_plan or "team" in org_plan: - return "org" - - return "free" +def _normalize_user_plan(whoami: Any) -> str: + """Normalize a whoami-v2 payload to the app's supported plan tiers.""" + return normalize_hf_user_plan(whoami) or "free" async def _fetch_user_plan(token: str) -> str: """Look up the user's HF plan via /api/whoami-v2. - Returns 'free' | 'pro' | 'org'. Non-200, network errors, or an unknown - payload shape all collapse to 'free' — safe default; we'd rather under- - grant the Pro cap than over-grant it on bad data. + Returns 'free' | 'pro'. Non-200, network errors, or an unknown + payload shape all collapse to 'free' — safe default; we'd rather avoid + selecting the Pro default on bad data. """ global _WHOAMI_SHAPE_LOGGED - async with httpx.AsyncClient(timeout=5.0) as client: - try: - resp = await client.get( - f"{OPENID_PROVIDER_URL}/api/whoami-v2", - headers={"Authorization": f"Bearer {token}"}, - ) - if resp.status_code != 200: - return "free" - whoami = resp.json() - except httpx.HTTPError: - return "free" - except ValueError: - return "free" + whoami = await fetch_whoami_v2(token) + if whoami is None: + return "free" if not _WHOAMI_SHAPE_LOGGED: _WHOAMI_SHAPE_LOGGED = True logger.debug( - "whoami-v2 payload keys: %s (sample values: plan=%r type=%r isPro=%r)", - sorted(whoami.keys()) if isinstance(whoami, dict) else type(whoami).__name__, - whoami.get("plan") if isinstance(whoami, dict) else None, - whoami.get("type") if isinstance(whoami, dict) else None, + "whoami-v2 payload keys: %s (sample values: isPro=%r)", + sorted(whoami.keys()) + if isinstance(whoami, dict) + else type(whoami).__name__, whoami.get("isPro") if isinstance(whoami, dict) else None, ) - if not isinstance(whoami, dict): - return "free" - return _normalize_plan(whoami) + return _normalize_user_plan(whoami) async def _extract_user_from_token(token: str) -> dict[str, Any] | None: @@ -159,32 +170,41 @@ async def _extract_user_from_token(token: str) -> dict[str, Any] | None: return None user = _user_from_info(user_info) user["plan"] = await _fetch_user_plan(token) + user[INTERNAL_HF_TOKEN_KEY] = clean_hf_token(token) return user -async def check_org_membership(token: str, org_name: str) -> bool: - """Check if the token owner belongs to an HF org. Only caches positive results.""" - now = time.time() - key = token + org_name - cached = _org_member_cache.get(key) - if cached and cached > now: - return True +async def _dev_user_from_env() -> dict[str, Any]: + """Use HF_TOKEN as the dev identity when available. - async with httpx.AsyncClient(timeout=10.0) as client: - try: - resp = await client.get( - f"{OPENID_PROVIDER_URL}/api/whoami-v2", - headers={"Authorization": f"Bearer {token}"}, - ) - if resp.status_code != 200: - return False - orgs = {o.get("name") for o in resp.json().get("orgs", [])} - if org_name in orgs: - _org_member_cache[key] = now + TOKEN_CACHE_TTL - return True - return False - except httpx.HTTPError: - return False + Local dev often runs without OAuth, but session trace uploads still need a + real HF namespace. Deriving the dev user from HF_TOKEN keeps local uploads + pointed at the token owner's dataset instead of dev/ml-intern-sessions. + """ + token = clean_hf_token(os.environ.get("HF_TOKEN", "")) + if not token: + return dict(DEV_USER) + + whoami = await fetch_whoami_v2(token) + if not isinstance(whoami, dict): + return dict(DEV_USER) + + username = None + for key in ("name", "user", "preferred_username"): + value = whoami.get(key) + if isinstance(value, str) and value: + username = value + break + if not username: + return dict(DEV_USER) + + return { + "user_id": username, + "username": username, + "authenticated": True, + "plan": await _fetch_user_plan(token), + INTERNAL_HF_TOKEN_KEY: token, + } async def get_current_user(request: Request) -> dict[str, Any]: @@ -194,15 +214,15 @@ async def get_current_user(request: Request) -> dict[str, Any]: 1. Authorization: Bearer header 2. hf_access_token cookie - In dev mode (AUTH_ENABLED=False), returns a default dev user. + In dev mode (AUTH_ENABLED=False), uses HF_TOKEN as the user when possible. """ if not AUTH_ENABLED: - return DEV_USER + return await _dev_user_from_env() - # Try Authorization header - auth_header = request.headers.get("Authorization", "") - if auth_header.startswith("Bearer "): - token = auth_header[7:] + # Bearer callers manage token lifecycle themselves; only browser cookie + # auth is forced through the scope-freshness marker below. + token = bearer_token_from_header(request.headers.get("Authorization", "")) + if token: user = await _extract_user_from_token(token) if user: return user @@ -210,6 +230,15 @@ async def get_current_user(request: Request) -> dict[str, Any]: # Try cookie token = request.cookies.get("hf_access_token") if token: + if not _cookie_has_current_oauth_scope_marker(request): + logger.info( + "Rejecting stale HF OAuth cookie; current scopes require refresh." + ) + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Authentication scopes changed. Please log in again.", + headers={"WWW-Authenticate": "Bearer"}, + ) user = await _extract_user_from_token(token) if user: return user @@ -219,31 +248,3 @@ async def get_current_user(request: Request) -> dict[str, Any]: detail="Not authenticated. Please log in via /auth/login.", headers={"WWW-Authenticate": "Bearer"}, ) - - -def _extract_token(request: Request) -> str | None: - """Pull the HF access token from the Authorization header or cookie. - - Mirrors the lookup order used by ``get_current_user``. - """ - auth_header = request.headers.get("Authorization", "") - if auth_header.startswith("Bearer "): - return auth_header[7:] - return request.cookies.get("hf_access_token") - - -async def require_huggingface_org_member(request: Request) -> bool: - """Return True if the caller is a member of the ``huggingface`` org. - - Used to gate endpoints that can push a session onto an Anthropic model - billed to the Space's ``ANTHROPIC_API_KEY``. Returns True unconditionally - in dev mode so local testing isn't blocked. - """ - if not AUTH_ENABLED: - return True - token = _extract_token(request) - if not token: - return False - return await check_org_membership(token, HF_EMPLOYEE_ORG) - - diff --git a/backend/kpis_scheduler.py b/backend/kpis_scheduler.py new file mode 100644 index 000000000..9b2199c69 --- /dev/null +++ b/backend/kpis_scheduler.py @@ -0,0 +1,148 @@ +"""In-process hourly KPI rollup, owned by the backend Space lifespan. + +Replaces an external GitHub Actions cron so the rollup lives next to the data +and reuses the Space's existing HF token — no production secrets on the +public source repo. See ``scripts/build_kpis.py`` for the data-flow diagram +and metric definitions. + +Behaviour:: + + lifespan startup → start APScheduler with cron("5 * * * *", UTC) + → fire a best-effort 6-hour backfill (fire-and-forget) + each :05 → run ``build_kpis.run_for_hour`` for the just-completed hour + lifespan shutdown → scheduler.shutdown(wait=False) + +Environment:: + + HF_KPI_WRITE_TOKEN | HF_SESSION_UPLOAD_TOKEN | HF_TOKEN | HF_ADMIN_TOKEN + First one found is used. Least-privilege first. + KPI_SOURCE_REPO default smolagents/ml-intern-sessions + KPI_TARGET_REPO default smolagents/ml-intern-kpis + ML_INTERN_KPIS_DISABLED if truthy, the scheduler is not started +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import logging +import os +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Optional + +logger = logging.getLogger(__name__) + +_PROJECT_ROOT = Path(__file__).resolve().parent.parent + +# Hold strong refs to backfill tasks so asyncio doesn't GC them mid-run. +_background_tasks: set[asyncio.Task] = set() + +_scheduler = None # AsyncIOScheduler instance (lazy import) + + +def _resolve_token() -> Optional[str]: + """Pick the first available HF token. Least-privilege first.""" + for var in ( + "HF_KPI_WRITE_TOKEN", + "HF_SESSION_UPLOAD_TOKEN", + "HF_TOKEN", + "HF_ADMIN_TOKEN", + ): + val = os.environ.get(var) + if val: + return val + return None + + +def _load_build_kpis(): + """Import ``scripts/build_kpis.py`` without putting ``scripts/`` on sys.path.""" + spec = importlib.util.spec_from_file_location( + "build_kpis", + _PROJECT_ROOT / "scripts" / "build_kpis.py", + ) + mod = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(mod) + return mod + + +async def _run_hour(hour_dt: datetime) -> None: + """Run one hourly rollup off the event loop. Best-effort, never raises.""" + token = _resolve_token() + if not token: + logger.warning("kpis_scheduler: no HF token available, skipping %s", hour_dt) + return + try: + mod = _load_build_kpis() + from huggingface_hub import HfApi + + api = HfApi() + source = os.environ.get("KPI_SOURCE_REPO", "smolagents/ml-intern-sessions") + target = os.environ.get("KPI_TARGET_REPO", "smolagents/ml-intern-kpis") + await asyncio.to_thread(mod.run_for_hour, api, source, target, hour_dt, token) + except Exception as e: + logger.warning("kpis_scheduler: rollup for %s failed: %s", hour_dt, e) + + +async def run_last_completed_hour() -> None: + """The scheduled-at-:05 job. Rolls up the previous whole hour.""" + now = datetime.now(timezone.utc).replace(minute=0, second=0, microsecond=0) + await _run_hour(now - timedelta(hours=1)) + + +async def backfill(hours: int = 6) -> None: + """Catch-up pass for hours the Space was down. Idempotent (overwrites).""" + now = datetime.now(timezone.utc).replace(minute=0, second=0, microsecond=0) + for i in range(1, hours + 1): + await _run_hour(now - timedelta(hours=i)) + + +def start(backfill_hours: int = 6) -> None: + """Called from FastAPI lifespan startup.""" + global _scheduler + if os.environ.get("ML_INTERN_KPIS_DISABLED"): + logger.info("kpis_scheduler: disabled via ML_INTERN_KPIS_DISABLED") + return + if _scheduler is not None: + return + + try: + from apscheduler.schedulers.asyncio import AsyncIOScheduler + from apscheduler.triggers.cron import CronTrigger + except ImportError: + logger.warning("kpis_scheduler: apscheduler not installed, skipping") + return + + _scheduler = AsyncIOScheduler(timezone="UTC") + _scheduler.add_job( + run_last_completed_hour, + CronTrigger(minute=5), + id="kpis_hourly", + misfire_grace_time=600, # tolerate a 10-min misfire window + coalesce=True, # collapse multiple missed fires into one + max_instances=1, + replace_existing=True, + ) + _scheduler.start() + logger.info("kpis_scheduler: started (cron '5 * * * *' UTC)") + + # Non-blocking backfill. Hold a strong ref until done so asyncio doesn't + # GC the task before it finishes. + try: + task = asyncio.get_running_loop().create_task(backfill(backfill_hours)) + _background_tasks.add(task) + task.add_done_callback(_background_tasks.discard) + except RuntimeError: + # Not in an event loop (tests); skip backfill. + pass + + +async def shutdown() -> None: + """Called from FastAPI lifespan shutdown.""" + global _scheduler + if _scheduler is None: + return + _scheduler.shutdown(wait=False) + _scheduler = None + logger.info("kpis_scheduler: stopped") diff --git a/backend/main.py b/backend/main.py index 888740e53..f4f5ab529 100644 --- a/backend/main.py +++ b/backend/main.py @@ -1,5 +1,6 @@ """FastAPI application for HF Agent web interface.""" +import asyncio import logging import os from contextlib import asynccontextmanager @@ -9,33 +10,92 @@ from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles -from routes.agent import router as agent_router -from routes.auth import router as auth_router -# Load .env from project root (parent directory) +# Load .env before importing routes/session_manager so persistence and model +# modules see local settings during startup. load_dotenv(Path(__file__).parent.parent / ".env") +from routes.agent import router as agent_router # noqa: E402 +from routes.auth import router as auth_router # noqa: E402 +from session_manager import session_manager # noqa: E402 + # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) logger = logging.getLogger(__name__) +SHUTDOWN_USAGE_REFRESH_CONCURRENCY = 32 + + +async def _flush_session_on_shutdown(sid: str, agent_session, semaphore) -> None: + sess = agent_session.session + if not sess.config.save_sessions: + return + try: + async with semaphore: + await session_manager.refresh_session_usage_metrics( + agent_session, + error_code="lifespan_billing_snapshot_error", + ) + sess.save_and_upload_detached(sess.config.session_dataset_repo) + logger.info("Flushed session %s on shutdown", sid) + except Exception as e: + logger.warning("Failed to flush session %s: %s", sid, e) @asynccontextmanager async def lifespan(app: FastAPI): """Application lifespan handler.""" logger.info("Starting HF Agent backend...") + await session_manager.start() + # Start in-process hourly KPI rollup. Replaces an external cron so the + # rollup lives next to the data and reuses the Space's HF token. + try: + import kpis_scheduler + + kpis_scheduler.start() + except Exception as e: + logger.warning("KPI scheduler failed to start: %s", e) yield - logger.info("Shutting down HF Agent backend...") + logger.info("Shutting down HF Agent backend...") + try: + import kpis_scheduler + + await kpis_scheduler.shutdown() + except Exception as e: + logger.warning("KPI scheduler shutdown failed: %s", e) + + # Final-flush: save every still-active session so we don't lose traces on + # server restart. Billing refreshes are timeboxed and bounded; uploads are + # detached subprocesses. + try: + semaphore = asyncio.Semaphore(SHUTDOWN_USAGE_REFRESH_CONCURRENCY) + await asyncio.gather( + *( + _flush_session_on_shutdown(sid, agent_session, semaphore) + for sid, agent_session in list(session_manager.sessions.items()) + ) + ) + except Exception as e: + logger.warning("Lifespan final-flush skipped: %s", e) + await session_manager.close() + + +# Disable FastAPI auto-docs when running on HF Spaces (SPACE_ID is set by the +# platform) to avoid exposing the full API surface to anonymous visitors. Local +# dev keeps /docs and /redoc available. +_DOCS_DISABLED = os.environ.get("SPACE_ID") is not None app = FastAPI( title="HF Agent", description="ML Engineering Assistant API", version="1.0.0", lifespan=lifespan, + docs_url=None if _DOCS_DISABLED else "/docs", + redoc_url=None if _DOCS_DISABLED else "/redoc", + openapi_url=None if _DOCS_DISABLED else "/openapi.json", ) # CORS middleware for development diff --git a/backend/models.py b/backend/models.py index 954779f6a..0c06cef33 100644 --- a/backend/models.py +++ b/backend/models.py @@ -1,9 +1,9 @@ """Pydantic models for API requests and responses.""" from enum import Enum -from typing import Any +from typing import Any, Literal -from pydantic import BaseModel +from pydantic import BaseModel, Field class OpType(str, Enum): @@ -11,7 +11,6 @@ class OpType(str, Enum): USER_INPUT = "user_input" EXEC_APPROVAL = "exec_approval" - INTERRUPT = "interrupt" UNDO = "undo" COMPACT = "compact" SHUTDOWN = "shutdown" @@ -38,6 +37,7 @@ class ToolApproval(BaseModel): approved: bool feedback: str | None = None edited_script: str | None = None + namespace: str | None = None class ApprovalRequest(BaseModel): @@ -51,7 +51,10 @@ class SubmitRequest(BaseModel): """Request to submit user input.""" session_id: str - text: str + # Cap text size to prevent context-bloat / cost-amplification: a malicious + # or runaway client could otherwise attach megabytes that then ride along + # in every subsequent turn until /api/compact is called. + text: str = Field(..., min_length=1, max_length=100_000) class TruncateRequest(BaseModel): @@ -65,6 +68,7 @@ class SessionResponse(BaseModel): session_id: str ready: bool = True + model: str | None = None class PendingApprovalTool(BaseModel): @@ -75,17 +79,132 @@ class PendingApprovalTool(BaseModel): arguments: dict[str, Any] = {} +class SessionAutoApprovalInfo(BaseModel): + """Per-session auto-approval budget state.""" + + enabled: bool = False + cost_cap_usd: float | None = None + estimated_spend_usd: float = 0.0 + remaining_usd: float | None = None + + class SessionInfo(BaseModel): """Session metadata.""" session_id: str created_at: str + usage_window_started_at: str | None = None is_active: bool is_processing: bool = False message_count: int user_id: str = "dev" pending_approval: list[PendingApprovalTool] | None = None model: str | None = None + title: str | None = None + notification_destinations: list[str] = Field(default_factory=list) + auto_approval: SessionAutoApprovalInfo = Field( + default_factory=SessionAutoApprovalInfo + ) + + +class SessionNotificationsRequest(BaseModel): + """Replace the session's auto-notification destinations.""" + + destinations: list[str] + + +class SessionYoloRequest(BaseModel): + """Update a session's auto-approval policy.""" + + enabled: bool + cost_cap_usd: float | None = Field(default=None, ge=0) + + +class UsageBucket(BaseModel): + """App-attributed usage totals for a session.""" + + session_id: str | None = None + total_usd: float = 0.0 + inference_usd: float = 0.0 + hf_jobs_estimated_usd: float = 0.0 + sandbox_estimated_usd: float = 0.0 + llm_calls: int = 0 + hf_jobs_count: int = 0 + sandbox_count: int = 0 + prompt_tokens: int = 0 + completion_tokens: int = 0 + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 + total_tokens: int = 0 + hf_jobs_billable_seconds_estimate: int = 0 + sandbox_billable_seconds_estimate: int = 0 + + +class HfAccountUsageBucket(BaseModel): + """HF account billing usage for a time window.""" + + window_start: str | None = None + window_end: str | None = None + timezone: str | None = None + total_usd: float = 0.0 + inference_providers_usd: float = 0.0 + hf_jobs_usd: float = 0.0 + inference_provider_requests: int = 0 + hf_jobs_minutes: float = 0.0 + + +class HfInferenceProvidersCredits(BaseModel): + """Included and configured Inference Providers account credits.""" + + included_usd: float = 0.0 + used_usd: float = 0.0 + remaining_included_usd: float = 0.0 + limit_usd: float = 0.0 + remaining_limit_usd: float = 0.0 + num_requests: int = 0 + period_start: str | None = None + period_end: str | None = None + + +class HfAccountUsage(BaseModel): + """Authoritative HF account billing usage from the signed-in token.""" + + source: Literal["hf_billing"] + available: bool = False + error: str | None = None + current_session: HfAccountUsageBucket | None = None + month: HfAccountUsageBucket | None = None + inference_providers_credits: HfInferenceProvidersCredits | None = None + + +class UsageResponse(BaseModel): + """Current-user app-attributed usage response.""" + + source: Literal["app_telemetry"] + currency: Literal["USD"] + generated_at: str + timezone: str + session: UsageBucket | None = None + hf_account: HfAccountUsage | None = None + auto_approval: SessionAutoApprovalInfo | None = None + links: dict[str, str] = Field(default_factory=dict) + + +class DatasetUploadResponse(BaseModel): + """Response for a dataset file uploaded to the Hub.""" + + session_id: str + repo_id: str + repo_type: Literal["dataset"] = "dataset" + private: bool = True + upload_id: str + config_name: str + filename: str + path_in_repo: str + size_bytes: int + format: Literal["csv", "json", "jsonl"] + hub_url: str + load_dataset_snippet: str class HealthResponse(BaseModel): @@ -99,7 +218,9 @@ class HealthResponse(BaseModel): class LLMHealthResponse(BaseModel): """LLM provider health check response.""" - status: str # "ok" | "error" + status: str # "ok" | "error" | "skipped" model: str error: str | None = None - error_type: str | None = None # "auth" | "credits" | "rate_limit" | "network" | "unknown" + error_type: str | None = ( + None # "auth" | "credits" | "rate_limit" | "network" | "unknown" + ) diff --git a/backend/routes/agent.py b/backend/routes/agent.py index 7f5779952..06cf1c384 100644 --- a/backend/routes/agent.py +++ b/backend/routes/agent.py @@ -7,142 +7,267 @@ import asyncio import json import logging -import os +from datetime import datetime from typing import Any -from dependencies import get_current_user, require_huggingface_org_member +from dependencies import ( + INTERNAL_HF_TOKEN_KEY, + get_current_user, +) from fastapi import ( APIRouter, Depends, HTTPException, Request, ) +from fastapi.exceptions import RequestValidationError from fastapi.responses import StreamingResponse -from litellm import acompletion +from huggingface_hub.errors import HfHubHTTPError +from litellm import Message, acompletion +from pydantic import ValidationError +from starlette.datastructures import FormData, UploadFile +from dataset_uploads import ( + MAX_DATASET_UPLOAD_BYTES, + dataset_context_note, + push_dataset_upload_to_hub, +) from models import ( ApprovalRequest, + DatasetUploadResponse, HealthResponse, LLMHealthResponse, SessionInfo, + SessionNotificationsRequest, SessionResponse, + SessionYoloRequest, SubmitRequest, TruncateRequest, + UsageResponse, +) +from session_manager import ( + MAX_SESSIONS, + AgentSession, + SessionCapacityError, + session_manager, ) -from session_manager import MAX_SESSIONS, AgentSession, SessionCapacityError, session_manager - -import user_quotas +from agent.core.hf_access import get_jobs_access +from agent.core.hf_tokens import resolve_hf_request_token +from agent.core.local_models import local_model_provider from agent.core.llm_params import _resolve_llm_params +from agent.core.model_ids import ( + CLAUDE_OPUS_48_MODEL_ID, + DEEPSEEK_V4_PRO_MODEL_ID, + GLM_52_MODEL_ID, + GPT_55_MODEL_ID, + KIMI_K27_CODE_MODEL_ID, + MINIMAX_M3_MODEL_ID, + strip_huggingface_model_prefix, +) +from agent.core.prompt_caching import with_prompt_cache_params +from usage import build_usage_response logger = logging.getLogger(__name__) router = APIRouter(prefix="/api", tags=["agent"]) +_background_route_tasks: set[asyncio.Task] = set() -AVAILABLE_MODELS = [ - { - "id": "moonshotai/Kimi-K2.6", - "label": "Kimi K2.6", - "provider": "huggingface", - "tier": "free", - "recommended": True, - }, - { - "id": "bedrock/us.anthropic.claude-opus-4-6-v1", - "label": "Claude Opus 4.6", - "provider": "anthropic", - "tier": "pro", - "recommended": True, - }, - { - "id": "MiniMaxAI/MiniMax-M2.7", - "label": "MiniMax M2.7", - "provider": "huggingface", - "tier": "free", - }, - { - "id": "zai-org/GLM-5.1", - "label": "GLM 5.1", - "provider": "huggingface", - "tier": "free", - }, -] - - -def _is_anthropic_model(model_id: str) -> bool: - return "anthropic" in model_id - - -async def _require_hf_for_anthropic(request: Request, model_id: str) -> None: - """403 if a non-``huggingface``-org user tries to select an Anthropic model. - - Anthropic models are billed to the Space's ``ANTHROPIC_API_KEY``; every - other model in ``AVAILABLE_MODELS`` is routed through HF Router and - billed via ``X-HF-Bill-To``. The gate only fires for Anthropic so - non-HF users can still freely switch between the free models. - - Pattern: https://github.com/huggingface/ml-intern/pull/63 - """ - if not _is_anthropic_model(model_id): - return - if not await require_huggingface_org_member(request): - raise HTTPException( - status_code=403, - detail={ - "error": "anthropic_restricted", - "message": ( - "Opus is gated to HF staff. Pick a free model — " - "Kimi K2.6, MiniMax M2.7, or GLM 5.1 — instead." - ), - }, +DEFAULT_GPT_MODEL_ID = GPT_55_MODEL_ID +DEFAULT_MODEL_ID = GLM_52_MODEL_ID +DATASET_UPLOAD_MULTIPART_SLACK_BYTES = 1024 * 1024 + + +async def _reset_usage_window(session_id: str) -> dict[str, Any] | None: + return await session_manager.reset_session_usage_window( + session_id, + started_at=datetime.utcnow(), + ) + + +async def _refresh_usage_and_upload( + agent_session: AgentSession, + *, + error_code: str, +) -> None: + session = agent_session.session + try: + await session_manager.refresh_session_usage_metrics( + agent_session, + error_code=error_code, + ) + session.save_and_upload_detached(session.config.session_dataset_repo) + except Exception as e: + logger.warning( + "Background usage refresh/upload failed for %s: %s", + agent_session.session_id, + e, ) -async def _enforce_claude_quota( - user: dict[str, Any], +def _schedule_usage_refresh_and_upload( agent_session: AgentSession, + *, + error_code: str, ) -> None: - """Charge the user's daily Claude quota on first use of Anthropic in a session. + task = asyncio.create_task( + _refresh_usage_and_upload(agent_session, error_code=error_code) + ) + _background_route_tasks.add(task) + task.add_done_callback(_background_route_tasks.discard) - Runs at *message-submit* time, not session-create time — so spinning up a - Claude session to look around doesn't burn quota. The ``claude_counted`` - flag on ``AgentSession`` guards against re-counting the same session. - No-ops when the session's current model isn't Anthropic, or when this - session has already been charged. Raises 429 when the user has hit - their daily cap. +def _available_models() -> list[dict[str, Any]]: + models = [ + { + "id": CLAUDE_OPUS_48_MODEL_ID, + "label": "Claude Opus 4.8", + }, + { + "id": DEFAULT_GPT_MODEL_ID, + "label": "GPT-5.5", + }, + { + "id": KIMI_K27_CODE_MODEL_ID, + "label": "Kimi K2.7 Code", + }, + { + "id": MINIMAX_M3_MODEL_ID, + "label": "MiniMax M3", + }, + { + "id": DEFAULT_MODEL_ID, + "label": "GLM 5.2", + "recommended": True, + }, + { + "id": DEEPSEEK_V4_PRO_MODEL_ID, + "label": "DeepSeek V4 Pro", + }, + ] + return models + + +AVAILABLE_MODELS = _available_models() + + +def _valid_model_ids() -> set[str]: + return {m["id"] for m in AVAILABLE_MODELS} + + +def _validate_model_id(model_id: str | None) -> None: + if not model_id or model_id in _valid_model_ids(): + return + raise HTTPException(status_code=400, detail=f"Unknown model: {model_id}") + + +def _default_model() -> str: + return DEFAULT_MODEL_ID + + +def _model_override_for_new_session(requested_model: str | None) -> str | None: + """Return the model override to use when creating a new session. + + Explicit model requests are honored. Empty web requests default to GLM 5.2. """ - if agent_session.claude_counted: + return requested_model or _default_model() + + +def _user_hf_token(user: dict[str, Any] | None) -> str | None: + if not isinstance(user, dict): + return None + return user.get(INTERNAL_HF_TOKEN_KEY) + + +def _model_requires_hf_router_token(model_id: str | None) -> bool: + normalized = strip_huggingface_model_prefix(model_id) or model_id or "" + return local_model_provider(normalized) is None + + +def _reject_oversize_dataset_upload(request: Request) -> None: + raw_content_length = request.headers.get("content-length") + if raw_content_length is None: return - model_name = agent_session.session.config.model_name - if not _is_anthropic_model(model_name): + try: + content_length = int(raw_content_length) + except (TypeError, ValueError): return - user_id = user["user_id"] - used = await user_quotas.get_claude_used_today(user_id) - cap = user_quotas.daily_cap_for(user.get("plan")) - if used >= cap: + if content_length > MAX_DATASET_UPLOAD_BYTES + DATASET_UPLOAD_MULTIPART_SLACK_BYTES: raise HTTPException( - status_code=429, - detail={ - "error": "claude_daily_cap", - "plan": user.get("plan", "free"), - "cap": cap, - "message": ( - "Daily Claude limit reached. Upgrade to HF Pro for " - f"{user_quotas.CLAUDE_PRO_DAILY}/day or use a free model." - ), - }, + status_code=413, + detail="Dataset upload exceeds the 100 MB limit.", ) - await user_quotas.increment_claude(user_id) - agent_session.claude_counted = True -def _check_session_access(session_id: str, user: dict[str, Any]) -> None: - """Verify the user has access to the given session. Raises 403 or 404.""" - info = session_manager.get_session_info(session_id) - if not info: +def _dataset_upload_file_from_form(form: FormData) -> UploadFile: + uploaded_files = [ + (key, value) + for key, value in form.multi_items() + if isinstance(value, UploadFile) + ] + if len(uploaded_files) != 1: + raise HTTPException( + status_code=400, + detail="Upload exactly one dataset file.", + ) + field_name, upload = uploaded_files[0] + if field_name != "file": + raise HTTPException( + status_code=400, + detail="Missing 'file' upload field.", + ) + return upload + + +def _dataset_upload_hub_http_exception(error: HfHubHTTPError) -> HTTPException: + status_code = getattr(error.response, "status_code", None) + if status_code == 401: + detail = "Hugging Face rejected the token used for the dataset upload." + return HTTPException(status_code=401, detail=detail) + if status_code == 403: + detail = ( + "Hugging Face denied permission to create or write to the dataset repo." + ) + return HTTPException(status_code=403, detail=detail) + if status_code == 404: + detail = "Could not find the Hugging Face namespace or dataset repo." + return HTTPException(status_code=404, detail=detail) + if status_code == 429: + detail = "Hugging Face Hub rate limit reached while uploading the dataset." + return HTTPException(status_code=429, detail=detail) + return HTTPException( + status_code=502, + detail="Hugging Face Hub upload failed. Please try again.", + ) + + +async def _check_session_access( + session_id: str, + user: dict[str, Any], + request: Request | None = None, + preload_sandbox: bool = True, +) -> AgentSession: + """Verify and lazily load the user's session. Raises 403 or 404.""" + hf_token = ( + resolve_hf_request_token(request) + if request is not None + else _user_hf_token(user) + ) + agent_session = await session_manager.ensure_session_loaded( + session_id, + user["user_id"], + hf_token=hf_token, + hf_username=user.get("username"), + user_plan=user.get("plan"), + preload_sandbox=preload_sandbox, + ) + if not agent_session: raise HTTPException(status_code=404, detail="Session not found") - if not session_manager.verify_session_access(session_id, user["user_id"]): + if user["user_id"] != "dev" and agent_session.user_id not in { + user["user_id"], + "dev", + }: raise HTTPException(status_code=403, detail="Access denied to this session") + return agent_session @router.get("/health", response_model=HealthResponse) @@ -156,18 +281,32 @@ async def health_check() -> HealthResponse: @router.get("/health/llm", response_model=LLMHealthResponse) -async def llm_health_check() -> LLMHealthResponse: +async def llm_health_check( + request: Request, + user: dict = Depends(get_current_user), +) -> LLMHealthResponse: """Check if the LLM provider is reachable and the API key is valid. - Makes a minimal 1-token completion call. Catches common errors: + Makes a minimal 1-token completion call against the authenticated user's + default model when a token is available. For token-less HF Router requests, + returns ``status="skipped"`` instead of making an unauthenticated probe. + Catches common errors: - 401 → invalid API key - 402/insufficient_quota → out of credits - 429 → rate limited - timeout / network → provider unreachable """ - model = session_manager.config.model_name + model = _default_model() + hf_token = resolve_hf_request_token(request) + if _model_requires_hf_router_token(model) and not hf_token: + return LLMHealthResponse(status="skipped", model=model) + try: - llm_params = _resolve_llm_params(model, reasoning_effort="high") + llm_params = _resolve_llm_params( + model, + hf_token, + reasoning_effort="high", + ) await acompletion( messages=[{"role": "user", "content": "hi"}], max_tokens=1, @@ -232,18 +371,15 @@ async def generate_title( reasoning model — reasoning_effort=low keeps the reasoning budget small so the 60-token output budget isn't consumed before the title is written. """ - api_key = ( - os.environ.get("INFERENCE_TOKEN") - or (user.get("hf_token") if isinstance(user, dict) else None) - or os.environ.get("HF_TOKEN") - ) try: + await _check_session_access(request.session_id, user) + llm_params = _resolve_llm_params( + "openai/gpt-oss-120b:cerebras", + _user_hf_token(user), + reasoning_effort="low", + ) + llm_params = with_prompt_cache_params(llm_params) response = await acompletion( - # Double openai/ prefix: LiteLLM strips the first as its provider - # prefix, leaving the HF model id on the wire for the router. - model="openai/openai/gpt-oss-120b:cerebras", - api_base="https://router.huggingface.co/v1", - api_key=api_key, messages=[ { "role": "system", @@ -260,17 +396,31 @@ async def generate_title( max_tokens=60, temperature=0.3, timeout=10, - reasoning_effort="low", + **llm_params, ) title = response.choices[0].message.content.strip().strip('"').strip("'") title = title.translate(_TITLE_STRIP_CHARS).strip() if len(title) > 50: title = title[:50].rstrip() + "…" + try: + await session_manager.update_session_title(request.session_id, title) + except Exception: + logger.debug( + "Skipping title persistence for missing session %s", request.session_id + ) return {"title": title} except Exception as e: logger.warning(f"Title generation failed: {e}") fallback = request.text.strip() title = fallback[:40].rstrip() + "…" if len(fallback) > 40 else fallback + try: + await _check_session_access(request.session_id, user) + await session_manager.update_session_title(request.session_id, title) + except Exception: + logger.debug( + "Skipping fallback title persistence for missing session %s", + request.session_id, + ) return {"title": title} @@ -285,20 +435,12 @@ async def create_session( behalf of the user. Optional body ``{"model"?: }`` selects the session's LLM; unknown - ids are rejected (400). The Claude-quota gate runs at message-submit - time, not here — spinning up an Opus session to look around is free. + ids are rejected (400). Empty requests use the web default. Returns 503 if the server or user has reached the session limit. """ # Extract the user's HF token (Bearer header, HttpOnly cookie, or env var) - hf_token = None - auth_header = request.headers.get("Authorization", "") - if auth_header.startswith("Bearer "): - hf_token = auth_header[7:] - if not hf_token: - hf_token = request.cookies.get("hf_access_token") - if not hf_token: - hf_token = os.environ.get("HF_TOKEN") + hf_token = resolve_hf_request_token(request) # Optional model override. Empty body falls back to the config default. model: str | None = None @@ -309,23 +451,30 @@ async def create_session( if isinstance(body, dict): model = body.get("model") - valid_ids = {m["id"] for m in AVAILABLE_MODELS} - if model and model not in valid_ids: - raise HTTPException(status_code=400, detail=f"Unknown model: {model}") + _validate_model_id(model) - # Opus is gated to HF staff (PR #63). Only fires when the resolved model - # is Anthropic; free models pass through. - resolved_model = model or session_manager.config.model_name - await _require_hf_for_anthropic(request, resolved_model) + # Empty requests use the web default. + model = _model_override_for_new_session(model) try: session_id = await session_manager.create_session( - user_id=user["user_id"], hf_token=hf_token, model=model + user_id=user["user_id"], + hf_username=user.get("username"), + hf_token=hf_token, + user_plan=user.get("plan"), + model=model, + is_pro=user.get("plan") == "pro", ) except SessionCapacityError as e: raise HTTPException(status_code=503, detail=str(e)) - return SessionResponse(session_id=session_id, ready=True) + await _reset_usage_window(session_id) + + return SessionResponse( + session_id=session_id, + ready=True, + model=model, + ) @router.post("/session/restore-summary", response_model=SessionResponse) @@ -337,37 +486,40 @@ async def restore_session_summary( summarization prompt on them and drop the result into the new session's context as a user-role system note. - Optional ``"model"`` in the body overrides the session's LLM. The - Claude-quota gate runs at message-submit time, not here. + Optional ``"model"`` in the body overrides the session's LLM; otherwise + the new session uses the web default. """ messages = body.get("messages") if not isinstance(messages, list) or not messages: raise HTTPException(status_code=400, detail="Missing 'messages' array") - hf_token = None - auth_header = request.headers.get("Authorization", "") - if auth_header.startswith("Bearer "): - hf_token = auth_header[7:] - if not hf_token: - hf_token = request.cookies.get("hf_access_token") - if not hf_token: - hf_token = os.environ.get("HF_TOKEN") + hf_token = resolve_hf_request_token(request) model = body.get("model") - valid_ids = {m["id"] for m in AVAILABLE_MODELS} - if model and model not in valid_ids: - raise HTTPException(status_code=400, detail=f"Unknown model: {model}") + _validate_model_id(model) - resolved_model = model or session_manager.config.model_name - await _require_hf_for_anthropic(request, resolved_model) + model = _model_override_for_new_session(model) try: session_id = await session_manager.create_session( - user_id=user["user_id"], hf_token=hf_token, model=model + user_id=user["user_id"], + hf_username=user.get("username"), + hf_token=hf_token, + user_plan=user.get("plan"), + model=model, + is_pro=user.get("plan") == "pro", ) except SessionCapacityError as e: raise HTTPException(status_code=503, detail=str(e)) + await _reset_usage_window(session_id) + + await _check_session_access( + session_id, + user, + request, + preload_sandbox=False, + ) try: summarized = await session_manager.seed_from_summary(session_id, messages) except ValueError as e: @@ -380,7 +532,11 @@ async def restore_session_summary( f"Seeded session {session_id} for {user.get('username', 'unknown')} " f"(summary of {summarized} messages)" ) - return SessionResponse(session_id=session_id, ready=True) + return SessionResponse( + session_id=session_id, + ready=True, + model=model, + ) @router.get("/session/{session_id}", response_model=SessionInfo) @@ -388,11 +544,25 @@ async def get_session( session_id: str, user: dict = Depends(get_current_user) ) -> SessionInfo: """Get session information. Only accessible by the session owner.""" - _check_session_access(session_id, user) + await _check_session_access(session_id, user) info = session_manager.get_session_info(session_id) return SessionInfo(**info) +@router.post("/session/{session_id}/activate", response_model=SessionInfo) +async def activate_session( + session_id: str, + request: Request, + user: dict = Depends(get_current_user), +) -> SessionInfo: + """Mark a session as actively revisited without resetting usage.""" + await _check_session_access(session_id, user, request) + info = await session_manager.activate_session(session_id) + if not info: + raise HTTPException(status_code=404, detail="Session not found") + return SessionInfo(**info) + + @router.post("/session/{session_id}/model") async def set_session_model( session_id: str, @@ -403,24 +573,16 @@ async def set_session_model( """Switch the active model for a single session (tab-scoped). Takes effect on the next LLM call in that session — other sessions - (including other browser tabs) are unaffected. Model switches don't - charge quota — the Claude-quota gate only fires at message-submit time. - - Switching TO an Anthropic model requires HF org membership (PR #63); - free-model switches are unrestricted. + (including other browser tabs) are unaffected. """ - _check_session_access(session_id, user) + agent_session = await _check_session_access(session_id, user, request) model_id = body.get("model") if not model_id: raise HTTPException(status_code=400, detail="Missing 'model' field") - valid_ids = {m["id"] for m in AVAILABLE_MODELS} - if model_id not in valid_ids: - raise HTTPException(status_code=400, detail=f"Unknown model: {model_id}") - await _require_hf_for_anthropic(request, model_id) - agent_session = session_manager.sessions.get(session_id) + _validate_model_id(model_id) if not agent_session: raise HTTPException(status_code=404, detail="Session not found") - agent_session.session.update_model(model_id) + await session_manager.update_session_model(session_id, model_id) logger.info( f"Session {session_id} model → {model_id} " f"(by {user.get('username', 'unknown')})" @@ -428,33 +590,211 @@ async def set_session_model( return {"session_id": session_id, "model": model_id} -@router.get("/user/quota") -async def get_user_quota(user: dict = Depends(get_current_user)) -> dict: - """Return the user's plan tier and today's Claude-session quota state.""" - plan = user.get("plan", "free") - used = await user_quotas.get_claude_used_today(user["user_id"]) - cap = user_quotas.daily_cap_for(plan) +@router.post("/session/{session_id}/notifications") +async def set_session_notifications( + session_id: str, + body: SessionNotificationsRequest, + user: dict = Depends(get_current_user), +) -> dict: + """Replace the session's auto-notification destinations.""" + agent_session = await _check_session_access(session_id, user) + try: + destinations = session_manager.set_notification_destinations( + session_id, body.destinations + ) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + await session_manager.persist_session_snapshot(agent_session) + return { + "session_id": session_id, + "notification_destinations": destinations, + } + + +@router.post("/session/{session_id}/datasets", response_model=DatasetUploadResponse) +async def upload_session_dataset( + session_id: str, + request: Request, + user: dict = Depends(get_current_user), +) -> DatasetUploadResponse: + """Upload a CSV/JSON dataset file to a private Hub dataset for this session.""" + file: UploadFile | None = None + try: + _reject_oversize_dataset_upload(request) + agent_session = await _check_session_access(session_id, user, request) + if not agent_session or not agent_session.is_active: + raise HTTPException(status_code=404, detail="Session not found") + if agent_session.is_processing: + raise HTTPException( + status_code=409, + detail="Cannot upload a dataset while the agent is processing.", + ) + if agent_session.session.pending_approval: + raise HTTPException( + status_code=409, + detail="Resolve pending approvals before uploading a dataset.", + ) + + hf_token = ( + resolve_hf_request_token(request, include_env_fallback=False) + or _user_hf_token(user) + or resolve_hf_request_token(request) + ) + if not hf_token: + raise HTTPException( + status_code=401, + detail="A Hugging Face token is required to upload datasets.", + ) + + form = await request.form( + max_files=1, + max_fields=1, + max_part_size=MAX_DATASET_UPLOAD_BYTES, + ) + file = _dataset_upload_file_from_form(form) + hf_username = user.get("username") or agent_session.hf_username + uploaded = await push_dataset_upload_to_hub( + upload=file, + session_id=session_id, + hf_username=hf_username, + hf_token=hf_token, + ) + agent_session.session.context_manager.add_message( + Message(role="user", content=dataset_context_note(uploaded)) + ) + session_manager._touch(agent_session) + await session_manager.persist_session_snapshot(agent_session) + logger.info( + "Uploaded dataset file %s to %s for session %s", + uploaded.filename, + uploaded.repo_id, + session_id, + ) + return DatasetUploadResponse(**uploaded.response_payload()) + except HTTPException: + raise + except HfHubHTTPError as e: + logger.warning( + "Hub rejected dataset upload for session %s: status=%s request_id=%s", + session_id, + getattr(e.response, "status_code", None), + getattr(e, "request_id", None), + ) + raise _dataset_upload_hub_http_exception(e) + except Exception: + logger.exception("Dataset upload failed for session %s", session_id) + raise HTTPException( + status_code=502, + detail="Dataset upload failed. Please try again.", + ) + finally: + if file is not None: + await file.close() + + +@router.patch("/session/{session_id}/yolo") +async def set_session_yolo( + session_id: str, + body: SessionYoloRequest, + user: dict = Depends(get_current_user), +) -> dict: + """Update the session-scoped auto-approval policy.""" + await _check_session_access(session_id, user) + try: + summary = await session_manager.update_session_auto_approval( + session_id, + enabled=body.enabled, + cost_cap_usd=body.cost_cap_usd, + cap_provided="cost_cap_usd" in body.model_fields_set, + ) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + return {"session_id": session_id, **summary} + + +@router.get("/user/jobs-access") +async def get_jobs_access_info( + request: Request, user: dict = Depends(get_current_user) +) -> dict: + """Return the namespaces the current token can run HF Jobs under. + + Credits are enforced by the HF API at job-creation time, not here — + the response only describes which wallets the caller is allowed to + pick from. Pro is irrelevant. + """ + token = resolve_hf_request_token(request) + + access = await get_jobs_access(token or "") return { - "plan": plan, - "claude_used_today": used, - "claude_daily_cap": cap, - "claude_remaining": max(0, cap - used), + "eligible_namespaces": access.eligible_namespaces if access else [], + "default_namespace": access.default_namespace if access else None, + "billing_url": "https://huggingface.co/settings/billing", } +@router.get("/usage", response_model=UsageResponse) +async def get_usage( + request: Request, + session_id: str | None = None, + tz: str | None = None, + user: dict = Depends(get_current_user), +) -> dict: + """Return app-attributed usage for the current user.""" + if session_id: + await _check_session_access( + session_id, + user, + request, + preload_sandbox=False, + ) + usage = await build_usage_response( + session_manager, + user_id=user["user_id"], + hf_token=( + resolve_hf_request_token(request, include_env_fallback=False) + or _user_hf_token(user) + or resolve_hf_request_token(request) + ), + session_id=session_id, + timezone_name=tz, + ) + if session_id: + auto_approval = ( + await session_manager.reconcile_session_auto_approval_from_usage( + session_id, + usage, + ) + ) + if auto_approval is not None: + usage["auto_approval"] = auto_approval + return usage + + @router.get("/sessions", response_model=list[SessionInfo]) async def list_sessions(user: dict = Depends(get_current_user)) -> list[SessionInfo]: """List sessions belonging to the authenticated user.""" - sessions = session_manager.list_sessions(user_id=user["user_id"]) + sessions = await session_manager.list_sessions(user_id=user["user_id"]) return [SessionInfo(**s) for s in sessions] +@router.post("/session/{session_id}/sandbox/teardown") +async def teardown_session_sandbox( + session_id: str, user: dict = Depends(get_current_user) +) -> dict: + """Best-effort sandbox teardown that preserves durable chat history.""" + await _check_session_access(session_id, user, preload_sandbox=False) + task = asyncio.create_task(session_manager.teardown_sandbox(session_id)) + _background_route_tasks.add(task) + task.add_done_callback(_background_route_tasks.discard) + return {"status": "teardown_requested", "session_id": session_id} + + @router.delete("/session/{session_id}") async def delete_session( session_id: str, user: dict = Depends(get_current_user) ) -> dict: """Delete a session. Only accessible by the session owner.""" - _check_session_access(session_id, user) + await _check_session_access(session_id, user, preload_sandbox=False) success = await session_manager.delete_session(session_id) if not success: raise HTTPException(status_code=404, detail="Session not found") @@ -463,17 +803,40 @@ async def delete_session( @router.post("/submit") async def submit_input( - request: SubmitRequest, user: dict = Depends(get_current_user) + request: Request, user: dict = Depends(get_current_user) ) -> dict: """Submit user input to a session. Only accessible by the session owner.""" - _check_session_access(request.session_id, user) - agent_session = session_manager.sessions.get(request.session_id) - if agent_session is not None: - await _enforce_claude_quota(user, agent_session) - success = await session_manager.submit_user_input(request.session_id, request.text) + # Parse the body manually so session ownership can be checked before the + # text-length constraints fire — otherwise a non-owner sending an empty + # or oversized text gets a 422 leaking the constraint instead of the 404 + # they'd get for any other access to a session they don't own. + try: + payload = await request.json() + except (json.JSONDecodeError, TypeError) as exc: + raise HTTPException(status_code=422, detail=str(exc)) + if not isinstance(payload, dict): + raise HTTPException(status_code=422, detail="Body must be a JSON object") + raw_session_id = payload.get("session_id") + if not isinstance(raw_session_id, str) or not raw_session_id: + raise RequestValidationError( + [ + { + "type": "missing", + "loc": ("body", "session_id"), + "msg": "Field required", + "input": payload, + } + ] + ) + await _check_session_access(raw_session_id, user) + try: + body = SubmitRequest(**payload) + except ValidationError as exc: + raise RequestValidationError(exc.errors()) from exc + success = await session_manager.submit_user_input(body.session_id, body.text) if not success: raise HTTPException(status_code=404, detail="Session not found or inactive") - return {"status": "submitted", "session_id": request.session_id} + return {"status": "submitted", "session_id": body.session_id} @router.post("/approve") @@ -481,13 +844,14 @@ async def submit_approval( request: ApprovalRequest, user: dict = Depends(get_current_user) ) -> dict: """Submit tool approvals to a session. Only accessible by the session owner.""" - _check_session_access(request.session_id, user) + await _check_session_access(request.session_id, user) approvals = [ { "tool_call_id": a.tool_call_id, "approved": a.approved, "feedback": a.feedback, "edited_script": a.edited_script, + "namespace": a.namespace, } for a in request.approvals ] @@ -504,9 +868,7 @@ async def chat_sse( user: dict = Depends(get_current_user), ) -> StreamingResponse: """SSE endpoint: submit input or approval, then stream events until turn ends.""" - _check_session_access(session_id, user) - - agent_session = session_manager.sessions.get(session_id) + agent_session = await _check_session_access(session_id, user, request) if not agent_session or not agent_session.is_active: raise HTTPException(status_code=404, detail="Session not found or inactive") @@ -522,16 +884,6 @@ async def chat_sse( text = body.get("text") approvals = body.get("approvals") - # Gate user-message sends against the daily Claude quota. Approvals are - # continuations of an in-progress turn — the session was already charged - # on its first message, so we skip the gate there. - if text is not None and not approvals: - try: - await _enforce_claude_quota(user, agent_session) - except HTTPException: - broadcaster.unsubscribe(sub_id) - raise - try: if approvals: formatted = [ @@ -540,6 +892,7 @@ async def chat_sse( "approved": a["approved"], "feedback": a.get("feedback"), "edited_script": a.get("edited_script"), + "namespace": a.get("namespace"), } for a in approvals ] @@ -548,12 +901,15 @@ async def chat_sse( success = await session_manager.submit_user_input(session_id, text) else: broadcaster.unsubscribe(sub_id) - raise HTTPException(status_code=400, detail="Must provide 'text' or 'approvals'") + raise HTTPException( + status_code=400, detail="Must provide 'text' or 'approvals'" + ) if not success: broadcaster.unsubscribe(sub_id) raise HTTPException(status_code=404, detail="Session not found or inactive") except HTTPException: + broadcaster.unsubscribe(sub_id) raise except Exception: broadcaster.unsubscribe(sub_id) @@ -562,19 +918,92 @@ async def chat_sse( return _sse_response(broadcaster, event_queue, sub_id) +@router.post("/pro-click/{session_id}") +async def record_pro_click( + session_id: str, + body: dict, + user: dict = Depends(get_current_user), +) -> dict: + """Record a click on a Pro upgrade CTA shown from inside a session.""" + agent_session = await _check_session_access(session_id, user) + + from agent.core import telemetry + + await telemetry.record_pro_cta_click( + agent_session.session, + source=str(body.get("source") or "unknown"), + target=str(body.get("target") or "pro_pricing"), + ) + if agent_session.session.config.save_sessions: + _schedule_usage_refresh_and_upload( + agent_session, + error_code="pro_click_billing_snapshot_error", + ) + return {"status": "ok"} + + # --------------------------------------------------------------------------- # Shared SSE helpers # --------------------------------------------------------------------------- -_TERMINAL_EVENTS = {"turn_complete", "approval_required", "error", "interrupted", "shutdown"} +_TERMINAL_EVENTS = { + "turn_complete", + "approval_required", + "error", + "interrupted", + "shutdown", +} _SSE_KEEPALIVE_SECONDS = 15 -def _sse_response(broadcaster, event_queue, sub_id) -> StreamingResponse: +def _last_event_seq(request: Request) -> int: + raw = ( + request.headers.get("last-event-id") or request.query_params.get("after") or "0" + ) + try: + return max(0, int(raw)) + except (TypeError, ValueError): + return 0 + + +def _format_sse(msg: dict[str, Any]) -> str: + seq = msg.get("seq") + body = {"event_type": msg.get("event_type"), "data": msg.get("data") or {}} + if seq is not None: + body["seq"] = seq + return f"id: {seq}\ndata: {json.dumps(body)}\n\n" + return f"data: {json.dumps(body)}\n\n" + + +def _event_doc_to_msg(doc: dict[str, Any]) -> dict[str, Any]: + return { + "event_type": doc.get("event_type"), + "data": doc.get("data") or {}, + "seq": doc.get("seq"), + } + + +def _sse_response( + broadcaster, + event_queue, + sub_id, + *, + replay_events: list[dict[str, Any]] | None = None, + after_seq: int = 0, +) -> StreamingResponse: """Build a StreamingResponse that drains *event_queue* as SSE, sending keepalive comments every 15 s to prevent proxy timeouts.""" async def event_generator(): try: + for doc in replay_events or []: + msg = _event_doc_to_msg(doc) + seq = msg.get("seq") + if isinstance(seq, int) and seq <= after_seq: + continue + yield _format_sse(msg) + if msg.get("event_type", "") in _TERMINAL_EVENTS: + return + while True: try: msg = await asyncio.wait_for( @@ -585,7 +1014,7 @@ async def event_generator(): yield ": keepalive\n\n" continue event_type = msg.get("event_type", "") - yield f"data: {json.dumps(msg)}\n\n" + yield _format_sse(msg) if event_type in _TERMINAL_EVENTS: break finally: @@ -605,6 +1034,7 @@ async def event_generator(): @router.get("/events/{session_id}") async def subscribe_events( session_id: str, + request: Request, user: dict = Depends(get_current_user), ) -> StreamingResponse: """Subscribe to events for a running session without submitting new input. @@ -612,15 +1042,23 @@ async def subscribe_events( Used by the frontend to re-attach after a connection drop (e.g. screen sleep). Returns 404 if the session isn't active or isn't processing. """ - _check_session_access(session_id, user) - - agent_session = session_manager.sessions.get(session_id) + agent_session = await _check_session_access(session_id, user, request) if not agent_session or not agent_session.is_active: raise HTTPException(status_code=404, detail="Session not found or inactive") + after_seq = _last_event_seq(request) + replay_events = await session_manager._store().load_events_after( + session_id, after_seq + ) broadcaster = agent_session.broadcaster sub_id, event_queue = broadcaster.subscribe() - return _sse_response(broadcaster, event_queue, sub_id) + return _sse_response( + broadcaster, + event_queue, + sub_id, + replay_events=replay_events, + after_seq=after_seq, + ) @router.post("/interrupt/{session_id}") @@ -628,7 +1066,7 @@ async def interrupt_session( session_id: str, user: dict = Depends(get_current_user) ) -> dict: """Interrupt the current operation in a session.""" - _check_session_access(session_id, user) + await _check_session_access(session_id, user) success = await session_manager.interrupt(session_id) if not success: raise HTTPException(status_code=404, detail="Session not found or inactive") @@ -640,17 +1078,19 @@ async def get_session_messages( session_id: str, user: dict = Depends(get_current_user) ) -> list[dict]: """Return the session's message history from memory.""" - _check_session_access(session_id, user) - agent_session = session_manager.sessions.get(session_id) + agent_session = await _check_session_access(session_id, user) if not agent_session or not agent_session.is_active: raise HTTPException(status_code=404, detail="Session not found or inactive") - return [msg.model_dump() for msg in agent_session.session.context_manager.items] + return [ + msg.model_dump(mode="json") + for msg in agent_session.session.context_manager.items + ] @router.post("/undo/{session_id}") async def undo_session(session_id: str, user: dict = Depends(get_current_user)) -> dict: """Undo the last turn in a session.""" - _check_session_access(session_id, user) + await _check_session_access(session_id, user) success = await session_manager.undo(session_id) if not success: raise HTTPException(status_code=404, detail="Session not found or inactive") @@ -659,13 +1099,30 @@ async def undo_session(session_id: str, user: dict = Depends(get_current_user)) @router.post("/truncate/{session_id}") async def truncate_session( - session_id: str, body: TruncateRequest, user: dict = Depends(get_current_user) + session_id: str, + request: Request, + user: dict = Depends(get_current_user), ) -> dict: """Truncate conversation to before a specific user message.""" - _check_session_access(session_id, user) + # Check session ownership before parsing the request body so a 404 on a + # non-existent / non-owned session_id beats the 422 schema-validation error + # (otherwise the response leaks the required field name to non-owners). + await _check_session_access(session_id, user) + try: + body = TruncateRequest(**(await request.json())) + except ValidationError as exc: + # Re-raise as RequestValidationError so FastAPI returns its standard + # structured 422 schema (`{"detail": [{"type":..., "loc":..., ...}]}`) + # instead of a string-stringified Pydantic dump. + raise RequestValidationError(exc.errors()) from exc + except (json.JSONDecodeError, TypeError) as exc: + raise HTTPException(status_code=422, detail=str(exc)) success = await session_manager.truncate(session_id, body.user_message_index) if not success: - raise HTTPException(status_code=404, detail="Session not found, inactive, or message index out of range") + raise HTTPException( + status_code=404, + detail="Session not found, inactive, or message index out of range", + ) return {"status": "truncated", "session_id": session_id} @@ -674,7 +1131,7 @@ async def compact_session( session_id: str, user: dict = Depends(get_current_user) ) -> dict: """Compact the context in a session.""" - _check_session_access(session_id, user) + await _check_session_access(session_id, user) success = await session_manager.compact(session_id) if not success: raise HTTPException(status_code=404, detail="Session not found or inactive") @@ -686,10 +1143,45 @@ async def shutdown_session( session_id: str, user: dict = Depends(get_current_user) ) -> dict: """Shutdown a session.""" - _check_session_access(session_id, user) + await _check_session_access(session_id, user) success = await session_manager.shutdown_session(session_id) if not success: raise HTTPException(status_code=404, detail="Session not found or inactive") return {"status": "shutdown_requested", "session_id": session_id} +@router.post("/feedback/{session_id}") +async def submit_feedback( + session_id: str, + body: dict, + user: dict = Depends(get_current_user), +) -> dict: + """Attach a user feedback signal to a session's event log. + + Body: {rating: "up"|"down"|"outcome_success"|"outcome_fail", + turn_index?: int, comment?: str, message_id?: str} + Appended as a `feedback` event and saved with the session trajectory. + """ + agent_session = await _check_session_access(session_id, user) + + rating = body.get("rating") + if rating not in {"up", "down", "outcome_success", "outcome_fail"}: + raise HTTPException(status_code=400, detail="invalid rating") + + from agent.core import telemetry + + await telemetry.record_feedback( + agent_session.session, + rating=rating, + turn_index=body.get("turn_index"), + message_id=body.get("message_id"), + comment=body.get("comment"), + ) + # Fire-and-forget save so feedback reaches the dataset even if the user + # closes the tab right after clicking. + if agent_session.session.config.save_sessions: + _schedule_usage_refresh_and_upload( + agent_session, + error_code="feedback_billing_snapshot_error", + ) + return {"status": "ok"} diff --git a/backend/routes/auth.py b/backend/routes/auth.py index dce10fe2f..d736deff1 100644 --- a/backend/routes/auth.py +++ b/backend/routes/auth.py @@ -4,28 +4,47 @@ After successful auth, sets an HttpOnly cookie with the access token. """ +import logging import os import secrets import time from urllib.parse import urlencode import httpx -from dependencies import AUTH_ENABLED, check_org_membership, get_current_user +from dependencies import ( + AUTH_ENABLED, + OAUTH_SCOPE_COOKIE, + REQUIRED_OAUTH_SCOPES, + configured_oauth_scopes, + get_current_user, + oauth_scope_fingerprint, +) from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import RedirectResponse router = APIRouter(prefix="/auth", tags=["auth"]) +logger = logging.getLogger(__name__) # OAuth configuration from environment OAUTH_CLIENT_ID = os.environ.get("OAUTH_CLIENT_ID", "") OAUTH_CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET", "") OPENID_PROVIDER_URL = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co") +OAUTH_SCOPES = configured_oauth_scopes() # In-memory OAuth state store with expiry (5 min TTL) _OAUTH_STATE_TTL = 300 oauth_states: dict[str, dict] = {} +def _missing_required_scopes(token_data: dict) -> set[str]: + raw_scopes = token_data.get("scope") + if not isinstance(raw_scopes, str) or not raw_scopes.strip(): + logger.debug("OAuth token response omitted a usable scope field") + return set() + granted = set(raw_scopes.split()) + return set(REQUIRED_OAUTH_SCOPES) - granted + + def _cleanup_expired_states() -> None: """Remove expired OAuth states to prevent memory growth.""" now = time.time() @@ -63,16 +82,15 @@ async def oauth_login(request: Request) -> RedirectResponse: "expires_at": time.time() + _OAUTH_STATE_TTL, } - # Build authorization URL + # Build authorization URL. We no longer suggest a default `orgIds` — + # users no longer need to join the ML Agent Explorers org to use the + # app, and HF Jobs are billed per-namespace via credits. params = { "client_id": OAUTH_CLIENT_ID, "redirect_uri": get_redirect_uri(request), - "scope": "openid profile read-repos write-repos contribute-repos manage-repos inference-api jobs write-discussions", + "scope": " ".join(OAUTH_SCOPES), "response_type": "code", "state": state, - "orgIds": os.environ.get( - "HF_OAUTH_ORG_ID", "698dbf55845d85df163175f1" - ), # ml-agent-explorers } auth_url = f"{OPENID_PROVIDER_URL}/oauth/authorize?{urlencode(params)}" @@ -120,6 +138,15 @@ async def oauth_callback( status_code=500, detail="Token exchange succeeded but no access_token was returned.", ) + missing_scopes = _missing_required_scopes(token_data) + if missing_scopes: + raise HTTPException( + status_code=403, + detail=( + "OAuth token is missing required scopes: " + + ", ".join(sorted(missing_scopes)) + ), + ) # Fetch user info (optional — failure is not fatal) async with httpx.AsyncClient() as client: @@ -145,6 +172,15 @@ async def oauth_callback( max_age=3600 * 24 * 7, # 7 days path="/", ) + response.set_cookie( + key=OAUTH_SCOPE_COOKIE, + value=oauth_scope_fingerprint(OAUTH_SCOPES), + httponly=True, + secure=is_production, + samesite="lax", + max_age=3600 * 24 * 7, + path="/", + ) return response @@ -153,6 +189,7 @@ async def logout() -> RedirectResponse: """Log out the user by clearing the auth cookie.""" response = RedirectResponse(url="/") response.delete_cookie(key="hf_access_token", path="/") + response.delete_cookie(key=OAUTH_SCOPE_COOKIE, path="/") return response @@ -168,21 +205,4 @@ async def get_me(user: dict = Depends(get_current_user)) -> dict: Uses the shared auth dependency which handles cookie + Bearer token. """ - return user - - -ORG_NAME = "ml-agent-explorers" - - -@router.get("/org-membership") -async def org_membership( - request: Request, user: dict = Depends(get_current_user) -) -> dict: - """Check if the authenticated user belongs to the ml-agent-explorers org.""" - if not AUTH_ENABLED: - return {"is_member": True} - token = request.cookies.get("hf_access_token") or "" - if not token: - return {"is_member": False} - is_member = await check_org_membership(token, ORG_NAME) - return {"is_member": is_member} + return {key: value for key, value in user.items() if not key.startswith("_")} diff --git a/backend/session_manager.py b/backend/session_manager.py index 7293f9cf3..260eb8ffe 100644 --- a/backend/session_manager.py +++ b/backend/session_manager.py @@ -1,21 +1,48 @@ """Session manager for handling multiple concurrent agent sessions.""" import asyncio +import json import logging +import os import uuid from dataclasses import dataclass, field -from datetime import datetime +from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any, Optional from agent.config import load_config from agent.core.agent_loop import process_submission +from agent.core.model_ids import ( + GLM_52_MODEL_ID, + is_known_router_model_id, + strip_huggingface_model_prefix, +) from agent.core.session import Event, OpType, Session +from agent.core.session_persistence import get_session_store from agent.core.tools import ToolRouter +from agent.core.usage_thresholds import ( + USAGE_THRESHOLD_TOOL_NAME, + USAGE_WARNING_FIRST_THRESHOLD_USD, + is_usage_threshold_pending, + next_usage_warning_threshold, + normalize_usage_threshold, + usage_threshold_pending_to_tool, +) +from agent.core.yolo_budget import ( + YOLO_BUDGET_TOOL_NAME, + is_yolo_budget_pending, + request_yolo_budget_exceeded_approval, + seed_session_spend, + session_spend_usd, + yolo_budget_pending_to_tool, +) +from agent.messaging.gateway import NotificationGateway # Get project root (parent of backend directory) PROJECT_ROOT = Path(__file__).parent.parent -DEFAULT_CONFIG_PATH = str(PROJECT_ROOT / "configs" / "main_agent_config.json") +DEFAULT_CONFIG_PATH = str(PROJECT_ROOT / "configs" / "frontend_agent_config.json") +USAGE_WARNING_SPEND_CACHE_TTL_SECONDS = 30.0 +USAGE_BILLING_REFRESH_TIMEOUT_SECONDS = 2.0 # These dataclasses match agent/main.py structure @@ -41,9 +68,8 @@ class Submission: class EventBroadcaster: """Reads from the agent's event queue and fans out to SSE subscribers. - Events that arrive when no subscribers are listening are discarded. - With SSE each turn is a separate request, so there is no reconnect - scenario that would need buffered replay. + Events that arrive when no subscribers are listening are discarded by + this in-memory fanout. Durable replay is handled by session_persistence. """ def __init__(self, event_queue: asyncio.Queue): @@ -67,7 +93,11 @@ async def run(self) -> None: while True: try: event: Event = await self._source.get() - msg = {"event_type": event.event_type, "data": event.data} + msg = { + "event_type": event.event_type, + "data": event.data, + "seq": event.seq, + } for q in self._subscribers.values(): await q.put(msg) except asyncio.CancelledError: @@ -85,16 +115,59 @@ class AgentSession: tool_router: ToolRouter submission_queue: asyncio.Queue user_id: str = "dev" # Owner of this session + hf_username: str | None = None # HF namespace used for personal trace uploads hf_token: str | None = None # User's HF OAuth token for tool execution + user_plan: str | None = None # Active HF account plan for plan-aware agent CTAs task: asyncio.Task | None = None created_at: datetime = field(default_factory=datetime.utcnow) + # Last genuine activity (submit/turn-start/turn-finish/direct user write). + # Drives the idle reaper. Defaults to load time so a freshly-restored but + # untouched session isn't reaped for a full idle window. + last_active_at: datetime = field(default_factory=datetime.utcnow) is_active: bool = True is_processing: bool = False # True while a submission is being executed + # Set under the lock by the reaper while tearing this session down. Blocks + # submit() from enqueueing onto a session that's being evicted. + is_reaping: bool = False broadcaster: Any = None - # True once this session has been counted against the user's daily - # Claude quota. Guards double-counting when the user re-selects an - # Anthropic model mid-session. - claude_counted: bool = False + title: str | None = None + usage_window_started_at: datetime | None = None + inference_billing_session_id: str | None = None + usage_warning_next_threshold_usd: float = USAGE_WARNING_FIRST_THRESHOLD_USD + usage_warning_spend_cache: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.usage_window_started_at is None: + self.usage_window_started_at = self.created_at + if not self.inference_billing_session_id or not _is_uuid( + self.inference_billing_session_id + ): + self.inference_billing_session_id = new_inference_billing_session_id( + self.session_id, + self.usage_window_started_at, + ) + try: + self.session.inference_billing_session_id = ( + self.inference_billing_session_id + ) + except AttributeError: + pass + + +def new_inference_billing_session_id( + session_id: str, # noqa: ARG001 - kept for a stable call signature. + started_at: datetime | None = None, # noqa: ARG001 - kept for a stable call signature. +) -> str: + """Return a Router billing session ID scoped to one visible usage window.""" + return str(uuid.uuid4()) + + +def _is_uuid(value: str) -> bool: + try: + uuid.UUID(value) + except ValueError: + return False + return True class SessionCapacityError(Exception): @@ -112,6 +185,19 @@ def __init__(self, message: str, error_type: str = "global") -> None: # and per-request overhead. MAX_SESSIONS: int = 200 MAX_SESSIONS_PER_USER: int = 10 +DEFAULT_YOLO_COST_CAP_USD: float = 5.0 +SANDBOX_SHUTDOWN_CLEANUP_CONCURRENCY: int = 10 +SANDBOX_SHUTDOWN_CLEANUP_TIMEOUT_S: float = 60.0 + +# ── Idle-session reaper ───────────────────────────────────────────── +# A live session idle ≥ REAPER_IDLE_MINUTES with no in-flight work has its +# sandbox + RAM released and is evicted from the live pool, while staying +# fully resumable from Mongo (it reappears as a normal chat, never "ended"). +# This frees both the global pool and the user's concurrent slots. +REAPER_IDLE_MINUTES: float = float(os.environ.get("REAPER_IDLE_MINUTES", "15")) +REAPER_INTERVAL_S: float = float(os.environ.get("REAPER_INTERVAL_S", "300")) +REAP_TEARDOWN_TIMEOUT_S: float = float(os.environ.get("REAP_TEARDOWN_TIMEOUT_S", "30")) +REAPER_IDLE = timedelta(minutes=REAPER_IDLE_MINUTES) class SessionManager: @@ -119,22 +205,1057 @@ class SessionManager: def __init__(self, config_path: str | None = None) -> None: self.config = load_config(config_path or DEFAULT_CONFIG_PATH) + normalized_default = strip_huggingface_model_prefix(self.config.model_name) + if normalized_default: + self.config.model_name = normalized_default + self.messaging_gateway = NotificationGateway(self.config.messaging) self.sessions: dict[str, AgentSession] = {} self._lock = asyncio.Lock() + self.persistence_store = None + # In-flight create_session calls that have passed the capacity check + # but not yet inserted their session. Counted alongside + # active_session_count to hard-cap the global pool against the + # check-then-create race (see create_session). + self._pending_creates: int = 0 + self._reaper_task: asyncio.Task | None = None + + async def start(self) -> None: + """Start shared background resources.""" + self.persistence_store = get_session_store() + await self.persistence_store.init() + await self.messaging_gateway.start() + self._reaper_task = asyncio.create_task(self._reaper_loop()) + + async def close(self) -> None: + """Flush and close shared background resources.""" + if self._reaper_task is not None: + self._reaper_task.cancel() + try: + await self._reaper_task + except asyncio.CancelledError: + pass + self._reaper_task = None + await self._cleanup_all_sandboxes_on_close() + await self.messaging_gateway.close() + if self.persistence_store is not None: + await self.persistence_store.close() + + def _store(self): + if self.persistence_store is None: + self.persistence_store = get_session_store() + return self.persistence_store def _count_user_sessions(self, user_id: str) -> int: """Count active sessions owned by a specific user.""" return sum( - 1 - for s in self.sessions.values() - if s.user_id == user_id and s.is_active + 1 for s in self.sessions.values() if s.user_id == user_id and s.is_active + ) + + @staticmethod + def _touch(agent_session: "AgentSession") -> None: + """Stamp genuine activity so the idle reaper's clock resets. + + Call on real user/agent activity (submit, turn start/finish, direct + user-initiated writes) — never on passive reads or hydration, which + would keep an otherwise-idle session alive forever. + """ + agent_session.last_active_at = datetime.utcnow() + + @staticmethod + def _model_from_saved_metadata( + model: str | None, + ) -> str: + normalized = strip_huggingface_model_prefix(model) + if normalized and is_known_router_model_id(normalized): + return normalized + + fallback_model = GLM_52_MODEL_ID + logger.warning( + "Saved session model %r failed validation; using %r", + model, + fallback_model, + ) + return fallback_model + + def _create_session_sync( + self, + *, + session_id: str, + user_id: str, + hf_username: str | None, + hf_token: str | None, + user_plan: str | None, + model: str | None, + event_queue: asyncio.Queue, + notification_destinations: list[str] | None = None, + ) -> tuple[ToolRouter, Session]: + """Build blocking per-session resources in a worker thread.""" + import time as _time + + t0 = _time.monotonic() + tool_router = ToolRouter(self.config.mcpServers, hf_token=hf_token) + # Deep-copy config so each session's model switches independently — + # tab A picking GLM doesn't flip tab B off the default model. + session_config = self.config.model_copy(deep=True) + normalized_model = strip_huggingface_model_prefix(model) + if normalized_model: + session_config.model_name = normalized_model + session = Session( + event_queue=event_queue, + config=session_config, + tool_router=tool_router, + hf_token=hf_token, + user_id=user_id, + hf_username=hf_username, + user_plan=user_plan, + notification_gateway=self.messaging_gateway, + notification_destinations=notification_destinations or [], + session_id=session_id, + persistence_store=self._store(), + ) + t1 = _time.monotonic() + logger.info("Session initialized in %.2fs", t1 - t0) + return tool_router, session + + def _serialize_messages(self, session: Session) -> list[dict[str, Any]]: + return [msg.model_dump(mode="json") for msg in session.context_manager.items] + + def _serialize_pending_approval(self, session: Session) -> list[dict[str, Any]]: + pending = session.pending_approval or {} + if is_usage_threshold_pending(pending) or is_yolo_budget_pending(pending): + return [dict(pending)] + tool_calls = pending.get("tool_calls") or [] + serialized: list[dict[str, Any]] = [] + for tc in tool_calls: + if hasattr(tc, "model_dump"): + serialized.append(tc.model_dump(mode="json")) + elif isinstance(tc, dict): + serialized.append(tc) + return serialized + + @staticmethod + def _pending_tools_for_api(session: Session) -> list[dict[str, Any]] | None: + pending = session.pending_approval or {} + if is_usage_threshold_pending(pending): + return [usage_threshold_pending_to_tool(pending)] + if is_yolo_budget_pending(pending): + return [yolo_budget_pending_to_tool(pending)] + tool_calls = pending.get("tool_calls") or [] + if not tool_calls: + return None + result: list[dict[str, Any]] = [] + for tc in tool_calls: + try: + args = json.loads(tc.function.arguments) + except (json.JSONDecodeError, AttributeError, TypeError): + args = {} + result.append( + { + "tool": getattr(tc.function, "name", None), + "tool_call_id": getattr(tc, "id", None), + "arguments": args, + } + ) + return result + + def _restore_pending_approval( + self, session: Session, pending_approval: list[dict[str, Any]] | None + ) -> None: + if not pending_approval: + session.pending_approval = None + return + first = pending_approval[0] + if isinstance(first, dict) and first.get("kind") in { + USAGE_THRESHOLD_TOOL_NAME, + YOLO_BUDGET_TOOL_NAME, + }: + session.pending_approval = dict(first) + return + from litellm import ChatCompletionMessageToolCall as ToolCall + + restored = [] + for raw in pending_approval: + try: + if "function" in raw: + restored.append(ToolCall(**raw)) + else: + restored.append( + ToolCall( + id=raw["tool_call_id"], + type="function", + function={ + "name": raw["tool"], + "arguments": json.dumps(raw.get("arguments") or {}), + }, + ) + ) + except Exception as e: + logger.warning("Dropping malformed pending approval: %s", e) + session.pending_approval = {"tool_calls": restored} if restored else None + + @staticmethod + def _pending_docs_for_api( + pending_approval: list[dict[str, Any]] | None, + ) -> list[dict[str, Any]] | None: + if not pending_approval: + return None + first = pending_approval[0] + if isinstance(first, dict) and first.get("kind") == USAGE_THRESHOLD_TOOL_NAME: + return [usage_threshold_pending_to_tool(first)] + if isinstance(first, dict) and first.get("kind") == YOLO_BUDGET_TOOL_NAME: + return [yolo_budget_pending_to_tool(first)] + result: list[dict[str, Any]] = [] + for raw in pending_approval: + if "function" in raw: + function = raw.get("function") or {} + try: + args = json.loads(function.get("arguments") or "{}") + except (json.JSONDecodeError, TypeError): + args = {} + result.append( + { + "tool": function.get("name"), + "tool_call_id": raw.get("id"), + "arguments": args, + } + ) + elif {"tool", "tool_call_id"}.issubset(raw): + result.append( + { + "tool": raw.get("tool"), + "tool_call_id": raw.get("tool_call_id"), + "arguments": raw.get("arguments") or {}, + } + ) + return result or None + + @staticmethod + def _runtime_state(agent_session: AgentSession) -> str: + if agent_session.session.pending_approval: + return "waiting_approval" + if agent_session.is_processing: + return "processing" + if not agent_session.is_active: + return "ended" + return "idle" + + @staticmethod + def _auto_approval_summary(session: Session) -> dict[str, Any]: + if hasattr(session, "auto_approval_policy_summary"): + return session.auto_approval_policy_summary() + cap = getattr(session, "auto_approval_cost_cap_usd", None) + estimated = float( + getattr(session, "auto_approval_estimated_spend_usd", 0.0) or 0.0 + ) + remaining = None if cap is None else round(max(0.0, float(cap) - estimated), 4) + return { + "enabled": bool(getattr(session, "auto_approval_enabled", False)), + "cost_cap_usd": cap, + "estimated_spend_usd": round(estimated, 4), + "remaining_usd": remaining, + } + + def _install_usage_threshold_checker(self, agent_session: AgentSession) -> None: + threshold = normalize_usage_threshold( + getattr( + agent_session.session, + "usage_warning_next_threshold_usd", + agent_session.usage_warning_next_threshold_usd, + ) + ) + agent_session.usage_warning_next_threshold_usd = threshold + agent_session.session.usage_warning_next_threshold_usd = threshold + + async def _checker(payload: dict[str, Any]) -> bool: + return await self._maybe_request_usage_threshold_approval( + agent_session.session_id, + payload, + ) + + agent_session.session.usage_threshold_checker = _checker + + def _install_yolo_budget_checker(self, agent_session: AgentSession) -> None: + async def _checker(payload: dict[str, Any]) -> bool: + return await self._maybe_request_yolo_budget_approval( + agent_session.session_id, + payload, + ) + + agent_session.session.yolo_budget_checker = _checker + + @staticmethod + def _set_inference_billing_session_id( + agent_session: AgentSession, + inference_billing_session_id: str, + ) -> None: + agent_session.inference_billing_session_id = inference_billing_session_id + try: + agent_session.session.inference_billing_session_id = ( + inference_billing_session_id + ) + except AttributeError: + pass + + @staticmethod + def _usage_spend_from_response(response: dict[str, Any]) -> tuple[float, str]: + def coerce_spend(value: Any) -> float | None: + if isinstance(value, bool) or value is None: + return None + try: + return max(0.0, float(value)) + except (TypeError, ValueError): + return None + + hf_account = response.get("hf_account") + session_bucket = response.get("session") + if isinstance(hf_account, dict): + current_session = hf_account.get("current_session") + if isinstance(current_session, dict): + spend = coerce_spend( + current_session.get("inference_providers_usd") + if "inference_providers_usd" in current_session + else current_session.get("total_usd") + ) + if spend is not None: + if isinstance(session_bucket, dict): + for key in ( + "hf_jobs_estimated_usd", + "sandbox_estimated_usd", + ): + spend += coerce_spend(session_bucket.get(key)) or 0.0 + return spend, "hf_billing_current_session" + + if isinstance(session_bucket, dict): + spend = coerce_spend(session_bucket.get("total_usd")) + if spend is not None: + return spend, "app_telemetry_session" + return 0.0, "app_telemetry_session" + + async def _current_session_usage_spend( + self, + agent_session: AgentSession, + *, + use_cache: bool = True, + ) -> tuple[float, str]: + now = datetime.now(UTC) + cache = agent_session.usage_warning_spend_cache + cache_expires_at = cache.get("expires_at") + if ( + use_cache + and isinstance(cache_expires_at, datetime) + and cache_expires_at > now + ): + return ( + float(cache.get("spend_usd") or 0.0), + str(cache.get("billing_source") or "app_telemetry_session"), + ) + + from usage import build_usage_response + + response = await build_usage_response( + self, + user_id=agent_session.user_id, + hf_token=agent_session.hf_token, + session_id=agent_session.session_id, + timezone_name="UTC", + ) + spend, billing_source = self._usage_spend_from_response(response) + agent_session.usage_warning_spend_cache = { + "spend_usd": spend, + "billing_source": billing_source, + "expires_at": now + + timedelta(seconds=USAGE_WARNING_SPEND_CACHE_TTL_SECONDS), + } + return spend, billing_source + + @staticmethod + def _fallback_hf_billing_snapshot(error: str) -> dict[str, Any]: + return { + "billing_scope": "account_window_delta", + "hf_billing": { + "source": "hf_billing_usage_v2", + "available": False, + "error": error, + "current_session": None, + }, + } + + async def refresh_session_usage_metrics( + self, + agent_session: AgentSession, + *, + error_code: str = "billing_snapshot_error", + billing_timeout_s: float | None = USAGE_BILLING_REFRESH_TIMEOUT_SECONDS, + ) -> dict[str, Any]: + """Refresh the dataset usage snapshot stored on the runtime session.""" + from agent.core.usage_metrics import ( + normalize_hf_billing_snapshot, + summarize_usage_events, + ) + from usage import build_hf_billing_snapshot + + session = agent_session.session + try: + billing_snapshot = build_hf_billing_snapshot( + self, + hf_token=agent_session.hf_token or getattr(session, "hf_token", None), + session_id=agent_session.session_id, + timezone_name="UTC", + ) + if billing_timeout_s is not None and billing_timeout_s > 0: + hf_billing_snapshot = await asyncio.wait_for( + billing_snapshot, + timeout=billing_timeout_s, + ) + else: + hf_billing_snapshot = await billing_snapshot + except TimeoutError: + logger.debug( + "HF billing snapshot refresh timed out for %s after %.2fs", + agent_session.session_id, + billing_timeout_s or 0, + ) + hf_billing_snapshot = self._fallback_hf_billing_snapshot(error_code) + except Exception as e: + logger.debug( + "HF billing snapshot refresh failed for %s: %s", + agent_session.session_id, + e, + ) + hf_billing_snapshot = self._fallback_hf_billing_snapshot(error_code) + + hf_billing_snapshot = normalize_hf_billing_snapshot(hf_billing_snapshot) + session.usage_hf_billing_snapshot = hf_billing_snapshot + metrics = summarize_usage_events( + getattr(session, "logged_events", []) or [], + session_id=agent_session.session_id, + hf_billing_snapshot=hf_billing_snapshot, + ) + session.usage_metrics = metrics + return metrics + + @staticmethod + def _runtime_session_usage_spend(agent_session: AgentSession) -> float: + from usage import aggregate_usage_events, event_created_at + + window_start = agent_session.usage_window_started_at + if isinstance(window_start, datetime): + if window_start.tzinfo is None: + window_start = window_start.replace(tzinfo=UTC) + else: + window_start = window_start.astimezone(UTC) + events = [] + for raw_event in getattr(agent_session.session, "logged_events", []) or []: + if raw_event.get("event_type") not in { + "llm_call", + "hf_job_complete", + "sandbox_create", + "sandbox_destroy", + }: + continue + if isinstance(window_start, datetime): + created_at = event_created_at(raw_event, timezone_name="UTC") + if created_at is not None and created_at < window_start: + continue + events.append(raw_event) + bucket = aggregate_usage_events( + events, + session_id=agent_session.session_id, + ) + return float(bucket.get("total_usd") or 0.0) + + async def _maybe_request_usage_threshold_approval( + self, + session_id: str, + continuation_payload: dict[str, Any], + ) -> bool: + agent_session = self.sessions.get(session_id) + if not agent_session or not agent_session.is_active: + return False + + session = agent_session.session + if session.pending_approval: + return False + + threshold = normalize_usage_threshold( + getattr( + session, + "usage_warning_next_threshold_usd", + agent_session.usage_warning_next_threshold_usd, + ) + ) + force_check = bool(continuation_payload.get("force_check")) + local_spend = self._runtime_session_usage_spend(agent_session) + if not force_check and local_spend < threshold: + return False + + current_spend, billing_source = await self._current_session_usage_spend( + agent_session, + use_cache=not force_check, + ) + if current_spend < threshold: + return False + + next_threshold = next_usage_warning_threshold(current_spend, threshold) + tool_call_id = f"usage-threshold-{uuid.uuid4().hex[:10]}" + pending: dict[str, Any] = { + "kind": USAGE_THRESHOLD_TOOL_NAME, + "tool_call_id": tool_call_id, + "threshold_usd": round(threshold, 4), + "current_spend_usd": round(current_spend, 6), + "next_threshold_usd": next_threshold, + "billing_source": billing_source, + "continuation": continuation_payload.get("continuation") + or "continue_agent", + "history_size": int( + continuation_payload.get("history_size") + or len(session.context_manager.items) + ), + } + final_response = continuation_payload.get("final_response") + if isinstance(final_response, str): + pending["final_response"] = final_response + + session.pending_approval = pending + self._touch(agent_session) + await session.send_event( + Event( + event_type="approval_required", + data={ + "tools": [usage_threshold_pending_to_tool(pending)], + "count": 1, + "usage_threshold": True, + }, + ) + ) + return True + + async def _maybe_request_yolo_budget_approval( + self, + session_id: str, + payload: dict[str, Any], + ) -> bool: + agent_session = self.sessions.get(session_id) + if not agent_session or not agent_session.is_active: + return False + + session = agent_session.session + if session.pending_approval: + return False + if not bool(getattr(session, "auto_approval_enabled", False)): + return False + cap = getattr(session, "auto_approval_cost_cap_usd", None) + if cap is None: + return False + try: + cap_usd = max(0.0, float(cap)) + except (TypeError, ValueError): + return False + + current_spend, billing_source = await self._current_session_usage_spend( + agent_session, + use_cache=False, ) + raw_observed_cost = payload.get("observed_cost_usd") + observed_cost = ( + max(0.0, float(raw_observed_cost)) + if isinstance(raw_observed_cost, (int, float)) + and not isinstance(raw_observed_cost, bool) + else 0.0 + ) + previous_ledger_spend = session_spend_usd(session) + seed_session_spend(session, current_spend) + ledger_spend = session_spend_usd(session) + effective_spend = max(current_spend, ledger_spend) + if effective_spend < cap_usd: + if ledger_spend != previous_ledger_spend or observed_cost > 0: + self._touch(agent_session) + await session.send_event( + Event( + event_type="session_update", + data={ + "session_id": session_id, + "auto_approval": self._auto_approval_summary(session), + }, + ) + ) + return False + + spend_kind = str(payload.get("spend_kind") or "session usage") + final_response = payload.get("final_response") + created = await request_yolo_budget_exceeded_approval( + session, + spend_kind=spend_kind, + current_spend_usd=effective_spend, + cap_usd=cap_usd, + billing_source=billing_source, + continuation=payload.get("continuation"), + final_response=final_response if isinstance(final_response, str) else None, + history_size=payload.get("history_size"), + reason=( + "YOLO cap paused session usage after " + f"{spend_kind}: current session spend ${effective_spend:.2f} " + f"has reached the ${cap_usd:.2f} cap." + ), + ) + if created: + self._touch(agent_session) + return created + + async def _start_agent_session( + self, + *, + agent_session: AgentSession, + event_queue: asyncio.Queue, + tool_router: ToolRouter, + ) -> AgentSession: + async with self._lock: + existing = self.sessions.get(agent_session.session_id) + if existing: + return existing + self.sessions[agent_session.session_id] = agent_session + + task = asyncio.create_task( + self._run_session( + agent_session.session_id, + agent_session.submission_queue, + event_queue, + tool_router, + ) + ) + agent_session.task = task + return agent_session + + @staticmethod + def _start_cpu_sandbox_preload(agent_session: AgentSession) -> None: + """Kick off a best-effort cpu-basic sandbox for the session.""" + try: + from agent.tools.sandbox_tool import start_cpu_sandbox_preload + + start_cpu_sandbox_preload(agent_session.session) + except Exception as e: + logger.warning( + "Failed to start CPU sandbox preload for %s: %s", + agent_session.session_id, + e, + ) + + @staticmethod + def _can_access_session(agent_session: AgentSession, user_id: str) -> bool: + return ( + user_id == "dev" + or agent_session.user_id == "dev" + or agent_session.user_id == user_id + ) + + @staticmethod + def _update_hf_identity( + agent_session: AgentSession, + *, + hf_token: str | None, + hf_username: str | None, + user_plan: str | None = None, + ) -> None: + if hf_token: + agent_session.hf_token = hf_token + agent_session.session.hf_token = hf_token + if hf_username: + agent_session.hf_username = hf_username + agent_session.session.hf_username = hf_username + if user_plan is not None: + agent_session.user_plan = user_plan + agent_session.session.user_plan = user_plan + + @staticmethod + def _has_active_sandbox_preload(agent_session: AgentSession) -> bool: + task = getattr(agent_session.session, "sandbox_preload_task", None) + return bool(task and not task.done()) + + @staticmethod + def _preload_failed_for_missing_hf_token(agent_session: AgentSession) -> bool: + error = getattr(agent_session.session, "sandbox_preload_error", None) + return isinstance(error, str) and error.startswith("No HF token available") + + def _restart_cpu_preload_if_token_recovered( + self, + agent_session: AgentSession, + *, + preload_sandbox: bool, + ) -> None: + if not preload_sandbox: + return + session = agent_session.session + if getattr(session, "sandbox", None): + return + if self._has_active_sandbox_preload(agent_session): + return + if not (agent_session.hf_token or getattr(session, "hf_token", None)): + return + + if not self._preload_failed_for_missing_hf_token(agent_session): + return + + session.sandbox_preload_error = None + session.sandbox_preload_task = None + session.sandbox_preload_cancel_event = None + self._start_cpu_sandbox_preload(agent_session) + + async def _clear_persisted_sandbox_metadata(self, session_id: str) -> None: + try: + await self._store().update_session_fields( + session_id, + sandbox_space_id=None, + sandbox_hardware=None, + sandbox_owner=None, + sandbox_created_at=None, + sandbox_status="destroyed", + ) + except Exception as e: + logger.warning("Failed to clear sandbox metadata for %s: %s", session_id, e) + + async def _cleanup_persisted_sandbox( + self, + session_id: str, + metadata: dict[str, Any], + *, + hf_token: str | None, + ) -> None: + """Delete a sandbox recorded by a previous backend process, if any.""" + space_id = metadata.get("sandbox_space_id") + if not isinstance(space_id, str) or not space_id: + return + if metadata.get("sandbox_status") == "destroyed": + return + + tokens: list[tuple[str, str]] = [] + seen: set[str] = set() + for label, token in ( + ("user", hf_token), + ("admin", os.environ.get("HF_ADMIN_TOKEN")), + ): + if token and token not in seen: + tokens.append((label, token)) + seen.add(token) + + if not tokens: + logger.warning( + "Cannot clean persisted sandbox %s for session %s: no HF token available", + space_id, + session_id, + ) + return + + last_err: Exception | None = None + for label, token in tokens: + try: + from huggingface_hub import HfApi + + api = HfApi(token=token) + await asyncio.to_thread( + api.delete_repo, + repo_id=space_id, + repo_type="space", + ) + logger.info( + "Deleted persisted sandbox %s for session %s with %s token", + space_id, + session_id, + label, + ) + await self._clear_persisted_sandbox_metadata(session_id) + return + except Exception as e: + status_code = getattr(getattr(e, "response", None), "status_code", None) + if status_code == 404: + logger.info( + "Persisted sandbox %s for session %s is already gone", + space_id, + session_id, + ) + await self._clear_persisted_sandbox_metadata(session_id) + return + last_err = e + + logger.warning( + "Failed to delete persisted sandbox %s for session %s: %s", + space_id, + session_id, + last_err, + ) + + async def persist_session_snapshot( + self, + agent_session: AgentSession, + *, + runtime_state: str | None = None, + status: str = "active", + raise_on_error: bool = False, + ) -> None: + """Persist the current runtime context snapshot. + + Best-effort by default: a disabled store is a no-op and write failures + are swallowed. Pass ``raise_on_error=True`` when the caller must know + the snapshot was durably written (e.g. the reaper, which only evicts a + session after confirming it stayed resumable) — then a disabled store + or a write failure raises instead of silently dropping the snapshot. + """ + store = self._store() + if not getattr(store, "enabled", False): + if raise_on_error: + raise RuntimeError("persistence store is disabled") + return + try: + await store.save_snapshot( + session_id=agent_session.session_id, + user_id=agent_session.user_id, + model=agent_session.session.config.model_name, + title=agent_session.title, + messages=self._serialize_messages(agent_session.session), + runtime_state=runtime_state or self._runtime_state(agent_session), + status=status, + turn_count=agent_session.session.turn_count, + pending_approval=self._serialize_pending_approval( + agent_session.session + ), + created_at=agent_session.created_at, + usage_window_started_at=agent_session.usage_window_started_at, + inference_billing_session_id=( + agent_session.inference_billing_session_id + ), + notification_destinations=list( + agent_session.session.notification_destinations + ), + auto_approval_enabled=bool( + getattr(agent_session.session, "auto_approval_enabled", False) + ), + auto_approval_cost_cap_usd=getattr( + agent_session.session, "auto_approval_cost_cap_usd", None + ), + auto_approval_estimated_spend_usd=float( + getattr( + agent_session.session, + "auto_approval_estimated_spend_usd", + 0.0, + ) + or 0.0 + ), + usage_warning_next_threshold_usd=normalize_usage_threshold( + getattr( + agent_session.session, + "usage_warning_next_threshold_usd", + agent_session.usage_warning_next_threshold_usd, + ) + ), + raise_on_error=raise_on_error, + ) + except Exception as e: + if raise_on_error: + raise + logger.warning( + "Failed to persist snapshot for %s: %s", + agent_session.session_id, + e, + ) + + async def ensure_session_loaded( + self, + session_id: str, + user_id: str, + hf_token: str | None = None, + hf_username: str | None = None, + user_plan: str | None = None, + preload_sandbox: bool = True, + ) -> AgentSession | None: + """Return a live runtime session, lazily restoring it from Mongo.""" + async with self._lock: + existing = self.sessions.get(session_id) + if existing: + if self._can_access_session(existing, user_id): + self._update_hf_identity( + existing, + hf_token=hf_token, + hf_username=hf_username, + user_plan=user_plan, + ) + self._install_usage_threshold_checker(existing) + self._install_yolo_budget_checker(existing) + self._restart_cpu_preload_if_token_recovered( + existing, + preload_sandbox=preload_sandbox, + ) + return existing + return None + + store = self._store() + loaded = await store.load_session(session_id) + if not loaded: + return None + + async with self._lock: + existing = self.sessions.get(session_id) + if existing: + if self._can_access_session(existing, user_id): + self._update_hf_identity( + existing, + hf_token=hf_token, + hf_username=hf_username, + user_plan=user_plan, + ) + self._install_usage_threshold_checker(existing) + self._install_yolo_budget_checker(existing) + self._restart_cpu_preload_if_token_recovered( + existing, + preload_sandbox=preload_sandbox, + ) + return existing + return None + + meta = loaded.get("metadata") or {} + owner = str(meta.get("user_id") or "") + if user_id != "dev" and owner != "dev" and owner != user_id: + return None + + await self._cleanup_persisted_sandbox( + session_id, + meta, + hf_token=hf_token, + ) + + from litellm import Message + + model = self._model_from_saved_metadata( + meta.get("model") or self.config.model_name, + ) + event_queue: asyncio.Queue = asyncio.Queue() + submission_queue: asyncio.Queue = asyncio.Queue() + tool_router, session = await asyncio.to_thread( + self._create_session_sync, + session_id=session_id, + user_id=owner or user_id, + hf_username=hf_username, + hf_token=hf_token, + user_plan=user_plan, + model=model, + event_queue=event_queue, + notification_destinations=meta.get("notification_destinations") or [], + ) + + restored_messages: list[Message] = [] + for raw in loaded.get("messages") or []: + if not isinstance(raw, dict) or raw.get("role") == "system": + continue + try: + restored_messages.append(Message.model_validate(raw)) + except Exception as e: + logger.warning("Dropping malformed restored message: %s", e) + if restored_messages: + # Keep the freshly-rendered system prompt, then attach the durable + # non-system context so tools/date/user context stay current. + session.context_manager.items = [ + session.context_manager.items[0], + *restored_messages, + ] + + # If this session ever had a sandbox, its container did not survive the + # resume (a fresh, empty one is lazily recreated). Tell the agent so it + # recreates files/packages instead of assuming /app/train.py et al. still + # exist. Gated on sandbox_status so pure Q&A chats get no note. Mirrors + # the seed_from_summary note convention. + # + # Skip it when an approval is pending: the restored context ends with an + # assistant tool-call message awaiting results, so injecting a user + # message here would sit between the tool_calls and their results. On + # approval the real results get appended after the note, leaving them + # orphaned (the context manager stubs the "missing" result right after + # the assistant message) — which the provider rejects. The agent still + # learns the sandbox is empty when the approved tool runs against it. + if meta.get("sandbox_status") and not meta.get("pending_approval"): + session.context_manager.items.append( + Message( + role="user", + content=( + "[SYSTEM: This session was resumed and its sandbox was " + "reset. Any files, installed packages, or running " + "processes from earlier are gone — recreate what you " + "need before using the sandbox.]" + ), + ) + ) + + self._restore_pending_approval(session, meta.get("pending_approval") or []) + session.turn_count = int(meta.get("turn_count") or 0) + session.auto_approval_enabled = bool(meta.get("auto_approval_enabled", False)) + raw_cap = meta.get("auto_approval_cost_cap_usd") + session.auto_approval_cost_cap_usd = ( + float(raw_cap) if isinstance(raw_cap, int | float) else None + ) + session.auto_approval_estimated_spend_usd = float( + meta.get("auto_approval_estimated_spend_usd") or 0.0 + ) + session.usage_warning_next_threshold_usd = normalize_usage_threshold( + meta.get("usage_warning_next_threshold_usd") + ) + + created_at = meta.get("created_at") + if not isinstance(created_at, datetime): + created_at = datetime.utcnow() + usage_window_started_at = meta.get("usage_window_started_at") + if not isinstance(usage_window_started_at, datetime): + usage_window_started_at = created_at + inference_billing_session_id = meta.get("inference_billing_session_id") + if not isinstance(inference_billing_session_id, str) or not _is_uuid( + inference_billing_session_id + ): + inference_billing_session_id = new_inference_billing_session_id( + session_id, + usage_window_started_at, + ) + + agent_session = AgentSession( + session_id=session_id, + session=session, + tool_router=tool_router, + submission_queue=submission_queue, + user_id=owner or user_id, + hf_username=hf_username, + hf_token=hf_token, + user_plan=user_plan, + created_at=created_at, + usage_window_started_at=usage_window_started_at, + inference_billing_session_id=inference_billing_session_id, + usage_warning_next_threshold_usd=session.usage_warning_next_threshold_usd, + is_active=True, + is_processing=False, + title=meta.get("title"), + ) + self._install_usage_threshold_checker(agent_session) + self._install_yolo_budget_checker(agent_session) + started = await self._start_agent_session( + agent_session=agent_session, + event_queue=event_queue, + tool_router=tool_router, + ) + if started is not agent_session: + self._update_hf_identity( + started, + hf_token=hf_token, + hf_username=hf_username, + user_plan=user_plan, + ) + return started + if preload_sandbox: + self._start_cpu_sandbox_preload(agent_session) + logger.info("Restored session %s for user %s", session_id, owner or user_id) + return agent_session async def create_session( self, user_id: str = "dev", + hf_username: str | None = None, hf_token: str | None = None, + user_plan: str | None = None, model: str | None = None, + is_pro: bool | None = None, ) -> str: """Create a new agent session and return its ID. @@ -144,21 +1265,29 @@ async def create_session( Args: user_id: The ID of the user who owns this session. + hf_username: The HF username/namespace used for personal trace uploads. hf_token: The user's HF OAuth token, stored for tool execution. + user_plan: The active HF account plan used for plan-aware agent CTAs. model: Optional model override. When set, replaces ``model_name`` on the per-session config clone. None falls back to the config default. Raises: SessionCapacityError: If the server or user has reached the - maximum number of concurrent sessions. + maximum number of live sessions. """ # ── Capacity checks ────────────────────────────────────────── + # Reserve a global slot under the lock so concurrent creates can't all + # pass the check then over-admit past MAX_SESSIONS (the build + insert + # happen later, outside the lock). active_session_count won't reflect + # this session until _start_agent_session inserts it, so we count + # _pending_creates alongside it to close that gap. async with self._lock: active_count = self.active_session_count - if active_count >= MAX_SESSIONS: + projected = active_count + self._pending_creates + if projected >= MAX_SESSIONS: raise SessionCapacityError( - f"Server is at capacity ({active_count}/{MAX_SESSIONS} sessions). " + f"Server is at capacity ({projected}/{MAX_SESSIONS} sessions). " "Please try again later.", error_type="global", ) @@ -167,9 +1296,12 @@ async def create_session( if user_count >= MAX_SESSIONS_PER_USER: raise SessionCapacityError( f"You have reached the maximum of {MAX_SESSIONS_PER_USER} " - "concurrent sessions. Please close an existing session first.", + f"live sessions. Close an existing session, or wait " + f"{REAPER_IDLE_MINUTES:g} minutes after your last activity " + "for an idle session to be released.", error_type="per_user", ) + self._pending_creates += 1 session_id = str(uuid.uuid4()) @@ -177,50 +1309,88 @@ async def create_session( submission_queue: asyncio.Queue = asyncio.Queue() event_queue: asyncio.Queue = asyncio.Queue() - # Run blocking constructors in a thread to keep the event loop responsive. - # Without this, Session.__init__ → ContextManager → litellm.get_max_tokens() - # blocks all HTTP/SSE handling. - import time as _time + reserved = True + try: + # Run blocking constructors in a thread to keep the event loop responsive. + tool_router, session = await asyncio.to_thread( + self._create_session_sync, + session_id=session_id, + user_id=user_id, + hf_username=hf_username, + hf_token=hf_token, + user_plan=user_plan, + model=model, + event_queue=event_queue, + ) - def _create_session_sync(): - t0 = _time.monotonic() - tool_router = ToolRouter(self.config.mcpServers, hf_token=hf_token) - # Deep-copy config so each session's model switches independently — - # tab A picking GLM doesn't flip tab B off Claude. - session_config = self.config.model_copy(deep=True) - if model: - session_config.model_name = model - session = Session( - event_queue, config=session_config, tool_router=tool_router, + # Create wrapper + agent_session = AgentSession( + session_id=session_id, + session=session, + tool_router=tool_router, + submission_queue=submission_queue, + user_id=user_id, + hf_username=hf_username, hf_token=hf_token, + user_plan=user_plan, ) - t1 = _time.monotonic() - logger.info(f"Session initialized in {t1 - t0:.2f}s") - return tool_router, session + self._install_usage_threshold_checker(agent_session) + self._install_yolo_budget_checker(agent_session) - tool_router, session = await asyncio.to_thread(_create_session_sync) + await self._start_agent_session( + agent_session=agent_session, + event_queue=event_queue, + tool_router=tool_router, + ) + # The session is now in self.sessions, so active_session_count + # reflects it — release the reservation before the slower (and + # non-capacity) persistence + preload work. + async with self._lock: + self._pending_creates -= 1 + reserved = False - # Create wrapper - agent_session = AgentSession( - session_id=session_id, - session=session, - tool_router=tool_router, - submission_queue=submission_queue, - user_id=user_id, - hf_token=hf_token, - ) + await self.persist_session_snapshot(agent_session, runtime_state="idle") + self._start_cpu_sandbox_preload(agent_session) - async with self._lock: - self.sessions[session_id] = agent_session + if is_pro is not None and user_id and user_id != "dev": + await self._track_pro_status(agent_session, is_pro=is_pro) - # Start the agent loop task - task = asyncio.create_task( - self._run_session(session_id, submission_queue, event_queue, tool_router) - ) - agent_session.task = task + logger.info(f"Created session {session_id} for user {user_id}") + return session_id + finally: + # Build/start failed before the session was inserted — always + # release the reservation so a failed create can't permanently + # shrink the pool. + if reserved: + async with self._lock: + self._pending_creates -= 1 + + async def _track_pro_status( + self, agent_session: AgentSession, *, is_pro: bool + ) -> None: + """Update Mongo per-user Pro state and emit a one-shot conversion + event if the store reports a free→Pro transition. Best-effort: any + Mongo failure is swallowed so we never fail session creation on + telemetry.""" + store = self._store() + if not getattr(store, "enabled", False): + return + try: + result = await store.mark_pro_seen(agent_session.user_id, is_pro=is_pro) + except Exception as e: + logger.debug("mark_pro_seen failed: %s", e) + return + if not result or not result.get("converted"): + return + try: + from agent.core import telemetry - logger.info(f"Created session {session_id} for user {user_id}") - return session_id + await telemetry.record_pro_conversion( + agent_session.session, + first_seen_at=result.get("first_seen_at"), + ) + except Exception as e: + logger.debug("record_pro_conversion failed: %s", e) async def seed_from_summary(self, session_id: str, messages: list[dict]) -> int: """Rehydrate a session from cached prior messages via summarization. @@ -253,9 +1423,8 @@ async def seed_from_summary(self, session_id: str, messages: list[dict]) -> int: session = agent_session.session # Pass the real tool specs so the summarizer sees what the agent - # actually has — otherwise Anthropic's modify_params injects a - # dummy tool and the summarizer editorializes that the original - # tool calls were fabricated. + # actually has. Without them, the summarizer can editorialize that + # original tool calls were fabricated. tool_specs = None try: tool_specs = agent_session.tool_router.get_tool_specs_for_llm() @@ -269,6 +1438,8 @@ async def seed_from_summary(self, session_id: str, messages: list[dict]) -> int: max_tokens=4000, prompt=_RESTORE_PROMPT, tool_specs=tool_specs, + session=session, + kind="restore", ) except Exception as e: logger.error("Summary call failed during seed: %s", e) @@ -283,18 +1454,207 @@ async def seed_from_summary(self, session_id: str, messages: list[dict]) -> int: ), ) session.context_manager.items.append(seed) + self._touch(agent_session) + await self.persist_session_snapshot(agent_session, runtime_state="idle") return len(parsed) @staticmethod async def _cleanup_sandbox(session: Session) -> None: - """Delete the sandbox Space if one was created for this session.""" - sandbox = getattr(session, "sandbox", None) - if sandbox and getattr(sandbox, "_owns_space", False): + """Delete the sandbox Space if one was created for this session. + + Retries on transient failures (HF API 5xx, rate-limit, network blips) + with exponential backoff. A single missed delete = a permanently + orphaned Space, so the cost of an extra retry beats the alternative. + """ + from agent.tools.sandbox_tool import teardown_session_sandbox + + await teardown_session_sandbox(session) + + async def _cleanup_all_sandboxes_on_close(self) -> None: + """Best-effort sandbox cleanup for graceful backend shutdown.""" + async with self._lock: + agent_sessions = list(self.sessions.values()) + if not agent_sessions: + return + + semaphore = asyncio.Semaphore(SANDBOX_SHUTDOWN_CLEANUP_CONCURRENCY) + + async def _cleanup_one(agent_session: AgentSession) -> None: + async with semaphore: + try: + await self._cleanup_sandbox(agent_session.session) + except Exception as e: + logger.warning( + "Shutdown sandbox cleanup failed for %s: %s", + agent_session.session_id, + e, + ) + + tasks = [ + asyncio.create_task(_cleanup_one(agent_session)) + for agent_session in agent_sessions + ] + try: + await asyncio.wait_for( + asyncio.gather(*tasks, return_exceptions=True), + timeout=SANDBOX_SHUTDOWN_CLEANUP_TIMEOUT_S, + ) + except asyncio.TimeoutError: + logger.warning( + "Timed out after %.0fs cleaning up sandboxes on shutdown; " + "orphan sweeper will handle any stragglers", + SANDBOX_SHUTDOWN_CLEANUP_TIMEOUT_S, + ) + + async def _reaper_loop(self) -> None: + """Periodically release resources held by idle sessions. + + Modeled on EventBroadcaster.run: a long-lived task started in start() + and cancelled in close(). Per-sweep exceptions are swallowed so one bad + sweep never kills the loop. + """ + while True: + try: + await asyncio.sleep(REAPER_INTERVAL_S) + await self._reap_idle_sessions() + except asyncio.CancelledError: + break + except Exception as e: + logger.error("Idle-session reaper sweep failed: %s", e) + + async def _reap_idle_sessions(self) -> None: + """Select idle candidates under the lock, then tear each down. + + Candidates are non-dev sessions that are live, not processing, not + awaiting tool approval (those are "approve later", not idle — reaping + would destroy the sandbox the approved tool needs), and untouched for + the idle window. We only snapshot IDs under the lock; the actual + teardown in _reap_one re-acquires it, because tearing a session down + while holding the lock would deadlock (the lock is non-reentrant). + """ + # Reaping is only safe when sessions stay resumable from Mongo. With no + # store, eviction would destroy non-dev chats outright, so don't reap. + if not getattr(self._store(), "enabled", False): + return + + cutoff = datetime.utcnow() - REAPER_IDLE + async with self._lock: + candidates = [ + agent_session.session_id + for agent_session in self.sessions.values() + if agent_session.is_active + and not agent_session.is_processing + and not agent_session.is_reaping + and agent_session.user_id != "dev" + and not agent_session.session.pending_approval + and agent_session.last_active_at <= cutoff + ] + if not candidates: + return + + reaped = 0 + for session_id in candidates: try: - logger.info(f"Deleting sandbox {sandbox.space_id}...") - await asyncio.to_thread(sandbox.delete) + if await self._reap_one(session_id, cutoff): + reaped += 1 except Exception as e: - logger.warning(f"Failed to delete sandbox {sandbox.space_id}: {e}") + logger.warning("Failed to reap idle session %s: %s", session_id, e) + if reaped: + logger.info("Reaped %d idle session(s)", reaped) + + async def _reap_one(self, session_id: str, cutoff: datetime) -> bool: + """Tear down one idle session, leaving it resumable from Mongo. + + Re-checks every idle condition under the lock (a user may have become + active in the gap since selection), marks the session reaping, persists + a resumable snapshot outside the lock, then does one final locked + re-check before eviction. The runtime task is cancelled *outside* the + lock: its own ``finally`` frees the sandbox, and its identity-gated + persist no-ops because the session is already popped — so it can't + overwrite our resumable snapshot with ``"ended"`` and there's no + deadlock. Returns True if the session was reaped. + """ + async with self._lock: + agent_session = self.sessions.get(session_id) + if ( + agent_session is None + or not agent_session.is_active + or agent_session.is_processing + or agent_session.is_reaping + or agent_session.session.pending_approval + or agent_session.last_active_at > cutoff + or not agent_session.submission_queue.empty() + ): + return False + agent_session.is_reaping = True + + # Persist a resumable snapshot *before* eviction so a concurrent reopen + # reloads clean state. status="active" (never "ended") keeps it a normal + # chat in the sidebar. Do this outside the manager lock: Mongo writes can + # take network round trips, and is_reaping=True is enough to block submit + # from enqueueing while the snapshot is in flight. + try: + await self.persist_session_snapshot( + agent_session, + runtime_state="idle", + status="active", + raise_on_error=True, + ) + except Exception as e: + async with self._lock: + if self.sessions.get(session_id) is agent_session: + agent_session.is_reaping = False + logger.warning( + "Skipping reap of %s: could not persist resumable snapshot: %s", + session_id, + e, + ) + return False + + async with self._lock: + current = self.sessions.get(session_id) + if current is not agent_session: + return False + if ( + not agent_session.is_active + or agent_session.is_processing + or agent_session.session.pending_approval + or agent_session.last_active_at > cutoff + or not agent_session.submission_queue.empty() + ): + agent_session.is_reaping = False + return False + self.sessions.pop(session_id, None) + task = agent_session.task + session = agent_session.session + + if task is not None and not task.done(): + task.cancel() + # Use asyncio.wait, not wait_for: wait_for re-raises the cancelled + # task's CancelledError, which we'd have to swallow — and that same + # bare except would also eat an *outer* cancel aimed at the reaper + # itself (close() cancelling _reaper_task), hanging shutdown. + # asyncio.wait returns the cancelled task in `done` and lets an + # outer cancel propagate cleanly. + done, _pending = await asyncio.wait({task}, timeout=REAP_TEARDOWN_TIMEOUT_S) + if not done: + logger.warning( + "Reaper teardown timed out after %.0fs for %s; orphan " + "sweeper will handle any sandbox straggler", + REAP_TEARDOWN_TIMEOUT_S, + session_id, + ) + elif not task.cancelled(): + # Surface (and retrieve, to avoid "exception never retrieved") + # any non-cancellation teardown error. + exc = task.exception() + if exc is not None: + logger.warning("Reaper teardown error for %s: %s", session_id, exc) + else: + # No live task to run the cleanup finally — free the sandbox here so + # a reaped session never leaves an orphaned Space behind. + await self._cleanup_sandbox(session) + return True async def _run_session( self, @@ -330,10 +1690,20 @@ async def _run_session( submission_queue.get(), timeout=1.0 ) agent_session.is_processing = True + self._touch(agent_session) try: - should_continue = await process_submission(session, submission) + should_continue = await process_submission( + session, submission + ) finally: agent_session.is_processing = False + # Stamp on turn finish too: a turn that ran longer + # than the idle window would otherwise be reaped the + # instant it completes. + self._touch(agent_session) + if session.config.save_sessions: + await self.refresh_session_usage_metrics(agent_session) + await self.persist_session_snapshot(agent_session) if not should_continue: break except asyncio.TimeoutError: @@ -356,23 +1726,53 @@ async def _run_session( await self._cleanup_sandbox(session) + # Final-flush: always save on session death so we capture ended + # sessions even if the client disconnects without /shutdown. + # Idempotent via session_id key; detached subprocess. + if session.config.save_sessions: + try: + await self.refresh_session_usage_metrics( + agent_session, + error_code="final_billing_snapshot_error", + ) + session.save_and_upload_detached( + session.config.session_dataset_repo + ) + except Exception as e: + logger.warning(f"Final-flush failed for {session_id}: {e}") + async with self._lock: - if session_id in self.sessions: - self.sessions[session_id].is_active = False + if self.sessions.get(session_id) is agent_session: + agent_session.is_active = False + await self.persist_session_snapshot( + agent_session, + runtime_state="ended", + status="ended", + ) logger.info(f"Session {session_id} ended") async def submit(self, session_id: str, operation: Operation) -> bool: - """Submit an operation to a session.""" - async with self._lock: - agent_session = self.sessions.get(session_id) - - if not agent_session or not agent_session.is_active: - logger.warning(f"Session {session_id} not found or inactive") - return False + """Submit an operation to a session. + Enqueues under the lock and rejects sessions being reaped, so submit + and reap can't interleave: either the message is enqueued before the + reaper's empty() re-check (which then aborts the reap), or the session + is already popped (we return False and the caller reloads a fresh + runtime from Mongo). The queue is unbounded, so put_nowait never blocks. + """ submission = Submission(id=f"sub_{uuid.uuid4().hex[:8]}", operation=operation) - await agent_session.submission_queue.put(submission) + async with self._lock: + agent_session = self.sessions.get(session_id) + if ( + not agent_session + or not agent_session.is_active + or agent_session.is_reaping + ): + logger.warning(f"Session {session_id} not found or inactive") + return False + agent_session.submission_queue.put_nowait(submission) + self._touch(agent_session) return True async def submit_user_input(self, session_id: str, text: str) -> bool: @@ -408,7 +1808,13 @@ async def truncate(self, session_id: str, user_message_index: int) -> bool: agent_session = self.sessions.get(session_id) if not agent_session or not agent_session.is_active: return False - return agent_session.session.context_manager.truncate_to_user_message(user_message_index) + success = agent_session.session.context_manager.truncate_to_user_message( + user_message_index + ) + if success: + self._touch(agent_session) + await self.persist_session_snapshot(agent_session, runtime_state="idle") + return success async def compact(self, session_id: str) -> bool: """Compact context in a session.""" @@ -433,12 +1839,15 @@ async def shutdown_session(self, session_id: str) -> bool: return success async def delete_session(self, session_id: str) -> bool: - """Delete a session entirely.""" + """Soft-delete a session and stop its runtime resources.""" async with self._lock: agent_session = self.sessions.pop(session_id, None) if not agent_session: - return False + await self._store().soft_delete_session(session_id) + return True + + await self._store().soft_delete_session(session_id) # Clean up sandbox Space before cancelling the task await self._cleanup_sandbox(agent_session.session) @@ -453,26 +1862,100 @@ async def delete_session(self, session_id: str) -> bool: return True - def get_session_owner(self, session_id: str) -> str | None: - """Get the user_id that owns a session, or None if session doesn't exist.""" + async def teardown_sandbox(self, session_id: str) -> bool: + """Delete only this session's sandbox runtime, preserving chat state.""" + async with self._lock: + agent_session = self.sessions.get(session_id) + + if not agent_session or not agent_session.is_active: + return False + + await self._cleanup_sandbox(agent_session.session) + await self.persist_session_snapshot(agent_session, runtime_state="idle") + return True + + async def update_session_title(self, session_id: str, title: str | None) -> None: + """Persist a user-visible title for sidebar rehydration.""" agent_session = self.sessions.get(session_id) - if not agent_session: + if agent_session: + agent_session.title = title + await self._store().update_session_fields(session_id, title=title) + + async def update_session_model(self, session_id: str, model_id: str) -> bool: + agent_session = self.sessions.get(session_id) + if not agent_session or not agent_session.is_active: + return False + agent_session.session.update_model(model_id) + self._touch(agent_session) + await self.persist_session_snapshot(agent_session, runtime_state="idle") + return True + + async def update_session_auto_approval( + self, + session_id: str, + *, + enabled: bool, + cost_cap_usd: float | None, + cap_provided: bool = False, + ) -> dict[str, Any]: + agent_session = self.sessions.get(session_id) + if not agent_session or not agent_session.is_active: + raise ValueError("Session not found or inactive") + + session = agent_session.session + seed_spend: float | None = None + if enabled: + try: + seed_spend, _ = await self._current_session_usage_spend( + agent_session, + use_cache=False, + ) + except Exception as e: + logger.debug("Could not seed YOLO spend for %s: %s", session_id, e) + if enabled: + if not cap_provided and cost_cap_usd is None: + cost_cap_usd = getattr(session, "auto_approval_cost_cap_usd", None) + if cost_cap_usd is None: + cost_cap_usd = DEFAULT_YOLO_COST_CAP_USD + elif cost_cap_usd is None: + cost_cap_usd = DEFAULT_YOLO_COST_CAP_USD + else: + if not cap_provided: + cost_cap_usd = getattr(session, "auto_approval_cost_cap_usd", None) + + if hasattr(session, "set_auto_approval_policy"): + session.set_auto_approval_policy( + enabled=enabled, + cost_cap_usd=cost_cap_usd, + ) + else: + session.auto_approval_enabled = bool(enabled) + session.auto_approval_cost_cap_usd = cost_cap_usd + if enabled and seed_spend is not None: + seed_session_spend(session, seed_spend) + self._touch(agent_session) + await self.persist_session_snapshot(agent_session) + return self._auto_approval_summary(session) + + async def reconcile_session_auto_approval_from_usage( + self, + session_id: str, + usage_response: dict[str, Any], + ) -> dict[str, Any] | None: + agent_session = self.sessions.get(session_id) + if not agent_session or not agent_session.is_active: return None - return agent_session.user_id - def verify_session_access(self, session_id: str, user_id: str) -> bool: - """Check if a user has access to a session. + session = agent_session.session + if not bool(getattr(session, "auto_approval_enabled", False)): + return self._auto_approval_summary(session) - Returns True if: - - The session exists AND the user owns it - - The user_id is "dev" (dev mode bypass) - """ - owner = self.get_session_owner(session_id) - if owner is None: - return False - if user_id == "dev" or owner == "dev": - return True - return owner == user_id + current_spend, _ = self._usage_spend_from_response(usage_response) + previous_spend = session_spend_usd(session) + seed_session_spend(session, current_spend) + if session_spend_usd(session) != previous_spend: + self._touch(agent_session) + return self._auto_approval_summary(session) def get_session_info(self, session_id: str) -> dict[str, Any] | None: """Get information about a session.""" @@ -480,42 +1963,175 @@ def get_session_info(self, session_id: str) -> dict[str, Any] | None: if not agent_session: return None - # Extract pending approval tools if any - pending_approval = None - pa = agent_session.session.pending_approval - if pa and pa.get("tool_calls"): - pending_approval = [] - for tc in pa["tool_calls"]: - import json - try: - args = json.loads(tc.function.arguments) - except (json.JSONDecodeError, AttributeError): - args = {} - pending_approval.append({ - "tool": tc.function.name, - "tool_call_id": tc.id, - "arguments": args, - }) + pending_approval = self._pending_tools_for_api(agent_session.session) return { "session_id": session_id, "created_at": agent_session.created_at.isoformat(), + "usage_window_started_at": ( + agent_session.usage_window_started_at or agent_session.created_at + ).isoformat(), "is_active": agent_session.is_active, "is_processing": agent_session.is_processing, "message_count": len(agent_session.session.context_manager.items), "user_id": agent_session.user_id, "pending_approval": pending_approval, "model": agent_session.session.config.model_name, + "title": agent_session.title, + "notification_destinations": list( + agent_session.session.notification_destinations + ), + "auto_approval": self._auto_approval_summary(agent_session.session), } - def list_sessions(self, user_id: str | None = None) -> list[dict[str, Any]]: + async def reset_session_usage_window( + self, + session_id: str, + *, + started_at: datetime | None = None, + ) -> dict[str, Any] | None: + """Reset the account-billing window used for the visible usage meter.""" + agent_session = self.sessions.get(session_id) + if not agent_session: + return None + + window_start = started_at or datetime.utcnow() + agent_session.usage_window_started_at = window_start + billing_session_id = new_inference_billing_session_id(session_id, window_start) + self._set_inference_billing_session_id(agent_session, billing_session_id) + agent_session.usage_warning_spend_cache = {} + self._touch(agent_session) + + store = self._store() + if getattr(store, "enabled", False): + await store.update_session_fields( + session_id, + usage_window_started_at=window_start, + inference_billing_session_id=billing_session_id, + last_active_at=agent_session.last_active_at, + ) + return self.get_session_info(session_id) + + async def activate_session(self, session_id: str) -> dict[str, Any] | None: + """Mark a session as revisited without resetting its usage window.""" + agent_session = self.sessions.get(session_id) + if not agent_session: + return None + + self._touch(agent_session) + + store = self._store() + if getattr(store, "enabled", False): + await store.update_session_fields( + session_id, + last_active_at=agent_session.last_active_at, + ) + return self.get_session_info(session_id) + + def set_notification_destinations( + self, session_id: str, destinations: list[str] + ) -> list[str]: + """Replace the session's opted-in auto-notification destinations.""" + agent_session = self.sessions.get(session_id) + if not agent_session or not agent_session.is_active: + raise ValueError("Session not found or inactive") + + normalized: list[str] = [] + seen: set[str] = set() + for raw_name in destinations: + name = raw_name.strip() + if not name: + raise ValueError("Destination names must not be empty") + destination = self.config.messaging.get_destination(name) + if destination is None: + raise ValueError(f"Unknown destination '{name}'") + if not destination.allow_auto_events: + raise ValueError(f"Destination '{name}' is not enabled for auto events") + if name not in seen: + normalized.append(name) + seen.add(name) + + agent_session.session.set_notification_destinations(normalized) + self._touch(agent_session) + return normalized + + async def list_sessions(self, user_id: str | None = None) -> list[dict[str, Any]]: """List sessions, optionally filtered by user. Args: user_id: If provided, only return sessions owned by this user. If "dev", return all sessions (dev mode). """ - results = [] + results: list[dict[str, Any]] = [] + store = self._store() + if getattr(store, "enabled", False): + for row in await store.list_sessions(user_id or "dev"): + sid = row.get("session_id") or row.get("_id") + if not sid: + continue + runtime_info = self.get_session_info(str(sid)) + if runtime_info: + results.append(runtime_info) + continue + created_at = row.get("created_at") + if isinstance(created_at, datetime): + created_at_str = created_at.isoformat() + else: + created_at_str = str(created_at or datetime.utcnow().isoformat()) + usage_window_started_at = ( + row.get("usage_window_started_at") or created_at + ) + if isinstance(usage_window_started_at, datetime): + usage_window_started_at_str = usage_window_started_at.isoformat() + else: + usage_window_started_at_str = str( + usage_window_started_at or created_at_str + ) + pending = self._pending_docs_for_api(row.get("pending_approval") or []) + results.append( + { + "session_id": str(sid), + "created_at": created_at_str, + "usage_window_started_at": usage_window_started_at_str, + "is_active": row.get("status") != "ended", + "is_processing": row.get("runtime_state") == "processing", + "message_count": int(row.get("message_count") or 0), + "user_id": row.get("user_id") or "dev", + "pending_approval": pending or None, + "model": row.get("model"), + "title": row.get("title"), + "notification_destinations": row.get( + "notification_destinations" + ) + or [], + "auto_approval": { + "enabled": bool(row.get("auto_approval_enabled", False)), + "cost_cap_usd": row.get("auto_approval_cost_cap_usd"), + "estimated_spend_usd": float( + row.get("auto_approval_estimated_spend_usd") or 0.0 + ), + "remaining_usd": ( + None + if row.get("auto_approval_cost_cap_usd") is None + else round( + max( + 0.0, + float( + row.get("auto_approval_cost_cap_usd") or 0.0 + ) + - float( + row.get("auto_approval_estimated_spend_usd") + or 0.0 + ), + ), + 4, + ) + ), + }, + } + ) + return results + for sid in self.sessions: info = self.get_session_info(sid) if not info: diff --git a/backend/usage.py b/backend/usage.py new file mode 100644 index 000000000..22758f4df --- /dev/null +++ b/backend/usage.py @@ -0,0 +1,751 @@ +"""Usage aggregation for app-attributed ML Intern spend.""" + +import asyncio +import logging +from datetime import UTC, datetime, timedelta +from typing import Any +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError + +import httpx + +from agent.core.usage_metrics import summarize_sandbox_lifecycle + +USAGE_EVENT_TYPES = ( + "llm_call", + "hf_job_complete", + "sandbox_create", + "sandbox_destroy", +) + +logger = logging.getLogger(__name__) + +HF_BILLING_USAGE_V2_URL = "https://huggingface.co/api/settings/billing/usage-v2" +HF_BILLING_USAGE_BY_INFERENCE_SESSION_URL = ( + "https://huggingface.co/api/settings/billing/usage-by-inference-session" +) +HF_BILLING_URL = "https://huggingface.co/settings/billing" +HF_INFERENCE_PROVIDERS_PRICING_URL = ( + "https://huggingface.co/docs/inference-providers/en/pricing" +) +HF_JOBS_PRICING_URL = "https://huggingface.co/docs/hub/jobs-pricing" + + +def _utc(dt: datetime) -> datetime: + if dt.tzinfo is None: + return dt.replace(tzinfo=UTC) + return dt.astimezone(UTC) + + +def _iso(dt: datetime | None) -> str | None: + if dt is None: + return None + return _utc(dt).isoformat().replace("+00:00", "Z") + + +def _coerce_float(value: Any) -> float: + if isinstance(value, bool) or value is None: + return 0.0 + try: + return float(value) + except (TypeError, ValueError): + return 0.0 + + +def _coerce_int(value: Any) -> int: + if isinstance(value, bool) or value is None: + return 0 + try: + return int(value) + except (TypeError, ValueError): + return 0 + + +def _nano_usd_to_usd(value: Any) -> float: + return _coerce_float(value) / 1_000_000_000 + + +def _micro_usd_to_usd(value: Any) -> float: + return _coerce_float(value) / 1_000_000 + + +def _cents_to_usd(value: Any) -> float: + return _coerce_float(value) / 100 + + +def _coerce_timezone(timezone_name: str | None) -> ZoneInfo | None: + if not timezone_name: + return None + try: + return ZoneInfo(timezone_name) + except (ZoneInfoNotFoundError, ValueError): + return None + + +def _normalize_event_timestamp( + dt: datetime, + *, + timezone_name: str | None = None, +) -> datetime: + if dt.tzinfo is not None: + return _utc(dt) + timezone = _coerce_timezone(timezone_name) + if timezone is not None: + return dt.replace(tzinfo=timezone).astimezone(UTC) + return dt.astimezone(UTC) + + +def _parse_timestamp( + value: Any, *, timezone_name: str | None = None +) -> datetime | None: + if isinstance(value, datetime): + return _normalize_event_timestamp(value, timezone_name=timezone_name) + if not isinstance(value, str) or not value: + return None + try: + return _normalize_event_timestamp( + datetime.fromisoformat(value.replace("Z", "+00:00")), + timezone_name=timezone_name, + ) + except ValueError: + return None + + +def event_created_at( + event: dict[str, Any], + *, + timezone_name: str | None = None, +) -> datetime | None: + return _parse_timestamp( + event.get("created_at") or event.get("timestamp"), + timezone_name=timezone_name, + ) + + +def resolve_usage_windows( + timezone_name: str | None, + *, + now: datetime | None = None, +) -> dict[str, datetime | str]: + """Return UTC month window for a browser timezone.""" + try: + tz = ZoneInfo(timezone_name or "UTC") + except (ZoneInfoNotFoundError, ValueError): + tz = ZoneInfo("UTC") + + now_utc = _utc(now or datetime.now(UTC)) + local_now = now_utc.astimezone(tz) + month_local = local_now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + return { + "timezone": tz.key, + "now_utc": now_utc, + "month_start_utc": month_local.astimezone(UTC), + } + + +def _empty_bucket( + *, + session_id: str | None = None, +) -> dict[str, Any]: + return { + "session_id": session_id, + "total_usd": 0.0, + "inference_usd": 0.0, + "hf_jobs_estimated_usd": 0.0, + "sandbox_estimated_usd": 0.0, + "llm_calls": 0, + "hf_jobs_count": 0, + "sandbox_count": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "cache_read_tokens": 0, + "cache_creation_tokens": 0, + "total_tokens": 0, + "hf_jobs_billable_seconds_estimate": 0, + "sandbox_billable_seconds_estimate": 0, + } + + +def _empty_hf_account_bucket( + *, + window_start: datetime | None = None, + window_end: datetime | None = None, + timezone: str | None = None, +) -> dict[str, Any]: + return { + "window_start": _iso(window_start), + "window_end": _iso(window_end), + "timezone": timezone, + "total_usd": 0.0, + "inference_providers_usd": 0.0, + "hf_jobs_usd": 0.0, + "inference_provider_requests": 0, + "hf_jobs_minutes": 0.0, + } + + +def aggregate_usage_events( + events: list[dict[str, Any]], + *, + session_id: str | None = None, +) -> dict[str, Any]: + bucket = _empty_bucket(session_id=session_id) + for event in events: + event_type = event.get("event_type") + data = event.get("data") or {} + if event_type == "llm_call": + bucket["llm_calls"] += 1 + bucket["inference_usd"] += _coerce_float(data.get("cost_usd")) + prompt_tokens = _coerce_int(data.get("prompt_tokens")) + completion_tokens = _coerce_int(data.get("completion_tokens")) + cache_read_tokens = _coerce_int(data.get("cache_read_tokens")) + cache_creation_tokens = _coerce_int(data.get("cache_creation_tokens")) + total_tokens = _coerce_int(data.get("total_tokens")) or ( + prompt_tokens + + completion_tokens + + cache_read_tokens + + cache_creation_tokens + ) + bucket["prompt_tokens"] += prompt_tokens + bucket["completion_tokens"] += completion_tokens + bucket["cache_read_tokens"] += cache_read_tokens + bucket["cache_creation_tokens"] += cache_creation_tokens + bucket["total_tokens"] += total_tokens + elif event_type == "hf_job_complete": + bucket["hf_jobs_count"] += 1 + bucket["hf_jobs_estimated_usd"] += _coerce_float( + data.get("estimated_cost_usd") + ) + bucket["hf_jobs_billable_seconds_estimate"] += _coerce_int( + data.get("billable_seconds_estimate") or data.get("wall_time_s") + ) + elif event_type == "sandbox_destroy": + # Sandbox costs are paired and added after the main pass so the + # create event can provide hardware pricing metadata. + continue + + _aggregate_sandbox_usage(events, bucket) + + bucket["inference_usd"] = round(bucket["inference_usd"], 6) + bucket["hf_jobs_estimated_usd"] = round(bucket["hf_jobs_estimated_usd"], 6) + bucket["sandbox_estimated_usd"] = round(bucket["sandbox_estimated_usd"], 6) + bucket["total_usd"] = round( + ( + bucket["inference_usd"] + + bucket["hf_jobs_estimated_usd"] + + bucket["sandbox_estimated_usd"] + ), + 6, + ) + return bucket + + +def _aggregate_sandbox_usage( + events: list[dict[str, Any]], + bucket: dict[str, Any], +) -> None: + lifecycle_events = [ + (index, event) + for index, event in enumerate(events) + if event.get("event_type") in {"sandbox_create", "sandbox_destroy"} + ] + sandbox = summarize_sandbox_lifecycle(lifecycle_events) + bucket["sandbox_count"] += sandbox["matched_pairs"] + bucket["sandbox_billable_seconds_estimate"] += sandbox["billable_seconds_estimate"] + bucket["sandbox_estimated_usd"] += sandbox["estimated_usd"] + + +def _account_bucket_from_billing_usage( + payload: dict[str, Any] | None, + *, + window_start: datetime, + window_end: datetime, + timezone: str, +) -> dict[str, Any]: + bucket = _empty_hf_account_bucket( + window_start=window_start, + window_end=window_end, + timezone=timezone, + ) + usage = payload.get("usage") if isinstance(payload, dict) else {} + if not isinstance(usage, dict): + return bucket + + inference = usage.get("inferenceProviders") + if not isinstance(inference, dict): + inference = {} + jobs = usage.get("jobs") + if not isinstance(jobs, dict): + jobs = {} + + bucket["inference_providers_usd"] = round( + _nano_usd_to_usd(inference.get("usedNanoUsd")), + 6, + ) + bucket["hf_jobs_usd"] = round(_micro_usd_to_usd(jobs.get("usedMicroUsd")), 6) + bucket["inference_provider_requests"] = _coerce_int(inference.get("numRequests")) + bucket["hf_jobs_minutes"] = round(_coerce_float(jobs.get("totalMinutes")), 3) + bucket["total_usd"] = round( + bucket["inference_providers_usd"] + bucket["hf_jobs_usd"], + 6, + ) + return bucket + + +def _session_bucket_from_inference_session_usage( + payload: dict[str, Any] | None, + *, + session_id: str, + window_start: datetime, + window_end: datetime, + timezone: str, +) -> dict[str, Any]: + bucket = _empty_hf_account_bucket( + window_start=window_start, + window_end=window_end, + timezone=timezone, + ) + periods = payload.get("periods") if isinstance(payload, dict) else [] + if not isinstance(periods, list): + return bucket + + cost_cents = 0.0 + request_count = 0 + for period in periods: + if not isinstance(period, dict): + continue + sessions = period.get("sessions") + if not isinstance(sessions, list): + continue + for session in sessions: + if not isinstance(session, dict) or session.get("id") != session_id: + continue + cost_cents += _coerce_float(session.get("costCents")) + request_count += _coerce_int(session.get("requestCount")) + + bucket["inference_providers_usd"] = round(_cents_to_usd(cost_cents), 6) + bucket["inference_provider_requests"] = request_count + bucket["total_usd"] = bucket["inference_providers_usd"] + return bucket + + +def _inference_credits_from_billing_usage( + payload: dict[str, Any] | None, +) -> dict[str, Any] | None: + usage = payload.get("usage") if isinstance(payload, dict) else {} + if not isinstance(usage, dict): + return None + inference = usage.get("inferenceProviders") + if not isinstance(inference, dict): + return None + + included_usd = _nano_usd_to_usd(inference.get("includedNanoUsd")) + used_usd = _nano_usd_to_usd(inference.get("usedNanoUsd")) + limit_usd = _nano_usd_to_usd(inference.get("limitNanoUsd")) + return { + "included_usd": round(included_usd, 6), + "used_usd": round(used_usd, 6), + "remaining_included_usd": round(max(0.0, included_usd - used_usd), 6), + "limit_usd": round(limit_usd, 6), + "remaining_limit_usd": round(max(0.0, limit_usd - used_usd), 6), + "num_requests": _coerce_int(inference.get("numRequests")), + "period_start": inference.get("periodStart"), + "period_end": inference.get("periodEnd"), + } + + +async def _fetch_hf_billing_usage_v2( + hf_token: str, + *, + start: datetime, + end: datetime, +) -> dict[str, Any] | None: + start_ts = max(1, int(_utc(start).timestamp())) + end_ts = max(start_ts + 1, int(_utc(end).timestamp())) + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.get( + HF_BILLING_USAGE_V2_URL, + params={"startDate": start_ts, "endDate": end_ts}, + headers={"Authorization": f"Bearer {hf_token}"}, + ) + if response.status_code != 200: + logger.debug( + "HF billing usage-v2 failed: status=%s body=%s", + response.status_code, + response.text[:200], + ) + return None + payload = response.json() + return payload if isinstance(payload, dict) else None + except (httpx.HTTPError, ValueError) as e: + logger.debug("HF billing usage-v2 failed: %s", e) + return None + + +async def _fetch_hf_inference_session_usage( + hf_token: str, + *, + start: datetime, + end: datetime, +) -> dict[str, Any] | None: + start_ts = _iso(start) + end_ts = _iso(max(_utc(end), _utc(start) + timedelta(seconds=1))) + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.get( + HF_BILLING_USAGE_BY_INFERENCE_SESSION_URL, + params={"startDate": start_ts, "endDate": end_ts}, + headers={"Authorization": f"Bearer {hf_token}"}, + ) + if response.status_code != 200: + logger.debug( + "HF inference session usage failed: status=%s body=%s", + response.status_code, + response.text[:200], + ) + return None + payload = response.json() + return payload if isinstance(payload, dict) else None + except (httpx.HTTPError, ValueError) as e: + logger.debug("HF inference session usage failed: %s", e) + return None + + +def _session_usage_window_started_at( + manager: Any, session_id: str | None +) -> datetime | None: + if not session_id: + return None + agent_session = getattr(manager, "sessions", {}).get(session_id) + usage_window_started_at = getattr(agent_session, "usage_window_started_at", None) + if isinstance(usage_window_started_at, datetime): + return _utc(usage_window_started_at) + created_at = getattr(agent_session, "created_at", None) + if isinstance(created_at, datetime): + return _utc(created_at) + return None + + +def _session_inference_billing_session_id( + manager: Any, session_id: str | None +) -> str | None: + if not session_id: + return None + agent_session = getattr(manager, "sessions", {}).get(session_id) + billing_session_id = getattr(agent_session, "inference_billing_session_id", None) + if isinstance(billing_session_id, str) and billing_session_id: + return billing_session_id + runtime_session = getattr(agent_session, "session", None) + billing_session_id = getattr(runtime_session, "inference_billing_session_id", None) + if isinstance(billing_session_id, str) and billing_session_id: + return billing_session_id + return None + + +async def _load_persisted_session_usage_window_metadata( + manager: Any, + session_id: str | None, +) -> tuple[datetime | None, str | None]: + if not session_id: + return None, None + store = manager._store() + if not getattr(store, "enabled", False) or not hasattr(store, "load_session"): + return None, None + loaded = await store.load_session(session_id) + metadata = loaded.get("metadata") if isinstance(loaded, dict) else None + started_at = None + billing_session_id = None + if isinstance(metadata, dict): + started_at = metadata.get("usage_window_started_at") or metadata.get( + "created_at" + ) + raw_billing_session_id = metadata.get("inference_billing_session_id") + if isinstance(raw_billing_session_id, str) and raw_billing_session_id: + billing_session_id = raw_billing_session_id + if isinstance(started_at, datetime): + return _utc(started_at), billing_session_id + parsed = _parse_timestamp(started_at) + return (_utc(parsed) if parsed is not None else None), billing_session_id + + +async def _build_hf_account_usage( + manager: Any, + *, + hf_token: str | None, + session_id: str | None, + timezone: str, + now_utc: datetime, + month_start: datetime, +) -> dict[str, Any]: + account_usage: dict[str, Any] = { + "source": "hf_billing", + "available": False, + "current_session": None, + "month": None, + "inference_providers_credits": None, + } + if not hf_token: + account_usage["error"] = "missing_hf_token" + return account_usage + + session_start = _session_usage_window_started_at(manager, session_id) + billing_session_id = _session_inference_billing_session_id(manager, session_id) + if session_start is None or billing_session_id is None: + ( + persisted_start, + persisted_billing_session_id, + ) = await _load_persisted_session_usage_window_metadata(manager, session_id) + if session_start is None: + session_start = persisted_start + if billing_session_id is None: + billing_session_id = persisted_billing_session_id + + window_tasks: dict[str, tuple[datetime, asyncio.Task[dict[str, Any] | None]]] = { + "month": ( + month_start, + asyncio.create_task( + _fetch_hf_billing_usage_v2(hf_token, start=month_start, end=now_utc) + ), + ), + } + if billing_session_id is not None and session_start is not None: + window_tasks["current_session"] = ( + session_start, + asyncio.create_task( + _fetch_hf_inference_session_usage( + hf_token, + start=session_start, + end=now_utc, + ) + ), + ) + + payloads: dict[str, dict[str, Any] | None] = {} + for name, (_, task) in window_tasks.items(): + payloads[name] = await task + + any_payload = any(isinstance(payload, dict) for payload in payloads.values()) + account_usage["available"] = any_payload + if not any_payload: + account_usage["error"] = "billing_usage_unavailable" + return account_usage + + for name, (start, _) in window_tasks.items(): + payload = payloads.get(name) + if payload is None: + continue + if name == "current_session" and billing_session_id is not None: + account_usage[name] = _session_bucket_from_inference_session_usage( + payload, + session_id=billing_session_id, + window_start=start, + window_end=now_utc, + timezone=timezone, + ) + else: + account_usage[name] = _account_bucket_from_billing_usage( + payload, + window_start=start, + window_end=now_utc, + timezone=timezone, + ) + + account_usage["inference_providers_credits"] = ( + _inference_credits_from_billing_usage(payloads.get("month")) + ) + return account_usage + + +async def build_hf_billing_snapshot( + manager: Any, + *, + hf_token: str | None, + session_id: str | None, + timezone_name: str | None = None, + now: datetime | None = None, +) -> dict[str, Any]: + """Return a dataset-safe HF billing rollup for the session window. + + This intentionally omits monthly account totals and credit-limit details. + The snapshot is an account-window delta, not per-call attribution. + """ + windows = resolve_usage_windows(timezone_name, now=now) + timezone = str(windows["timezone"]) + now_utc = windows["now_utc"] + snapshot: dict[str, Any] = { + "billing_scope": "account_window_delta", + "hf_billing": { + "source": "hf_billing_usage_v2", + "available": False, + "error": None, + "current_session": None, + }, + } + hf_billing = snapshot["hf_billing"] + + if not hf_token: + hf_billing["error"] = "missing_hf_token" + return snapshot + if not session_id: + hf_billing["error"] = "missing_session_id" + return snapshot + + session_start = _session_usage_window_started_at(manager, session_id) + if session_start is None: + session_start, _ = await _load_persisted_session_usage_window_metadata( + manager, + session_id, + ) + if session_start is None: + hf_billing["error"] = "missing_session_window" + return snapshot + + payload = await _fetch_hf_billing_usage_v2( + hf_token, + start=session_start, + end=now_utc, + ) + if not isinstance(payload, dict): + hf_billing["error"] = "billing_usage_unavailable" + return snapshot + + hf_billing["available"] = True + hf_billing["current_session"] = _account_bucket_from_billing_usage( + payload, + window_start=session_start, + window_end=now_utc, + timezone=timezone, + ) + return snapshot + + +def _event_in_window( + event: dict[str, Any], + *, + start: datetime | None = None, + end: datetime | None = None, + timezone_name: str | None = None, +) -> bool: + if start is None and end is None: + return True + created_at = event_created_at(event, timezone_name=timezone_name) + if created_at is None: + return False + if start is not None and created_at < _utc(start): + return False + if end is not None and created_at >= _utc(end): + return False + return True + + +def _events_from_runtime_session(agent_session: Any) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + for raw in getattr(agent_session.session, "logged_events", []) or []: + if raw.get("event_type") not in USAGE_EVENT_TYPES: + continue + events.append( + { + "session_id": agent_session.session_id, + "event_type": raw.get("event_type"), + "data": raw.get("data") or {}, + "timestamp": raw.get("timestamp"), + } + ) + return events + + +def _runtime_sessions_for_user(manager: Any, user_id: str) -> list[Any]: + sessions = list(getattr(manager, "sessions", {}).values()) + if user_id == "dev": + return sessions + return [session for session in sessions if session.user_id == user_id] + + +async def _load_usage_events( + manager: Any, + *, + user_id: str, + session_id: str | None = None, + start: datetime | None = None, + end: datetime | None = None, + timezone_name: str | None = None, +) -> list[dict[str, Any]]: + store = manager._store() + if getattr(store, "enabled", False): + return await store.load_usage_events( + user_id, + session_id=session_id, + start=start, + end=end, + ) + + events: list[dict[str, Any]] = [] + for agent_session in _runtime_sessions_for_user(manager, user_id): + if session_id is not None and agent_session.session_id != session_id: + continue + for event in _events_from_runtime_session(agent_session): + if _event_in_window( + event, + start=start, + end=end, + timezone_name=timezone_name, + ): + events.append(event) + return events + + +async def build_usage_response( + manager: Any, + *, + user_id: str, + hf_token: str | None = None, + session_id: str | None = None, + timezone_name: str | None = None, + now: datetime | None = None, +) -> dict[str, Any]: + windows = resolve_usage_windows(timezone_name, now=now) + timezone = str(windows["timezone"]) + now_utc = windows["now_utc"] + month_start = windows["month_start_utc"] + + session_events: list[dict[str, Any]] = [] + if session_id: + session_start = _session_usage_window_started_at(manager, session_id) + if session_start is None: + session_start, _ = await _load_persisted_session_usage_window_metadata( + manager, + session_id, + ) + session_events = await _load_usage_events( + manager, + user_id=user_id, + session_id=session_id, + start=session_start, + ) + + hf_account = await _build_hf_account_usage( + manager, + hf_token=hf_token, + session_id=session_id, + timezone=timezone, + now_utc=now_utc, + month_start=month_start, + ) + + return { + "source": "app_telemetry", + "currency": "USD", + "generated_at": _iso(now_utc), + "timezone": timezone, + "session": ( + aggregate_usage_events(session_events, session_id=session_id) + if session_id + else None + ), + "hf_account": hf_account, + "links": { + "hf_billing": HF_BILLING_URL, + "inference_providers_pricing": HF_INFERENCE_PROVIDERS_PRICING_URL, + "jobs_pricing": HF_JOBS_PRICING_URL, + }, + } diff --git a/backend/user_quotas.py b/backend/user_quotas.py deleted file mode 100644 index 2b38b1111..000000000 --- a/backend/user_quotas.py +++ /dev/null @@ -1,83 +0,0 @@ -"""In-memory daily quota for Claude session creations. - -Tracks per-user Claude session starts against a daily cap derived from the -user's HF plan. Caps reset at UTC midnight; the store itself is in-process -and wipes on restart (deliberate — the cost of occasional over-subsidy at -restart is much lower than running a DB). - -Unit: session *creations*, not messages. A user who selects Claude in a new -session consumes one quota point; switching an existing Claude session to -Claude again doesn't (`AgentSession.claude_counted` guards that). - -Cap tiers: - free user → CLAUDE_FREE_DAILY (1) - pro / org → CLAUDE_PRO_DAILY (20) -""" - -import asyncio -import os -from datetime import UTC, datetime - -CLAUDE_FREE_DAILY: int = int(os.environ.get("CLAUDE_FREE_DAILY", "1")) -CLAUDE_PRO_DAILY: int = int(os.environ.get("CLAUDE_PRO_DAILY", "20")) - -# user_id -> (day_utc_iso, count_for_that_day) -_claude_counts: dict[str, tuple[str, int]] = {} -_lock = asyncio.Lock() - - -def _today() -> str: - return datetime.now(UTC).date().isoformat() - - -def daily_cap_for(plan: str | None) -> int: - """Return the daily Claude-session cap for the given plan.""" - return CLAUDE_FREE_DAILY if (plan or "free") == "free" else CLAUDE_PRO_DAILY - - -async def get_claude_used_today(user_id: str) -> int: - """Return today's Claude session count for the user (0 if none / stale day).""" - async with _lock: - entry = _claude_counts.get(user_id) - if entry is None: - return 0 - day, count = entry - if day != _today(): - # Stale day — drop the entry so the first increment starts fresh. - _claude_counts.pop(user_id, None) - return 0 - return count - - -async def increment_claude(user_id: str) -> int: - """Bump today's Claude session count for the user. Returns the new value.""" - async with _lock: - today = _today() - day, count = _claude_counts.get(user_id, (today, 0)) - if day != today: - count = 0 - count += 1 - _claude_counts[user_id] = (today, count) - return count - - -async def refund_claude(user_id: str) -> None: - """Decrement today's count — used when session creation fails after a successful gate.""" - async with _lock: - entry = _claude_counts.get(user_id) - if entry is None: - return - day, count = entry - if day != _today(): - _claude_counts.pop(user_id, None) - return - new_count = max(0, count - 1) - if new_count == 0: - _claude_counts.pop(user_id, None) - else: - _claude_counts[user_id] = (day, new_count) - - -def _reset_for_tests() -> None: - """Test-only: clear the in-memory store.""" - _claude_counts.clear() diff --git a/configs/__init__.py b/configs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/configs/cli_agent_config.json b/configs/cli_agent_config.json new file mode 100644 index 000000000..97462cf6d --- /dev/null +++ b/configs/cli_agent_config.json @@ -0,0 +1,22 @@ +{ + "model_name": "zai-org/GLM-5.2:novita", + "save_sessions": true, + "session_dataset_repo": "smolagents/ml-intern-sessions", + "share_traces": true, + "personal_trace_repo_template": "{hf_user}/ml-intern-sessions", + "yolo_mode": false, + "confirm_cpu_jobs": true, + "auto_file_upload": true, + "tool_runtime": "local", + "messaging": { + "enabled": false, + "auto_event_types": ["approval_required", "error", "turn_complete"], + "destinations": {} + }, + "mcpServers": { + "hf-mcp-server": { + "transport": "http", + "url": "https://huggingface.co/mcp?login" + } + } +} diff --git a/configs/main_agent_config.json b/configs/frontend_agent_config.json similarity index 51% rename from configs/main_agent_config.json rename to configs/frontend_agent_config.json index af76608f3..5c7ada21a 100644 --- a/configs/main_agent_config.json +++ b/configs/frontend_agent_config.json @@ -1,7 +1,9 @@ { - "model_name": "bedrock/us.anthropic.claude-opus-4-6-v1", + "model_name": "${ML_INTERN_DEFAULT_MODEL_ID:-zai-org/GLM-5.2:novita}", "save_sessions": true, - "session_dataset_repo": "akseljoonas/hf-agent-sessions", + "session_dataset_repo": "smolagents/ml-intern-sessions", + "share_traces": true, + "personal_trace_repo_template": "{hf_user}/ml-intern-sessions", "yolo_mode": false, "confirm_cpu_jobs": true, "auto_file_upload": true, diff --git a/frontend/src/components/Chat/ActivityStatusBar.tsx b/frontend/src/components/Chat/ActivityStatusBar.tsx index 3dd0af534..61e502d2d 100644 --- a/frontend/src/components/Chat/ActivityStatusBar.tsx +++ b/frontend/src/components/Chat/ActivityStatusBar.tsx @@ -25,7 +25,7 @@ function formatResearchStatus(raw: string): string { const s = raw.replace(/^▸\s*/, ''); const jsonStart = s.indexOf('{'); const toolName = jsonStart > 0 ? s.slice(0, jsonStart).trim() : s.trim(); - let args: Record = {}; + const args: Record = {}; if (jsonStart > 0) { const jsonStr = s.slice(jsonStart); try { diff --git a/frontend/src/components/Chat/AssistantMessage.tsx b/frontend/src/components/Chat/AssistantMessage.tsx index 83bd8cae5..91c7b8c10 100644 --- a/frontend/src/components/Chat/AssistantMessage.tsx +++ b/frontend/src/components/Chat/AssistantMessage.tsx @@ -1,13 +1,19 @@ -import { useMemo } from 'react'; -import { Box, Stack, Typography } from '@mui/material'; +import { useMemo, useState } from 'react'; +import { Box, IconButton, Stack, Tooltip, Typography } from '@mui/material'; +import ThumbUpOutlined from '@mui/icons-material/ThumbUpOutlined'; +import ThumbUp from '@mui/icons-material/ThumbUp'; +import ThumbDownOutlined from '@mui/icons-material/ThumbDownOutlined'; +import ThumbDown from '@mui/icons-material/ThumbDown'; import MarkdownContent from './MarkdownContent'; import ToolCallGroup from './ToolCallGroup'; +import { apiFetch } from '@/utils/api'; import type { UIMessage } from 'ai'; import type { MessageMeta } from '@/types/agent'; interface AssistantMessageProps { message: UIMessage; isStreaming?: boolean; + sessionId?: string | null; approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise; } @@ -43,8 +49,27 @@ function groupParts(parts: UIMessage['parts']) { return groups; } -export default function AssistantMessage({ message, isStreaming = false, approveTools }: AssistantMessageProps) { +export default function AssistantMessage({ message, isStreaming = false, sessionId, approveTools }: AssistantMessageProps) { const groups = useMemo(() => groupParts(message.parts), [message.parts]); + const [feedback, setFeedback] = useState<'up' | 'down' | null>(null); + const [feedbackBusy, setFeedbackBusy] = useState(false); + + const sendFeedback = async (rating: 'up' | 'down') => { + if (!sessionId || feedbackBusy) return; + setFeedbackBusy(true); + // Optimistic toggle — feedback is observability, not a hard requirement. + setFeedback(rating); + try { + await apiFetch(`/api/feedback/${sessionId}`, { + method: 'POST', + body: JSON.stringify({ rating, message_id: message.id }), + }); + } catch { + // Silently swallow — don't block chat UX on a telemetry write. + } finally { + setFeedbackBusy(false); + } + }; // Find the last text group index for streaming cursor let lastTextIdx = -1; @@ -114,6 +139,24 @@ export default function AssistantMessage({ message, isStreaming = false, approve return null; })} + {!isStreaming && sessionId && ( + + + sendFeedback('up')}> + {feedback === 'up' ? : } + + + + sendFeedback('down')}> + {feedback === 'down' ? : } + + + + )} ); } diff --git a/frontend/src/components/Chat/ChatErrorBanner.tsx b/frontend/src/components/Chat/ChatErrorBanner.tsx new file mode 100644 index 000000000..38ad87990 --- /dev/null +++ b/frontend/src/components/Chat/ChatErrorBanner.tsx @@ -0,0 +1,154 @@ +import { useEffect, useMemo, useState } from 'react'; +import { Alert, AlertTitle, Box, Button, Link, Typography } from '@mui/material'; +import { useAgentStore } from '@/store/agentStore'; +import { apiFetch } from '@/utils/api'; +import { inferenceCreditCta, isInferenceCreditError } from '@/utils/inferenceBilling'; + +interface ChatErrorBannerProps { + error: string; + sessionId: string; + model?: string | null; + onDismiss: () => void; +} + +const DISCUSSIONS_URL = 'https://huggingface.co/spaces/smolagents/ml-intern/discussions'; + +export default function ChatErrorBanner({ error, sessionId, model, onDismiss }: ChatErrorBannerProps) { + const [copied, setCopied] = useState(false); + const [reportedAt, setReportedAt] = useState(() => new Date().toISOString()); + const userPlan = useAgentStore((s) => s.user?.plan); + const isCreditError = isInferenceCreditError(error); + const creditCta = isCreditError ? inferenceCreditCta(userPlan) : null; + + useEffect(() => { + setReportedAt(new Date().toISOString()); + setCopied(false); + }, [error]); + + const details = useMemo( + () => [ + 'ML Intern message failure', + `time: ${reportedAt}`, + `session: ${sessionId}`, + `model: ${model || 'unknown'}`, + `error: ${error}`, + ].join('\n'), + [error, model, reportedAt, sessionId], + ); + + const copyDetails = async () => { + try { + await navigator.clipboard.writeText(details); + setCopied(true); + } catch { + setCopied(false); + } + }; + + const trackProClick = () => { + if (userPlan === 'pro') return; + void apiFetch(`/api/pro-click/${sessionId}`, { + method: 'POST', + body: JSON.stringify({ source: 'inference_credit_error', target: 'hf_pro' }), + }).catch(() => {}); + }; + + return ( + + + + + + } + > + + {creditCta?.title ?? 'Message failed'} + + + {creditCta ? ( + creditCta.message + ) : ( + <> + The backend could not process the last message. Retry after a moment. If it keeps + happening,{' '} + + open a discussion + {' '} + with the copied details. + + )} + + {creditCta && ( + + + {creditCta.secondaryHref && creditCta.secondaryLabel && ( + + )} + + )} + + {error} + + + + ); +} diff --git a/frontend/src/components/Chat/ChatInput.tsx b/frontend/src/components/Chat/ChatInput.tsx index d9fe5c4dc..1cc642637 100644 --- a/frontend/src/components/Chat/ChatInput.tsx +++ b/frontend/src/components/Chat/ChatInput.tsx @@ -1,19 +1,42 @@ import { useState, useCallback, useEffect, useRef, KeyboardEvent } from 'react'; -import { Box, TextField, IconButton, CircularProgress, Typography, Menu, MenuItem, ListItemIcon, ListItemText, Chip } from '@mui/material'; +import { + Alert, + Box, + TextField, + IconButton, + CircularProgress, + Typography, + Menu, + MenuItem, + ListItemIcon, + ListItemText, + Chip, + LinearProgress, + Snackbar, + Tooltip, +} from '@mui/material'; import ArrowUpwardIcon from '@mui/icons-material/ArrowUpward'; import ArrowDropDownIcon from '@mui/icons-material/ArrowDropDown'; import StopIcon from '@mui/icons-material/Stop'; -import { apiFetch } from '@/utils/api'; -import { useUserQuota } from '@/hooks/useUserQuota'; -import ClaudeCapDialog from '@/components/ClaudeCapDialog'; +import AddIcon from '@mui/icons-material/Add'; +import { apiFetch, apiUpload } from '@/utils/api'; +import JobsUpgradeDialog from '@/components/JobsUpgradeDialog'; import { useAgentStore } from '@/store/agentStore'; -import { FIRST_FREE_MODEL_PATH } from '@/utils/model'; +import { useSessionStore } from '@/store/sessionStore'; +import { + CLAUDE_OPUS_48_MODEL_PATH, + DEEPSEEK_V4_PRO_MODEL_PATH, + GLM_52_MODEL_PATH, + GPT_55_MODEL_PATH, + KIMI_K27_CODE_MODEL_PATH, + MINIMAX_M3_MODEL_PATH, + isClaudePath, +} from '@/utils/model'; // Model configuration interface ModelOption { id: string; name: string; - description: string; modelPath: string; avatarUrl: string; recommended?: boolean; @@ -24,69 +47,203 @@ const getHfAvatarUrl = (modelId: string) => { return `https://huggingface.co/api/avatars/${org}`; }; -const MODEL_OPTIONS: ModelOption[] = [ +const DEFAULT_MODEL_OPTIONS: ModelOption[] = [ { - id: 'kimi-k2.6', - name: 'Kimi K2.6', - description: 'Novita', - modelPath: 'moonshotai/Kimi-K2.6', - avatarUrl: getHfAvatarUrl('moonshotai/Kimi-K2.6'), - recommended: true, + id: 'claude-opus-4-8', + name: 'Claude Opus 4.8', + modelPath: CLAUDE_OPUS_48_MODEL_PATH, + avatarUrl: getHfAvatarUrl(CLAUDE_OPUS_48_MODEL_PATH), }, { - id: 'claude-opus', - name: 'Claude Opus 4.6', - description: 'Anthropic', - modelPath: 'anthropic/claude-opus-4-6', - avatarUrl: 'https://huggingface.co/api/avatars/Anthropic', - recommended: true, + id: 'gpt-5.5', + name: 'GPT-5.5', + modelPath: GPT_55_MODEL_PATH, + avatarUrl: getHfAvatarUrl(GPT_55_MODEL_PATH), + }, + { + id: 'kimi-k2.7-code', + name: 'Kimi K2.7 Code', + modelPath: KIMI_K27_CODE_MODEL_PATH, + avatarUrl: getHfAvatarUrl(KIMI_K27_CODE_MODEL_PATH), }, { - id: 'minimax-m2.7', - name: 'MiniMax M2.7', - description: 'Novita', - modelPath: 'MiniMaxAI/MiniMax-M2.7', - avatarUrl: getHfAvatarUrl('MiniMaxAI/MiniMax-M2.7'), + id: 'minimax-m3', + name: 'MiniMax M3', + modelPath: MINIMAX_M3_MODEL_PATH, + avatarUrl: getHfAvatarUrl(MINIMAX_M3_MODEL_PATH), }, { - id: 'glm-5.1', - name: 'GLM 5.1', - description: 'Together', - modelPath: 'zai-org/GLM-5.1', - avatarUrl: getHfAvatarUrl('zai-org/GLM-5.1'), + id: 'glm-5.2', + name: 'GLM 5.2', + modelPath: GLM_52_MODEL_PATH, + avatarUrl: getHfAvatarUrl(GLM_52_MODEL_PATH), + recommended: true, + }, + { + id: 'deepseek-v4-pro', + name: 'DeepSeek V4 Pro', + modelPath: DEEPSEEK_V4_PRO_MODEL_PATH, + avatarUrl: getHfAvatarUrl('deepseek-ai/DeepSeek-V4-Pro'), }, ]; -const findModelByPath = (path: string): ModelOption | undefined => { - return MODEL_OPTIONS.find(m => m.modelPath === path || path?.includes(m.id)); +const DEFAULT_MODEL_PATH = GLM_52_MODEL_PATH; + +const normalizeModelPath = (path: string | undefined) => ( + (path ?? '') + .toLowerCase() + .replace(/^huggingface\//, '') + .replace(/claude-opus-4\.(\d)/g, 'claude-opus-4-$1') +); + +const findModelByPath = (path: string, options: ModelOption[]): ModelOption | undefined => { + const normalizedPath = normalizeModelPath(path); + const matched = options.find((m) => { + const normalizedModelPath = normalizeModelPath(m.modelPath); + const normalizedId = normalizeModelPath(m.id); + return ( + m.modelPath === path || + normalizedModelPath === normalizedPath || + normalizedPath.includes(normalizedId) + ); + }); + if (matched) return matched; + if (isClaudePath(path)) { + const claude = options.find(isClaudeModel); + if (claude) return claude; + } + return undefined; +}; + +const modelOptionId = (modelPath: string) => ( + normalizeModelPath(modelPath) + .replace(/[^a-z0-9]+/g, '-') + .replace(/^-|-$/g, '') +); + +const modelOptionFromApi = (model: { + id?: string; + label?: string; + recommended?: boolean; +}): ModelOption | null => { + if (!model.id) return null; + return { + id: modelOptionId(model.id), + name: model.label ?? model.id, + modelPath: model.id, + avatarUrl: getHfAvatarUrl(model.id.replace(/^huggingface\//, '')), + recommended: Boolean(model.recommended), + }; +}; + +const readApiErrorMessage = async (res: Response, fallback: string): Promise => { + try { + const data = await res.json(); + const detail = data?.detail; + if (typeof detail === 'string') return detail; + if (detail && typeof detail.message === 'string') return detail.message; + if (detail && typeof detail.error === 'string') return detail.error; + } catch { + /* ignore malformed error bodies */ + } + return fallback; }; interface ChatInputProps { sessionId?: string; + initialModelPath?: string | null; onSend: (text: string) => void; onStop?: () => void; + onDatasetUploaded?: () => Promise | boolean; isProcessing?: boolean; disabled?: boolean; placeholder?: string; } -const isClaudeModel = (m: ModelOption) => m.modelPath.startsWith('anthropic/'); -const firstFreeModel = () => MODEL_OPTIONS.find(m => !isClaudeModel(m)) ?? MODEL_OPTIONS[0]; +interface DatasetUploadResponse { + session_id: string; + repo_id: string; + repo_type: 'dataset'; + private: true; + upload_id: string; + config_name: string; + filename: string; + path_in_repo: string; + size_bytes: number; + format: 'csv' | 'json' | 'jsonl'; + hub_url: string; + load_dataset_snippet: string; +} + +const MAX_DATASET_UPLOAD_BYTES = 100 * 1024 * 1024; +const DATASET_UPLOAD_ACCEPT = '.csv,.json,.jsonl'; +const DATASET_UPLOAD_EXTENSIONS = new Set(['csv', 'json', 'jsonl']); + +const isClaudeModel = (m: ModelOption) => isClaudePath(m.modelPath); + +const formatBytes = (bytes: number) => { + if (bytes < 1024) return `${bytes} B`; + if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB`; + return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; +}; + +const datasetRepoUrl = (repoId: string) => ( + `https://huggingface.co/datasets/${repoId.split('/').map(encodeURIComponent).join('/')}` +); -export default function ChatInput({ sessionId, onSend, onStop, isProcessing = false, disabled = false, placeholder = 'Ask anything...' }: ChatInputProps) { +export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, onDatasetUploaded, isProcessing = false, disabled = false, placeholder = 'Ask anything...' }: ChatInputProps) { const [input, setInput] = useState(''); const inputRef = useRef(null); - const [selectedModelId, setSelectedModelId] = useState(MODEL_OPTIONS[0].id); + const fileInputRef = useRef(null); + const [modelOptions, setModelOptions] = useState(DEFAULT_MODEL_OPTIONS); + const modelOptionsRef = useRef(DEFAULT_MODEL_OPTIONS); + const sessionIdRef = useRef(sessionId); + const [selectedModelPath, setSelectedModelPath] = useState( + () => ( + findModelByPath(initialModelPath ?? '', DEFAULT_MODEL_OPTIONS)?.modelPath + ?? DEFAULT_MODEL_PATH + ), + ); const [modelAnchorEl, setModelAnchorEl] = useState(null); - const { quota, refresh: refreshQuota } = useUserQuota(); - // The daily-cap dialog is triggered from two places: (a) a 429 returned - // from the chat transport when the user tries to send on Opus over cap — - // surfaced via the agent-store flag — and (b) nothing else right now - // (switching models is free). Keeping the open state in the store means - // the hook layer can flip it without threading props through. - const claudeQuotaExhausted = useAgentStore((s) => s.claudeQuotaExhausted); - const setClaudeQuotaExhausted = useAgentStore((s) => s.setClaudeQuotaExhausted); - const lastSentRef = useRef(''); + const jobsUpgradeRequired = useAgentStore((s) => s.jobsUpgradeRequired); + const setJobsUpgradeRequired = useAgentStore((s) => s.setJobsUpgradeRequired); + const updateSessionModel = useSessionStore((s) => s.updateSessionModel); + const [awaitingTopUp, setAwaitingTopUp] = useState(false); + const [modelSwitchError, setModelSwitchError] = useState(null); + const [datasetUploadError, setDatasetUploadError] = useState(null); + const [datasetUploadSuccess, setDatasetUploadSuccess] = useState(null); + const [uploadedDatasets, setUploadedDatasets] = useState([]); + const [isUploadingDataset, setIsUploadingDataset] = useState(false); + const [datasetUploadProgress, setDatasetUploadProgress] = useState(null); + + useEffect(() => { + modelOptionsRef.current = modelOptions; + }, [modelOptions]); + + useEffect(() => { + sessionIdRef.current = sessionId; + }, [sessionId]); + + useEffect(() => { + let cancelled = false; + apiFetch('/api/config/model') + .then((res) => (res.ok ? res.json() : null)) + .then((data) => { + if (cancelled || !data?.available) return; + const next = data.available + .map(modelOptionFromApi) + .filter((model: ModelOption | null): model is ModelOption => model !== null); + if (!next.length) return; + modelOptionsRef.current = next; + setModelOptions(next); + if (!sessionIdRef.current) { + const current = data.current ? findModelByPath(data.current, next) : null; + if (current) setSelectedModelPath(current.modelPath); + } + }) + .catch(() => { /* ignore */ }); + return () => { cancelled = true; }; + }, []); // Model is per-session: fetch this tab's current model every time the // session changes. Other tabs keep their own selections independently. @@ -98,15 +255,24 @@ export default function ChatInput({ sessionId, onSend, onStop, isProcessing = fa .then((data) => { if (cancelled) return; if (data?.model) { - const model = findModelByPath(data.model); - if (model) setSelectedModelId(model.id); + const model = findModelByPath(data.model, modelOptionsRef.current); + setSelectedModelPath(model?.modelPath ?? data.model); + updateSessionModel(sessionId, data.model); } }) .catch(() => { /* ignore */ }); return () => { cancelled = true; }; - }, [sessionId]); + }, [sessionId, updateSessionModel]); - const selectedModel = MODEL_OPTIONS.find(m => m.id === selectedModelId) || MODEL_OPTIONS[0]; + const visibleModelOptions = modelOptions; + const selectedModel = ( + findModelByPath(selectedModelPath, visibleModelOptions) + || findModelByPath(selectedModelPath, modelOptions) + || visibleModelOptions.find(m => m.recommended) + || modelOptions.find(m => m.recommended) + || visibleModelOptions[0] + || modelOptions[0] + ); // Auto-focus the textarea when the session becomes ready useEffect(() => { @@ -116,27 +282,86 @@ export default function ChatInput({ sessionId, onSend, onStop, isProcessing = fa }, [disabled, isProcessing]); const handleSend = useCallback(() => { - if (input.trim() && !disabled) { - lastSentRef.current = input; + if (input.trim() && !disabled && !isUploadingDataset) { onSend(input); setInput(''); } - }, [input, disabled, onSend]); + }, [input, disabled, isUploadingDataset, onSend]); + + const handleDatasetUploadClick = useCallback(() => { + fileInputRef.current?.click(); + }, []); + + const handleDatasetFileChange = useCallback( + async (event: React.ChangeEvent) => { + const file = event.target.files?.[0]; + event.target.value = ''; + if (!file) return; + + if (!sessionId) { + setDatasetUploadError('Start a session before uploading a dataset.'); + return; + } + + const extension = file.name.split('.').pop()?.toLowerCase() || ''; + if (!DATASET_UPLOAD_EXTENSIONS.has(extension)) { + setDatasetUploadError('Only CSV, JSON, and JSONL dataset files are supported.'); + return; + } + if (file.size > MAX_DATASET_UPLOAD_BYTES) { + setDatasetUploadError( + `Dataset files must be 100 MB or smaller. ${file.name} is ${formatBytes(file.size)}.` + ); + return; + } + if (file.size === 0) { + setDatasetUploadError('Uploaded dataset file is empty.'); + return; + } + + const formData = new FormData(); + formData.append('file', file); + setIsUploadingDataset(true); + setDatasetUploadProgress(0); + setDatasetUploadError(null); + setDatasetUploadSuccess(null); + try { + const res = await apiUpload(`/api/session/${sessionId}/datasets`, formData, { + onProgress: ({ percent }) => { + setDatasetUploadProgress(percent !== null && percent < 100 ? percent : null); + }, + }); + if (!res.ok) { + setDatasetUploadError(await readApiErrorMessage(res, 'Dataset upload failed.')); + return; + } + const payload = await res.json() as DatasetUploadResponse; + setUploadedDatasets((previous) => [payload, ...previous]); + setDatasetUploadSuccess(`Uploaded ${payload.filename} to ${payload.repo_id}`); + await onDatasetUploaded?.(); + } catch (error) { + setDatasetUploadError( + error instanceof Error ? error.message : 'Dataset upload failed.' + ); + } finally { + setIsUploadingDataset(false); + setDatasetUploadProgress(null); + } + }, + [sessionId, onDatasetUploaded], + ); - // When the chat transport reports a Claude-quota 429, restore the typed - // text so the user doesn't lose their message. useEffect(() => { - if (claudeQuotaExhausted && lastSentRef.current) { - setInput(lastSentRef.current); - } - }, [claudeQuotaExhausted]); + if (!datasetUploadError) return; + const timeout = window.setTimeout(() => setDatasetUploadError(null), 7000); + return () => window.clearTimeout(timeout); + }, [datasetUploadError]); - // Refresh the quota display whenever the session changes (user might - // have started another tab that spent quota). useEffect(() => { - if (sessionId) refreshQuota(); - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [sessionId]); + if (!datasetUploadSuccess) return; + const timeout = window.setTimeout(() => setDatasetUploadSuccess(null), 5000); + return () => window.clearTimeout(timeout); + }, [datasetUploadSuccess]); const handleKeyDown = useCallback( (e: KeyboardEvent) => { @@ -164,48 +389,58 @@ export default function ChatInput({ sessionId, onSend, onStop, isProcessing = fa method: 'POST', body: JSON.stringify({ model: model.modelPath }), }); - if (res.ok) setSelectedModelId(model.id); - } catch { /* ignore */ } + if (res.ok) { + setSelectedModelPath(model.modelPath); + updateSessionModel(sessionId, model.modelPath); + setModelSwitchError(null); + return; + } + setModelSwitchError(await readApiErrorMessage(res, 'Could not switch model.')); + } catch (error) { + setModelSwitchError(error instanceof Error ? error.message : 'Could not switch model.'); + } }; - // Dialog close: just clear the flag. The typed text is already restored. - const handleCapDialogClose = useCallback(() => { - setClaudeQuotaExhausted(false); - }, [setClaudeQuotaExhausted]); + const handleJobsUpgradeClose = useCallback(() => { + setJobsUpgradeRequired(null); + setAwaitingTopUp(false); + }, [setJobsUpgradeRequired]); - // "Use a free model" — switch the current session to Kimi (or the first - // non-Anthropic option) and auto-retry the send that tripped the cap. - const handleUseFreeModel = useCallback(async () => { - setClaudeQuotaExhausted(false); - if (!sessionId) return; - const free = MODEL_OPTIONS.find(m => m.modelPath === FIRST_FREE_MODEL_PATH) - ?? firstFreeModel(); + const handleJobsUpgradeClick = useCallback(async () => { + setAwaitingTopUp(true); + if (!sessionId || !jobsUpgradeRequired) return; try { - const res = await apiFetch(`/api/session/${sessionId}/model`, { + await apiFetch(`/api/pro-click/${sessionId}`, { method: 'POST', - body: JSON.stringify({ model: free.modelPath }), + body: JSON.stringify({ source: 'hf_jobs_billing_dialog', target: 'hf_billing' }), }); - if (res.ok) { - setSelectedModelId(free.id); - const retryText = lastSentRef.current; - if (retryText) { - onSend(retryText); - setInput(''); - lastSentRef.current = ''; - } - } - } catch { /* ignore */ } - }, [sessionId, onSend, setClaudeQuotaExhausted]); - - // Hide the chip until the user has actually burned quota — an unused - // Opus session shouldn't populate a counter. - const claudeChip = (() => { - if (!quota || quota.claudeUsedToday === 0) return null; - if (quota.plan === 'free') { - return quota.claudeRemaining > 0 ? 'Free today' : 'Pro only'; + } catch { + /* tracking is best-effort */ } - return `${quota.claudeUsedToday}/${quota.claudeDailyCap} today`; - })(); + }, [sessionId, jobsUpgradeRequired]); + + const handleJobsRetry = useCallback(() => { + const namespace = jobsUpgradeRequired?.namespace; + setJobsUpgradeRequired(null); + setAwaitingTopUp(false); + const msg = namespace + ? `I just added credits to the \`${namespace}\` namespace. Please retry the previous job.` + : "I just added credits. Please retry the previous job."; + onSend(msg); + }, [jobsUpgradeRequired, setJobsUpgradeRequired, onSend]); + + // Auto-retry when the user comes back to this tab after clicking "Add credits". + // Browsers fire visibilitychange when the tab regains focus from a sibling tab. + useEffect(() => { + if (!awaitingTopUp || !jobsUpgradeRequired) return; + const onVisible = () => { + if (document.visibilityState === 'visible') { + handleJobsRetry(); + } + }; + document.addEventListener('visibilitychange', onVisible); + return () => document.removeEventListener('visibilitychange', onVisible); + }, [awaitingTopUp, jobsUpgradeRequired, handleJobsRetry]); return ( + + + + + + + + + + {isProcessing ? ( )} + {isUploadingDataset && ( + + + + )} + {(datasetUploadError || datasetUploadSuccess) && ( + + { + setDatasetUploadError(null); + setDatasetUploadSuccess(null); + }} + sx={{ fontSize: '0.8rem', maxWidth: 520, width: '100%' }} + > + {datasetUploadError ?? datasetUploadSuccess} + + + )} + {uploadedDatasets.length > 0 && ( + + {uploadedDatasets.map((dataset) => ( + + ))} + + )} {/* Powered By Badge */} - {MODEL_OPTIONS.map((model) => ( + {visibleModelOptions.map((model) => ( handleSelectModel(model)} - selected={selectedModelId === model.id} + selected={selectedModel.modelPath === model.modelPath} sx={{ py: 1.5, '&.Mui-selected': { @@ -405,37 +739,36 @@ export default function ChatInput({ sessionId, onSend, onStop, isProcessing = fa }} /> )} - {isClaudeModel(model) && claudeChip && ( - - )} } - secondary={model.description} - secondaryTypographyProps={{ - sx: { fontSize: '12px', color: 'var(--muted-text)' } - }} /> ))} - + setModelSwitchError(null)} + autoHideDuration={6000} + > + setModelSwitchError(null)} + sx={{ fontSize: '0.8rem', maxWidth: 480 }} + > + {modelSwitchError} + + ); diff --git a/frontend/src/components/Chat/ExpiredBanner.tsx b/frontend/src/components/Chat/ExpiredBanner.tsx index d2bd790ab..32f638c24 100644 --- a/frontend/src/components/Chat/ExpiredBanner.tsx +++ b/frontend/src/components/Chat/ExpiredBanner.tsx @@ -18,7 +18,7 @@ interface Props { } export default function ExpiredBanner({ sessionId }: Props) { - const { renameSession, deleteSession } = useSessionStore(); + const { renameSession, deleteSession, updateSessionModel } = useSessionStore(); const [busy, setBusy] = useState<'catch-up' | 'start-over' | null>(null); const [error, setError] = useState(null); @@ -50,12 +50,13 @@ export default function ExpiredBanner({ sessionId }: Props) { useAgentStore.getState().clearSessionState(sessionId); renameSession(sessionId, newId); + if (data.model) updateSessionModel(newId, data.model); } catch (e) { logger.warn('Catch-up failed:', e); setError("Couldn't catch up — try starting over."); setBusy(null); } - }, [sessionId, renameSession]); + }, [sessionId, renameSession, updateSessionModel]); const handleStartOver = useCallback(() => { setBusy('start-over'); diff --git a/frontend/src/components/Chat/MarkdownContent.tsx b/frontend/src/components/Chat/MarkdownContent.tsx index aaab83eb1..0d1e69171 100644 --- a/frontend/src/components/Chat/MarkdownContent.tsx +++ b/frontend/src/components/Chat/MarkdownContent.tsx @@ -1,4 +1,4 @@ -import { useMemo, useRef, useState, useEffect } from 'react'; +import { useMemo, useRef, useState, useEffect, type ComponentPropsWithoutRef } from 'react'; import { Box } from '@mui/material'; import ReactMarkdown from 'react-markdown'; import remarkGfm from 'remark-gfm'; @@ -166,9 +166,17 @@ export default function MarkdownContent({ content, sx, isStreaming = false }: Ma const remarkPlugins = useMemo(() => [remarkGfm], []); + const components = useMemo(() => ({ + a: ({ href, children, ...props }: ComponentPropsWithoutRef<'a'>) => ( + + {children} + + ), + }), []); + return ( - {displayContent} + {displayContent} ); } diff --git a/frontend/src/components/Chat/MessageBubble.tsx b/frontend/src/components/Chat/MessageBubble.tsx index ede0f8a7f..ab971205c 100644 --- a/frontend/src/components/Chat/MessageBubble.tsx +++ b/frontend/src/components/Chat/MessageBubble.tsx @@ -9,6 +9,7 @@ interface MessageBubbleProps { onEditAndRegenerate?: (messageId: string, newText: string) => void | Promise; isProcessing?: boolean; isStreaming?: boolean; + sessionId?: string | null; approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise; } @@ -19,6 +20,7 @@ export default function MessageBubble({ onEditAndRegenerate, isProcessing = false, isStreaming = false, + sessionId, approveTools, }: MessageBubbleProps) { if (message.role === 'user') { @@ -38,6 +40,7 @@ export default function MessageBubble({ ); diff --git a/frontend/src/components/Chat/MessageList.tsx b/frontend/src/components/Chat/MessageList.tsx index b50a66626..5e3efcaea 100644 --- a/frontend/src/components/Chat/MessageList.tsx +++ b/frontend/src/components/Chat/MessageList.tsx @@ -8,6 +8,7 @@ import type { UIMessage } from 'ai'; interface MessageListProps { messages: UIMessage[]; isProcessing: boolean; + sessionId?: string | null; approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise; onUndoLastTurn: () => void | Promise; onEditAndRegenerate?: (messageId: string, newText: string) => void | Promise; @@ -57,7 +58,7 @@ function WelcomeGreeting() { ); } -export default function MessageList({ messages, isProcessing, approveTools, onUndoLastTurn, onEditAndRegenerate }: MessageListProps) { +export default function MessageList({ messages, isProcessing, sessionId, approveTools, onUndoLastTurn, onEditAndRegenerate }: MessageListProps) { const scrollContainerRef = useRef(null); const stickToBottom = useRef(true); @@ -139,6 +140,7 @@ export default function MessageList({ messages, isProcessing, approveTools, onUn onEditAndRegenerate={onEditAndRegenerate} isProcessing={isProcessing} isStreaming={isProcessing && msg.id === lastAssistantId} + sessionId={sessionId} approveTools={approveTools} /> )) diff --git a/frontend/src/components/Chat/ThinkingIndicator.tsx b/frontend/src/components/Chat/ThinkingIndicator.tsx deleted file mode 100644 index b8c37181f..000000000 --- a/frontend/src/components/Chat/ThinkingIndicator.tsx +++ /dev/null @@ -1,48 +0,0 @@ -import { Box, Typography } from '@mui/material'; - -/** Pulsing dots shown while the agent is processing. */ -export default function ThinkingIndicator() { - return ( - - - Thinking - - - - - - - - ); -} diff --git a/frontend/src/components/Chat/ToolCallGroup.tsx b/frontend/src/components/Chat/ToolCallGroup.tsx index fc9fe35c1..36fe77617 100644 --- a/frontend/src/components/Chat/ToolCallGroup.tsx +++ b/frontend/src/components/Chat/ToolCallGroup.tsx @@ -1,5 +1,5 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; -import { Box, Stack, Typography, Chip, Button, TextField, IconButton, Link, CircularProgress } from '@mui/material'; +import { Alert, Box, Stack, Typography, Chip, Button, TextField, IconButton, Link, CircularProgress } from '@mui/material'; import CheckCircleOutlineIcon from '@mui/icons-material/CheckCircleOutline'; import ErrorOutlineIcon from '@mui/icons-material/ErrorOutline'; import OpenInNewIcon from '@mui/icons-material/OpenInNew'; @@ -11,6 +11,8 @@ import { useAgentStore, type ResearchAgentState } from '@/store/agentStore'; import { useLayoutStore } from '@/store/layoutStore'; import { logger } from '@/utils/logger'; import { RESEARCH_MAX_STEPS } from '@/lib/research-store'; +import { useSessionStore } from '@/store/sessionStore'; +import { apiFetch } from '@/utils/api'; import type { UIMessage } from 'ai'; // --------------------------------------------------------------------------- @@ -20,6 +22,47 @@ type DynamicToolPart = Extract | undefined): number { + const current = numberOrNull(args?.current_spend_usd) ?? 0; + const cap = numberOrNull(args?.cap_usd) ?? current; + const estimate = numberOrNull(args?.estimated_next_usd) ?? 0; + return Math.ceil(Math.max(cap, current + estimate, current) + 5); +} + /** Check if a tool part was cancelled (output-error with cancellation message). */ function isCancelledTool(tool: DynamicToolPart): boolean { return tool.state === 'output-error' && @@ -220,6 +263,194 @@ function ResearchSteps({ steps }: { steps: string[] }) { ); } +// --------------------------------------------------------------------------- +// Trackio dashboard embed +// --------------------------------------------------------------------------- + +// HF repo IDs are `/` where each segment is alphanumerics plus +// `_`, `.`, `-`. Anything else (slashes, spaces, query params, missing owner) +// would let an attacker-controlled string redirect the embed to a different +// Space, so we refuse to render rather than build a malformed URL. +const SPACE_ID_PATTERN = /^[a-zA-Z0-9_.-]+\/[a-zA-Z0-9_.-]+$/; + +function isValidSpaceId(spaceId: string): boolean { + return SPACE_ID_PATTERN.test(spaceId); +} + +/** HF Space embed subdomain: 'user/space_name' → 'user-space-name'. */ +function spaceIdToSubdomain(spaceId: string): string { + return spaceId + .toLowerCase() + .replace(/[/_.]/g, '-') + .replace(/-+/g, '-') + .replace(/^-|-$/g, ''); +} + +function buildTrackioEmbedUrl(spaceId: string, project?: string): string { + // __theme=dark is gradio's standard query param to force the embedded + // dashboard into dark mode so it blends with the surrounding chat instead + // of flashing a bright white panel inside the dark UI. + const params = new URLSearchParams({ + sidebar: 'hidden', + footer: 'false', + __theme: 'dark', + }); + if (project) params.set('project', project); + return `https://${spaceIdToSubdomain(spaceId)}.hf.space/?${params.toString()}`; +} + +function buildTrackioPageUrl(spaceId: string, project?: string): string { + const qs = project ? `?${new URLSearchParams({ project }).toString()}` : ''; + return `https://huggingface.co/spaces/${spaceId}${qs}`; +} + +function TrackioEmbed({ spaceId, project }: { spaceId: string; project?: string }) { + const [expanded, setExpanded] = useState(true); + const [iframeLoaded, setIframeLoaded] = useState(false); + const embedUrl = useMemo(() => buildTrackioEmbedUrl(spaceId, project), [spaceId, project]); + const pageUrl = useMemo(() => buildTrackioPageUrl(spaceId, project), [spaceId, project]); + const label = project ? `${spaceId} · ${project}` : spaceId; + + if (!isValidSpaceId(spaceId)) return null; + + return ( + + + e.stopPropagation()} + sx={{ + px: 1.25, + py: 0.5, + borderBottom: expanded ? '1px solid var(--tool-border)' : 'none', + }} + > + + trackio + + + {label} + + e.stopPropagation()} + sx={{ + display: 'inline-flex', + alignItems: 'center', + gap: 0.4, + color: 'var(--accent-yellow)', + fontSize: '0.65rem', + textDecoration: 'none', + '&:hover': { textDecoration: 'underline' }, + }} + > + + Open + + + + {expanded && ( + +