Skip to content

fix(cli): graceful BudgetExceededError — wrapper-only, no Agent param bloat (fixes #1627) #1645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MervinPraison merged 1 commit into from
May 9, 2026
Merged

fix(cli): graceful BudgetExceededError — wrapper-only, no Agent param bloat (fixes #1627) #1645

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 53 additions & 28 deletions src/praisonai/praisonai/cli/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4065,6 +4065,31 @@ def _run_inline_workflow(self, initial_prompt):

return results[-1].get("output", "") if results else ""

def _execute_agent_with_budget_handling(self, agent, method_name, *args, **kwargs):
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This helper method should also be applied to the agent.chat(full_prompt) call site within _run_inline_workflow (around line 4030) to ensure consistent budget error handling across all direct-prompt execution paths.

"""Run ``agent.<method_name>(*args, **kwargs)`` with a graceful
BudgetExceededError handler.

Wrapper-only fix (no core SDK changes). Users configure budgets via
``execution=ExecutionConfig(max_budget=...)`` on the Agent — per
AGENTS.md §5.3 there is NO top-level ``max_budget=`` parameter on
Agent.__init__ (avoids parameter bloat).

When the budget is hit this prints a single-line actionable error
message and exits with code 1 instead of leaking a raw traceback.
Any other exception is re-raised unchanged.
"""
from praisonaiagents.errors import BudgetExceededError
try:
return getattr(agent, method_name)(*args, **kwargs)
except BudgetExceededError as e:
from rich import print as rich_print
rich_print(
f"[red]Budget limit exceeded: {e!s}. "
"Hint: set budget via "
"execution=ExecutionConfig(max_budget=1.00) on your Agent.[/red]"
)
Comment on lines +4084 to +4090
Copy link
Copy Markdown

@greptile-apps greptile-apps Bot May 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 PraisonAIError.__str__() returns "[budget] <message> (agent: ..., run: ...)". When {e!s} is interpolated directly into a Rich markup string, the [budget] token is treated as a markup tag by Rich. Since budget is not a recognised style, Rich raises a MarkupError, causing the error handler itself to crash — the opposite of graceful recovery. The fix is to escape the exception message before embedding it in the markup string.

Suggested change
except BudgetExceededError as e:
from rich import print as rich_print
rich_print(
f"[red]Budget limit exceeded: {e!s}. "
"Hint: set budget via "
"execution=ExecutionConfig(max_budget=1.00) on your Agent.[/red]"
)
except BudgetExceededError as e:
from rich import print as rich_print
from rich.markup import escape as _escape
rich_print(
f"[red]Budget limit exceeded: {_escape(str(e))}. "
"Hint: set budget via "
"execution=ExecutionConfig(max_budget=1.00) on your Agent.[/red]"
)

sys.exit(1)
Comment on lines +4086 to +4091
Copy link
Copy Markdown

@greptile-apps greptile-apps Bot May 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Error output corrupts structured stdout modes

When display_mode is jsonl or json, the budget error message is printed to stdout via rich_print, injecting a plain-text line into what is supposed to be a machine-readable JSON or JSONL stream. Any downstream consumer that parses stdout will fail to decode the envelope. Printing the message to stderr instead (e.g. rich_print("...", file=sys.stderr)) would preserve the structured stdout contract while still surfacing the error to the user.

Comment on lines +4084 to +4091
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This error handling block has two issues:

  1. Cleanup Bypass: Calling sys.exit(1) here prevents the execution of cleanup logic at the call sites (e.g., disable_status_output(), flow.stop(), disable_editor_output()). To fix this, call sites should wrap the execution in a try...finally block to ensure terminal state is restored even on exit.
  2. Output Corruption: Printing the error to stdout can break machine-readable outputs like JSON/JSONL. It should be directed to stderr.

Additionally, it is recommended to escape the exception message to prevent rich from misinterpreting characters like [ or ] as markup tags.

Suggested change
except BudgetExceededError as e:
from rich import print as rich_print
rich_print(
f"[red]Budget limit exceeded: {e!s}. "
"Hint: set budget via "
"execution=ExecutionConfig(max_budget=1.00) on your Agent.[/red]"
)
sys.exit(1)
except BudgetExceededError as e:
from rich.console import Console
from rich.markup import escape
Console(stderr=True).print(
f"[red]Budget limit exceeded: {escape(str(e))}. "
"Hint: set budget via "
"execution=ExecutionConfig(max_budget=1.00) on your Agent.[/red]"
)
sys.exit(1)


def _extract_cli_config_for_yaml(self):
"""
Extract CLI configuration that should be passed to YAML processing.
Expand Down Expand Up @@ -4611,26 +4636,26 @@ def level_based_approve(function_name, arguments, risk_level):
from rich.panel import Panel

with Live(Panel(Spinner("dots", text="Generating..."), border_style="cyan"), refresh_per_second=10, transient=True):
result = auto_rag.chat(prompt)
result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
else:
result = auto_rag.chat(prompt)
result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
else:
# Resolve display mode from CLI flags
display_mode = self._resolve_display_mode()

if display_mode == 'silent':
# -qq: No output at all, exit code only
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

elif display_mode == 'quiet':
# -q: Result only, no spinners or status
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
if result is not None:
output = getattr(result, 'output', None) or (str(result) if result else None)
if output:
Expand All @@ -4642,31 +4667,31 @@ def level_based_approve(function_name, arguments, risk_level):
from praisonaiagents.output.status import enable_status_output, disable_status_output
enable_status_output(show_timestamps=True, show_metrics=True)
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
disable_status_output()
except ImportError:
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

elif display_mode == 'debug':
# -vv: SDK TraceOutput with markdown rendering
try:
from praisonaiagents.output.trace import enable_trace_output, disable_trace_output
enable_trace_output(use_markdown=True)
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
disable_trace_output()
except ImportError:
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

elif display_mode == 'jsonl':
# --output jsonl: JSONL structured output for CI/CD
Expand All @@ -4688,9 +4713,9 @@ def level_based_approve(function_name, arguments, risk_level):

start_time = time.time()
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

# Emit final result
reason = getattr(result, 'completion_reason', None) if hasattr(result, 'completion_reason') else 'complete'
Expand All @@ -4709,9 +4734,9 @@ def level_based_approve(function_name, arguments, risk_level):
import json as json_mod
start_time = time.time()
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

output = result.output if hasattr(result, 'output') else str(result)
envelope = {
Expand All @@ -4732,15 +4757,15 @@ def level_based_approve(function_name, arguments, risk_level):
flow = track_workflow()
flow.start()
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
flow.stop()
except ImportError:
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

elif display_mode == 'editor':
# --output editor: User-friendly step-by-step format
Expand All @@ -4750,9 +4775,9 @@ def level_based_approve(function_name, arguments, risk_level):

# Run agent
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

# SDK callbacks (interaction, llm_content) handle display —
# no explicit editor.output() needed here.
Expand All @@ -4774,15 +4799,15 @@ def level_based_approve(function_name, arguments, risk_level):
from praisonaiagents.output.status import enable_status_output, disable_status_output
enable_status_output(show_timestamps=False, show_metrics=False)
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
disable_status_output()
except ImportError:
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

# ===== POST-PROCESSING WITH NEW FEATURES =====

Expand Down
Loading