refactor: harden BudgetEnforcer with error handling and review fixes #182
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,59 @@ | ||
| """Billing period computation utilities. | ||
|
|
||
| Pure functions for determining billing period boundaries based on a | ||
| configurable reset day. Used by :class:`~ai_company.budget.enforcer.BudgetEnforcer` | ||
| to scope cost queries to the current billing cycle. | ||
| """ | ||
|
|
||
| from datetime import UTC, datetime | ||
|
|
||
|
|
||
| def billing_period_start( | ||
| reset_day: int, | ||
| *, | ||
| now: datetime | None = None, | ||
| ) -> datetime: | ||
| """Compute the UTC-aware start of the current billing period. | ||
|
|
||
| If ``now.day >= reset_day``, returns current month's ``reset_day`` | ||
| at 00:00 UTC. Otherwise, returns previous month's ``reset_day`` | ||
| at 00:00 UTC. | ||
|
|
||
| Args: | ||
| reset_day: Day of month when the billing period resets (1-28). | ||
| now: Reference timestamp. Defaults to ``datetime.now(UTC)``. | ||
|
|
||
| Returns: | ||
| UTC-aware datetime at midnight on the billing period start day. | ||
|
|
||
| Raises: | ||
| ValueError: If ``reset_day`` is not in ``[1, 28]``. | ||
| """ | ||
| if not 1 <= reset_day <= 28: # noqa: PLR2004 | ||
| msg = f"reset_day must be 1-28, got {reset_day}" | ||
| raise ValueError(msg) | ||
|
|
||
| if now is None: | ||
| now = datetime.now(UTC) | ||
|
|
||
| if now.day >= reset_day: | ||
| return datetime(now.year, now.month, reset_day, tzinfo=UTC) | ||
|
|
||
| # Roll back to previous month | ||
| if now.month == 1: | ||
| return datetime(now.year - 1, 12, reset_day, tzinfo=UTC) | ||
| return datetime(now.year, now.month - 1, reset_day, tzinfo=UTC) | ||
|
Comment on lines
+32
to
+45
Contributor
There was a problem hiding this comment. Harden the public input boundary for
Suggested hardening+def _normalize_utc_now(*, now: datetime | None) -> datetime:
+ if now is None:
+ return datetime.now(UTC)
+ if now.tzinfo is None:
+ msg = "now must be timezone-aware"
+ raise ValueError(msg)
+ return now.astimezone(UTC)
+
+
def billing_period_start(
reset_day: int,
*,
now: datetime | None = None,
) -> datetime:
@@
- if not 1 <= reset_day <= 28: # noqa: PLR2004
+ if (
+ isinstance(reset_day, bool)
+ or not isinstance(reset_day, int)
+ or not 1 <= reset_day <= 28 # noqa: PLR2004
+ ):
msg = f"reset_day must be 1-28, got {reset_day}"
raise ValueError(msg)
- if now is None:
- now = datetime.now(UTC)
+ now = _normalize_utc_now(now=now)
@@
def daily_period_start(*, now: datetime | None = None) -> datetime:
@@
- if now is None:
- now = datetime.now(UTC)
+ now = _normalize_utc_now(now=now)
return datetime(now.year, now.month, now.day, tzinfo=UTC)As per coding guidelines, "Validate at system boundaries (user input, external APIs, config files)". Also applies to: 57-59 🤖 Prompt for AI Agents |
||
|
|
||
|
|
||
| def daily_period_start(*, now: datetime | None = None) -> datetime: | ||
| """Compute the UTC-aware start of today (midnight UTC). | ||
|
|
||
| Args: | ||
| now: Reference timestamp. Defaults to ``datetime.now(UTC)``. | ||
|
|
||
| Returns: | ||
| UTC-aware datetime at midnight of the current day. | ||
| """ | ||
| if now is None: | ||
| now = datetime.now(UTC) | ||
| return datetime(now.year, now.month, now.day, tzinfo=UTC) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion | 🟠 Major
Initialize the standard logger for this module.
These helpers are now part of the budget-enforcement business logic path, but the module still doesn't define
logger = get_logger(__name__). That also leaves the invalid-input raise paths without the standard observability hook used elsewhere undersrc/ai_company/**.Suggested fix
As per coding guidelines, "Every module with business logic MUST have: from ai_company.observability import get_logger then logger = get_logger(name)".
📝 Committable suggestion
🤖 Prompt for AI Agents