Skip to content

feat(oauth_state_manager): public cleanup API, TTL cleanup, and tests#1371

Open
dive2tech wants to merge 2 commits intoeigent-ai:mainfrom
dive2tech:feat/oauth-state-manager-robustness
Open

feat(oauth_state_manager): public cleanup API, TTL cleanup, and tests#1371
dive2tech wants to merge 2 commits intoeigent-ai:mainfrom
dive2tech:feat/oauth-state-manager-robustness

Conversation

@dive2tech
Copy link
Contributor

Summary

This pull request hardens the OAuth authorization state manager by adding a public cleanup API, optional TTL-based pruning of completed flows, and a dedicated test suite. It improves the robustness of long-lived background authorization flows and removes the last direct access to internal state from callers.

Motivation

  • Robustness: Callers currently reach into oauth_state_manager._states to clear state, which is brittle and bypasses synchronization.
  • Lifecycle clarity: There is no dedicated mechanism to remove completed or abandoned OAuth states, so the in-memory map can accumulate stale entries.
  • Test coverage: The OAuth state manager previously had no unit tests, leaving cancellation, status transitions, and concurrency behavior unverified.

Changes

  • backend/app/utils/oauth_state_manager.py
    • Added TERMINAL_STATUSES constant (success, failed, cancelled) and used it in update_status and cleanup logic.
    • Added remove_state(provider):
      • Validates the provider name (rejects None/empty).
      • Cancels the state if it is still pending or authorizing before removal.
      • Removes the state from the internal map under the existing lock.
    • Added list_states():
      • Returns all current states as dictionaries via OAuthState.to_dict(), suitable for diagnostics or a future API.
    • Added clear_completed(max_age_seconds: int | None = None):
      • Removes states in a terminal status (success/failed/cancelled) that have a non-null completed_at.
      • When max_age_seconds is provided, only removes states that completed at least that many seconds ago.
      • Returns the number of states removed and emits a log entry when any states are pruned.
    • Updated update_status:
      • Uses TERMINAL_STATUSES to decide when to set completed_at.
      • Switched logging to parameterized style (logger.info("...", provider, status)).
    • backend/app/controller/tool_controller.py
      • Replaced direct access to the internal _states map when uninstalling Google Calendar:
        • Before: checked the state and called state.cancel() then oauth_state_manager._states.pop("google_calendar", None).
        • After: uses oauth_state_manager.remove_state("google_calendar") and logs when the cache is cleared.
      • This keeps all state transitions and cleanup inside OAuthStateManager, honoring its locking and lifecycle rules.

Tests

  • backend/tests/app/utils/test_oauth_state_manager.py (new)
    • OAuthState tests:
      • Initialization defaults (provider, status, error, completed_at).
      • to_dict() for both pending and completed states.
      • is_cancelled() before and after cancel().
      • cancel() sets the cancel event, status, and completed_at.
    • OAuthStateManager tests:
      • create_state / get_state create and retrieve a state for a provider.
      • Creating a second state for the same provider cancels the previous one and replaces it.
      • update_status sets status, error, and completed_at for success and failed.
      • update_status is a no-op for unknown providers.
      • remove_state:
        • Returns False for None/empty provider names and missing providers.
        • Returns True when a state is present and removes it.
        • Ensures pending/authorizing states are cancelled on removal.
      • list_states:
        • Returns an empty list when there are no states.
        • Returns dictionaries with provider, status, and timestamps when states exist.
      • clear_completed:
        • Returns 0 when there are no states.
        • Removes terminal states and leaves pending ones.
        • Honors max_age_seconds so very recent completions are retained while older ones are pruned.
      • Basic concurrency check:
        • Two threads call create_state for the same provider; the manager remains consistent and returns a non-None state at the end.

dive2tech and others added 2 commits February 25, 2026 14:59
- Add remove_state(provider) to drop state and cancel if in progress; validate provider
- Add list_states() for API/debugging
- Add clear_completed(max_age_seconds) to prune terminal states with optional TTL
- Add TERMINAL_STATUSES constant; use in update_status and clear_completed
- tool_controller: use remove_state('google_calendar') instead of _states.pop
- Add tests/app/utils/test_oauth_state_manager.py for OAuthState, create_state, get_state, update_status, remove_state, list_states, clear_completed, and concurrency

Co-authored-by: Cursor <cursoragent@cursor.com>
…, remove unused)

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant