fix: job token budget, iteration cap → Failed, web cancel stops worker

[ilblackdragon]

Summary

Fixes #698. Jobs could enter infinite retry loops burning millions of tokens because:

• No token budget enforcement — JobContext::add_tokens() existed but was never called. Added configurable max_tokens_per_job (settings.json agent.max_tokens_per_job / AGENT_MAX_TOKENS_PER_JOB env var, default 0 = unlimited). Token usage is tracked after respond_with_tools() and the job fails on budget exceeded. Per-job override via metadata["max_tokens"].
• Iteration cap → Stuck → self-repair loop — Changed mark_stuck() to mark_failed() for iteration cap and persistent rate limiting, so DefaultSelfRepair cannot restart them.
• Web cancel didn't stop the worker — Cancel handler now calls scheduler.stop(job_id) which updates in-memory ContextManager AND aborts the worker task. Falls back to DB-only update when scheduler is unavailable.

Test plan

• test_token_budget_exceeded_fails_job — token budget enforcement transitions to Failed
• test_iteration_cap_marks_failed_not_stuck — iteration cap transitions to Failed, not Stuck
• cargo clippy --all --all-features — zero warnings
• cargo test — 2758 passed, 0 failed
• Manual: set agent.max_tokens_per_job in settings.json, run a job, verify it stops at the budget
• Manual: cancel a running agent job from the web UI, verify the worker actually stops

🤖 Generated with Claude Code

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the stability and resource management of agent jobs. It introduces a token budget system to control LLM costs, refines job failure states to prevent undesirable retry loops, and improves the reliability of job cancellation from the web interface by ensuring worker tasks are properly terminated. These changes collectively contribute to more predictable and efficient agent operations.

Highlights

• Token Budget Enforcement: Introduced a configurable max_tokens_per_job setting to prevent jobs from consuming excessive LLM tokens, with an option for per-job overrides via metadata. Jobs exceeding this budget will now fail.
• Improved Job State Transitions: Modified the iteration cap and persistent rate limiting logic to transition jobs directly to a Failed state instead of Stuck, preventing agents from entering infinite self-repair loops.
• Enhanced Web Cancellation: Implemented a more robust job cancellation mechanism from the web UI, which now attempts to stop the associated worker task via the scheduler, ensuring immediate termination of the agent's activity.

Changelog

• .env.example
  • Added an example environment variable for AGENT_MAX_TOKENS_PER_JOB.
• src/agent/dispatcher.rs
  • Updated test configurations to include the new max_tokens_per_job setting.
• src/agent/scheduler.rs
  • Implemented logic to apply max_tokens_per_job from configuration or metadata when creating a new job.
  • Added a separate update to set the token budget in the job context.
• src/agent/worker.rs
  • Changed job state transition from Stuck to Failed when max_iterations are exceeded.
  • Changed job state transition from Stuck to Failed for persistent LLM rate limiting.
  • Integrated token usage tracking after respond_with_tools() and enforced the max_tokens budget, failing the job if exceeded.
  • Added new test cases: test_token_budget_exceeded_fails_job and test_iteration_cap_marks_failed_not_stuck.
• src/channels/web/handlers/jobs.rs
  • Modified jobs_cancel_handler to attempt stopping the worker via the scheduler first, then falling back to a DB-only status update if the scheduler is unavailable.
• src/config/agent.rs
  • Added max_tokens_per_job field to AgentConfig struct.
  • Initialized max_tokens_per_job to 0 in AgentConfig::default().
  • Implemented parsing for AGENT_MAX_TOKENS_PER_JOB from environment variables and settings.
• src/settings.rs
  • Added max_tokens_per_job field to AgentSettings struct.
  • Set default value of max_tokens_per_job to 0 in AgentSettings::default().

Activity

• Verified test_token_budget_exceeded_fails_job to confirm token budget enforcement leads to job failure.
• Verified test_iteration_cap_marks_failed_not_stuck to confirm iteration cap leads to job failure, not stuck state.
• Ran cargo clippy --all --all-features with zero warnings.
• Ran cargo test with 2758 tests passed and 0 failed.
• Planned manual testing for agent.max_tokens_per_job and web cancellation.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

Pull request overview

This PR addresses runaway agent jobs by introducing per-job token budgets, changing iteration-cap/rate-limit terminal behavior to Failed, and making web “cancel” actually stop the running worker via the scheduler.

Changes:

• Add configurable agent.max_tokens_per_job (settings + env) and plumb it into job creation with optional per-job metadata override.
• Change iteration-cap and persistent rate-limit terminal state from Stuck to Failed in the worker loop.
• Update web cancel handler to stop the worker via scheduler.stop(job_id) (with DB-only fallback).

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/settings.rs	Adds AgentSettings.max_tokens_per_job with default 0 (unlimited).
src/config/agent.rs	Resolves max_tokens_per_job from env/settings into AgentConfig.
src/agent/scheduler.rs	Applies token budget to new jobs, allowing metadata["max_tokens"] override.
src/agent/worker.rs	Marks iteration/rate-limit terminal conditions as Failed; adds token usage tracking; adds tests.
src/channels/web/handlers/jobs.rs	Uses scheduler stop on cancel to abort worker; falls back to DB update.
src/agent/dispatcher.rs	Updates tests to include new AgentConfig field.
.env.example	Documents AGENT_MAX_TOKENS_PER_JOB.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Token budget enforcement is only applied to the respond_with_tools() call when selections.is_empty(). Iterations that go through select_tools() (and any other LLM calls like planning) are not counted, so a job can still exceed the intended budget without being failed. Consider plumbing TokenUsage out of select_tools() / plan() (or adding a wrapper that records usage for every LLM call in the worker loop) so all LLM token consumption contributes to ctx.add_tokens(...).

Valid point. select_tools() makes an LLM call but does not expose TokenUsage in its return type. Plumbing usage out requires changing the Reasoning trait — a broader refactor. Added a NOTE comment acknowledging the gap in 7364930. Since there was zero token tracking before this PR, tracking respond_with_tools() alone is a meaningful improvement, and the iteration cap now also serves as a hard stop.

[Copilot]

These new tests validate JobContext::add_tokens() and Worker::mark_failed(), but they don't exercise the newly added runtime path that increments tokens based on respond_with_tools().usage inside Worker::run(). Adding a test that runs a minimal worker loop iteration with a stub LLM returning a non-zero usage would better cover the regression this PR is addressing (budget enforcement from actual LLM calls).

Acknowledged. Testing the full worker loop with a mock LLM that returns non-zero usage is substantial (requires mocking Reasoning, select_tools, respond_with_tools). The current tests cover the core path: add_tokens correctly returns an error on budget exceeded, and mark_failed correctly transitions the job. The integration between these pieces in the loop is straightforward (3 lines of glue code). A full integration test with a mock LLM would be valuable but is better suited as a follow-up.

[Copilot]

stopped_via_scheduler is computed with scheduler.stop(job_id).await.is_ok(), but Scheduler::stop() currently returns Ok(()) even when the job isn't in the in-memory scheduler map. That means this handler can skip the DB status update while the job remains active in the DB (and no worker was actually stopped). Fix by checking scheduler.is_running(job_id).await (or changing stop() to return NotFound) and only skipping the DB update when a running job was actually stopped (and/or always persist the Cancelled state).

Good catch. Fixed in 7364930 — the cancel handler now always persists Cancelled to the DB, using scheduler.stop() as an additional best-effort step (to abort the worker task) rather than a replacement for the DB update. This handles the edge case where stop() returns Ok(()) for jobs not in the scheduler map.

Code Review

This pull request introduces several important fixes to job lifecycle management, addressing potential infinite loops and token over-consumption. The changes to enforce a token budget, fail jobs on iteration caps, and ensure web cancellations properly stop worker tasks are well-implemented and crucial for stability. The new tests adequately cover the new failure conditions. I've kept the original suggestions to improve the conciseness and readability of the new logic, as they do not contradict any established rules.

_{Note: Security Review did not run due to the size of the PR.}

This logic for checking the token budget is correct, but it can be made more concise by combining the await? and the if let into a single statement. This improves readability by making the error handling path more direct.

Suggested change

	let budget_result = self
	.context_manager()
	.update_context(self.job_id, |ctx| ctx.add_tokens(total_tokens))
	.await?;
	if let Err(msg) = budget_result {
	self.mark_failed(&msg).await?;
	return Ok(());
	}
	if let Err(msg) = self
	.context_manager()
	.update_context(self.job_id, |ctx| ctx.add_tokens(total_tokens))
	.await?
	{
	self.mark_failed(&msg).await?;
	return Ok(());
	}

Applied — collapsed the nested if into a let-chain per clippy. Fixed in 7364930.

This logic to determine if the job was stopped via the scheduler is a bit nested. You can improve readability and reduce nesting by using an async block to scope the awaits and return the result directly.

let stopped_via_scheduler = async {
    if let Some(scheduler_slot) = &state.scheduler {
        if let Some(scheduler) = scheduler_slot.read().await.as_ref() {
            return scheduler.stop(job_id).await.is_ok();
        }
    }
    false
}.await;

Applied a similar simplification using let-chains (collapsed nested ifs). Fixed in 7364930.

Review

Three changes bundled here: token budget enforcement, iteration cap -> Failed (not Stuck), and web cancel stops worker via scheduler. All three are solid improvements.

Token Budget

• ctx.add_tokens() / ctx.max_tokens pattern is clean. The metadata override (max_tokens in job metadata) is a nice touch for per-job control.
• Concern: add_tokens returns Result via update_context, which returns Result<T, _> where T is the closure's return value. The double-unwrap pattern (await.unwrap() then check is_err()) in the test is fine, but in the worker loop the pattern is:

  let budget_result = self.context_manager().update_context(...).await?;
  if let Err(msg) = budget_result { ... }

  This is correct but would be clearer if add_tokens returned a named error type instead of Result<(), String>.

Iteration Cap -> Failed

• Good fix. mark_stuck triggers self-repair which retries, creating an infinite loop. mark_failed is the correct terminal state for "tried too many times."

Web Cancel via Scheduler

• The stopped_via_scheduler pattern correctly tries the in-memory scheduler first (which aborts the task handle), falling back to DB-only update.
• Edge case: If scheduler.stop() succeeds (returns Ok) but the worker hasn't persisted its final state yet, the response says "cancelled" but the DB might still show InProgress briefly. Acceptable for now but worth a comment.

Tests

Both new tests (test_token_budget_exceeded_fails_job, test_iteration_cap_marks_failed_not_stuck) are good regression tests.

Overall looks good -- approve once the existing CI passes.

Thanks for the review @zmanian!

Re: add_tokens returning Result<(), String> — agreed a named error type would be cleaner. That is pre-existing API surface in context/state.rs; changing it is a follow-up refactor.

Re: DB consistency edge case on cancel — fixed in 7364930. The handler now always persists Cancelled to the DB, using scheduler.stop() as best-effort to abort the worker. The fire-and-forget persist inside scheduler.stop() is a bonus but the handler no longer relies on it as the sole DB update.

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

.github/workflows/staging-ci.yml:122

• The actions/create-github-app-token@v2 step now runs unconditionally. If GH_RELEASES_MANAGER_APP_ID / GH_RELEASES_MANAGER_APP_PRIVATE_KEY secrets are missing or empty, this action fails the job before the later fallback to github.token can run. Restore the previous if: guard (or make the step continue-on-error and handle the empty-token case) so staging CI still works when the GitHub App secrets aren't configured.

      - name: Generate GitHub App token
        id: app-token
        uses: actions/create-github-app-token@v2
        with:
          app-id: ${{ secrets.GH_RELEASES_MANAGER_APP_ID }}
          private-key: ${{ secrets.GH_RELEASES_MANAGER_APP_PRIVATE_KEY }}

.github/workflows/staging-ci.yml:236

• Same issue here: generating the GitHub App token is unconditional, so missing/empty GH_RELEASES_MANAGER_* secrets will fail the gate job before the workflow can fall back to github.token. Add back an if: guard or make the step non-fatal and keep the fallback logic effective.

      - name: Generate GitHub App token
        id: app-token
        uses: actions/create-github-app-token@v2
        with:
          app-id: ${{ secrets.GH_RELEASES_MANAGER_APP_ID }}
          private-key: ${{ secrets.GH_RELEASES_MANAGER_APP_PRIVATE_KEY }}

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.