Add session scoping to caches. #13482

sfc-gh-jkinkead · 2025-12-29T20:50:57Z

Describe your changes

Add session-level scoping to st.cache_data and st.cache_resource.

Add the "scope" parameter through the caching stack.

Create session-scoped caches as appropriate.

Add methods to clear caches by session ID. Call those methods when sessions are disconnected.

Note: This clears caches on disconnect and shutdown. In the current websocket session manager, sessions only appear to be shut down when the backend process terminates - and so disconnection is the only hook that's actually invoked in the typical session lifecycle. I'm not sure if this is a bug or a design choice. Either way, like the docs note, this means that session caches might populate multiple times for a single user session in some edge-cases. I think this is fine.

GitHub Issue Link (if applicable)

Fix for Shut down the session after closing the browser tab by the user #8545. An on_release hook for a session-scoped resource can be used for disconnect hooks.
Can be used to implement Session State convenience function for initialization #10089. A session-scoped @st.cache_data or @st.cache_resource function invoked at the start of a script will work as an init function. Reading this more closely, the issue is for a helper function to make this easy, not for means to init at session start.
Implements one of the suggested fixes for st.cache_resource alternative for non-threadsafe object (that can't be serialized / used with st.cache_data) #6703.

Testing Plan

See unit tests.

Contribution License Agreement

By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license.

Add a "scope" parameter through the caching stack. Create session-scoped caches as appropriate. Add methods to clear caches by session ID. Call those methods when sessions expire.

snyk-io · 2025-12-29T20:51:06Z

✅ Snyk checks have passed. No issues have been found so far.

Status	Scanner	Critical	High	Medium	Low	Total (0)
✅	Open Source Security	0	0	0	0	0 issues
✅	Licenses	0	0	0	0	0 issues

💻 Catch issues earlier using the plugins for VS Code, JetBrains IDEs, Visual Studio, and Eclipse.

github-actions · 2025-12-29T20:51:11Z

✅ PR preview is ready!

Name	Link
📦 Wheel file	https://core-previews.s3-us-west-2.amazonaws.com/pr-13482/streamlit-1.52.2-py3-none-any.whl
📦 `@streamlit/component-v2-lib`	Download from artifacts
🕹️ Preview app	pr-13482.streamlit.app (☁️ Deploy here if not accessible)

lib/streamlit/runtime/caching/cache_data_api.py

Copilot

Pull request overview

This PR adds session-scoped caching to st.cache_data and st.cache_resource, allowing cache entries to be scoped either globally (default) or per-session. Session-scoped caches are automatically cleared when sessions disconnect or shut down, enabling resource cleanup and per-session initialization patterns.

Key changes:

Adds a new scope parameter ("global" or "session") to both caching decorators
Refactors cache storage from flat dictionaries to nested session-to-function-key mappings
Implements clear_session() methods to clean up session-specific caches
Integrates cache clearing into the session disconnect and shutdown lifecycle

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
lib/streamlit/runtime/caching/cache_utils.py	Adds `CacheScope` type alias and `get_session_id_or_throw()` helper function; updates `CachedFuncInfo` to accept scope parameter
lib/streamlit/runtime/caching/cache_data_api.py	Refactors DataCaches to use nested dictionaries for session scoping; adds scope parameter to decorator and `clear_session()` method
lib/streamlit/runtime/caching/cache_resource_api.py	Refactors ResourceCaches to use nested dictionaries for session scoping; adds scope parameter to decorator and `clear_session()` method
lib/streamlit/runtime/app_session.py	Adds `clear_session_caches()` method and integrates it into shutdown and disconnect flows
lib/streamlit/runtime/websocket_session_manager.py	Calls `clear_session_caches()` when sessions disconnect
lib/tests/streamlit/runtime/caching/common_cache_test.py	Adds comprehensive tests for session-scoped cache lookup, clearing, and invalid scope handling
lib/tests/streamlit/runtime/caching/cache_utils_test.py	New test file for `get_session_id_or_throw()` utility function
lib/tests/streamlit/runtime/app_session_test.py	Updates shutdown tests to verify cache clearing and adds test for `clear_session_caches()` method

lib/streamlit/runtime/caching/cache_data_api.py

lib/streamlit/runtime/caching/cache_resource_api.py

lib/streamlit/runtime/caching/cache_data_api.py

lib/streamlit/runtime/caching/cache_resource_api.py

lib/tests/streamlit/runtime/app_session_test.py

lib/tests/streamlit/runtime/caching/cache_utils_test.py

github-actions · 2025-12-31T09:14:45Z

Summary

This PR adds a new scope parameter to st.cache_data and st.cache_resource decorators, enabling session-level caching alongside the existing global caching. Key changes include:

New scope parameter: Accepts "global" (default) or "session" values
Session-scoped cache storage: Session caches are stored separately using the session ID as a key dimension
Automatic cleanup: Session caches are cleared when sessions disconnect or shut down, with on_release callbacks invoked
New clear_session() API method: Added to both cache APIs for programmatic session cache clearing

This feature addresses several GitHub issues (#8545, #6703) by enabling session-scoped resource management and disconnect hooks.

Code Quality

The implementation is clean and follows existing codebase patterns well:

Strengths:

Consistent implementation across both cache_data and cache_resource APIs
Proper type annotations using CacheScope type alias
Thread-safe implementation with appropriate locking in DataCaches and ResourceCaches
Clear error messages for invalid scope values
Well-documented with comprehensive docstrings explaining behavior and edge cases

Minor Issues:

Unused function in test file (lib/tests/streamlit/runtime/caching/cache_utils_test.py, lines 26-29):

def function_for_testing(
    pos_one: int, pos_two: int, _scope: str, pos_three: int
) -> str:
    """Dummy function for testing function caches."""
    return f"{pos_one}-{pos_two}-{_scope}-{pos_three}"

This function is defined but never used in the test file.

Missing return type annotation on clear_session methods (cache_data_api.py line 693, cache_resource_api.py line 603):
For consistency with other methods, consider adding explicit -> None return type annotation.
Double cache clearing on shutdown: clear_session_caches() is called both in app_session.py:shutdown() (line 277) and in _handle_scriptrunner_event_on_event_loop() (line 706) when shutdown is requested. While this is idempotent and safe, it may cause redundant work. This appears intentional to ensure cleanup in all scenarios (sync and async paths).

Test Coverage

The test coverage is comprehensive and well-structured:

Covered scenarios:

✅ Session scope cache lookups with multiple sessions (test_session_scope_handles_lookup)
✅ Clearing individual session caches without affecting others (test_session_scope_handles_clear)
✅ Bad scope parameter validation (test_bad_scope_raises_exception)
✅ get_session_id_or_throw function behavior (GetSessionIdOrThrowTest)
✅ clear_session_caches method in AppSession (test_clear_session_caches)
✅ Shutdown behavior calling clear_session (test_shutdown)
✅ Parameterized tests covering both cache_data and cache_resource

Tests follow best practices:

Uses @parameterized.expand to reduce code duplication
Uses proper mocking with patch.object
Tests include docstrings describing their purpose
Tests verify cache isolation between sessions

Potential additions (optional):

Test for on_release callback being invoked when session cache is cleared (for cache_resource)
Integration test for interaction between session-scoped and global-scoped caches with the same function key
Test for behavior when persist="disk" is combined with scope="session" (edge case)

Backwards Compatibility

✅ Fully backwards compatible:

The scope parameter defaults to "global", preserving existing behavior
No breaking changes to existing APIs or method signatures
Existing user code will continue to work without modification
Global cache behavior is unchanged

Security & Risk

No security concerns identified:

Memory is properly managed with cache cleanup on session disconnect
Thread safety is maintained with appropriate locking
No new attack vectors introduced

Low regression risk:

Changes are additive and well-isolated
The default behavior (scope="global") is unchanged
Edge cases are documented (e.g., reconnected sessions may repopulate cache)

Note on documented behavior:
The docstring correctly notes that "disconnected sessions can reconnect - so it is possible for the cache to populate multiple times in a single session for the same key." This is appropriate documentation of expected behavior.

Recommendations

Remove unused test function: Delete function_for_testing from cache_utils_test.py (lines 26-29) as it serves no purpose.
Add return type annotation (optional): Add -> None to clear_session methods in both API classes for consistency:
```
def clear_session(self, session_id: str) -> None:
```
Consider adding typing test (optional): Add a test case in lib/tests/streamlit/typing/cache_types.py to verify the scope parameter is properly typed:
```
@st.cache_data(scope="session")
def cached_data_fn_with_scope(arg1: int) -> bool:
    return True
```

Consider adding example (optional): Adding a usage example for session-scoped caching in the docstring would help users understand the feature:

>>> @st.cache_resource(scope="session", on_release=lambda x: x.close())
... def get_session_connection():
...     return create_connection()

Verdict

APPROVED: This is a well-implemented feature that addresses a real user need for session-scoped caching. The code quality is high, test coverage is comprehensive, and the implementation follows existing patterns in the codebase. The minor issues noted (unused test function, missing return type annotation) are not blockers and can be addressed in a follow-up if desired.

This is an automated AI review. Please verify the feedback and use your judgment.

cursor · 2025-12-31T09:26:51Z

lib/streamlit/runtime/caching/cache_data_api.py

+
+        if session_caches is not None:
+            for cache in session_caches.values():
+                cache.clear()


Storage not closed when session caches are cleared

The clear_session method in DataCaches calls cache.clear() but does not call cache.storage.close() afterwards. This is inconsistent with clear_all() (which calls both clear() and storage.close() in its fallback path) and get_cache() (which calls storage.close() when replacing a cache). For disk-persisted session-scoped caches, this could result in file handles not being released when sessions are disconnected, leading to a resource leak.

Good catch! Updating.

lukasmasuch · 2025-12-31T11:59:59Z

lib/streamlit/runtime/caching/cache_resource_api.py

        """Clear all cache_resource caches."""
        _resource_caches.clear_all()

+    def clear_session(self, session_id: str) -> None:


question: is clear_session supposed to be exposed to users via the public API? I think there isn't even an official API to retrieve the session ID. In case it should be exposed, it would be good to add a gather_metrics decorator to track its usage. cc @jrieke

Maybe its also cleaner to support this via a flag on the .clear method, e.g. clear(scope="session")

I'd casually considered this suitable for a public API - but I hadn't really thought through the fact that session ID isn't public info in an app. I don't think this is especially useful for a user - even clear has pretty limited utility - so it should likely just be made private.

I'll pull this into a helpful function at the module level instead, so that it's not an exported symbol by default.

API lifted to a non-public namespace.

This actually matches the scope of this function better, since it's operating on a module-level object ...

FWIW, if we wanted to give users the ability to clear caches, it could just be done through the AppSession method.

Clear backing store for session cache for the data cache.

sfc-gh-jkinkead

PTAL!

sfc-gh-jkinkead · 2026-01-05T19:23:05Z

lib/streamlit/runtime/caching/cache_data_api.py

+
+        if session_caches is not None:
+            for cache in session_caches.values():
+                cache.clear()


Good catch! Updating.

sfc-gh-jkinkead · 2026-01-05T20:03:36Z

lib/streamlit/runtime/caching/cache_resource_api.py

        """Clear all cache_resource caches."""
        _resource_caches.clear_all()

+    def clear_session(self, session_id: str) -> None:


API lifted to a non-public namespace.

This actually matches the scope of this function better, since it's operating on a module-level object ...

lukasmasuch

LGTM 👍

## Describe your changes Add session-level scoping to `st.cache_data` and `st.cache_resource`. Add the "scope" parameter through the caching stack. Create session-scoped caches as appropriate. Add methods to clear caches by session ID. Call those methods when sessions are disconnected. **Note**: This clears caches on disconnect _and_ shutdown. In the current websocket session manager, sessions only appear to be shut down when the backend process terminates - and so disconnection is the only hook that's actually invoked in the typical session lifecycle. I'm not sure if this is a bug or a design choice. Either way, like the docs note, this means that session caches might populate multiple times for a single user session in some edge-cases. I think this is fine. ## GitHub Issue Link (if applicable) - Fix for streamlit#8545. An `on_release` hook for a session-scoped resource can be used for disconnect hooks. - Implements one of the suggested fixes for streamlit#6703. ## Testing Plan See unit tests. --- **Contribution License Agreement** By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license.

## Describe your changes Fixes an issue with Snowflake connections not getting re-initialized after having been closed. <details> <summary>Claude issue analysis</summary> # Issue Report: Snowflake Connection "Connection is closed" Error ## Summary After recent PRs adding `on_release` to `st.cache_resource` and session-scoped connection support, users may encounter a `snowflake.connector.errors.DatabaseError: 250002 (08003): Connection is closed` error when using `st.connection("snowflake")` with cached data queries. ## Related PRs - **PR #13439**: Add `on_release` to `st.cache_resource` - **PR #13482**: Add session scoping to caches - **PR #13538**: Add `SnowflakeCallersRightsConnection` - **PR #13506**: Add session-scoped connection support ## Error Details ``` Traceback (most recent call last): File ".../streamlit/runtime/scriptrunner/exec_code.py", line 129, in exec_func_with_error_handling result = func() ... File ".../snowflake/snowpark/_internal/server_connection.py", line 205, in _cursor self._thread_store.cursor = self._conn.cursor() File ".../snowflake/connector/connection.py", line 1270, in cursor Error.errorhandler_wrapper(...) snowflake.connector.errors.DatabaseError: 250002 (08003): Connection is closed ``` ## Root Cause Analysis ### Background The PRs mentioned added important functionality: 1. **PR #13439** added the `on_release` callback to `st.cache_resource`, which is called when cache entries are evicted 2. **PR #13506** modified `connection_factory.py` to use this `on_release` callback to call `connection.close()` when a connection is evicted from the cache This is the relevant code in `connection_factory.py`: ```python def on_release_wrapped(connection: ConnectionClass) -> None: connection.close() __create_connection = cache_resource( max_entries=max_entries, show_spinner="Running `st.connection(...)`.", ttl=ttl, scope=scope, on_release=on_release_wrapped, # Calls close() when evicted )(__create_connection) ``` ### The Bug In `BaseSnowflakeConnection.close()`, after calling `self._raw_instance.close()`, the `_raw_instance` attribute was **NOT** reset to `None`: ```python def close(self) -> None: """Closes the underlying Snowflake connection.""" if self._raw_instance is not None: self._raw_instance.close() # BUG: _raw_instance was NOT set to None! ``` This caused the following issue: 1. When `close()` was called (e.g., via `on_release` when a cache entry is evicted), the underlying connection was closed 2. However, `_raw_instance` still referenced the **closed** connection object 3. The `_instance` property checks `if self._raw_instance is None` to decide whether to create a new connection: ```python @Property def _instance(self) -> RawConnectionT: if self._raw_instance is None: self._raw_instance = self._connect(**self._kwargs) return self._raw_instance ``` 4. Since `_raw_instance` wasn't `None`, subsequent access to `_instance` returned the **CLOSED** connection 5. Any operations on the closed connection failed with "Connection is closed" ### When This Bug Manifests The `on_release` callback (which calls `close()`) is triggered when: - Cache entries expire due to TTL - Cache is full and oldest entries are evicted (`max_entries`) - `st.cache_resource.clear()` is called - For session-scoped caches: when a session disconnects For global-scoped connections like `st.connection("snowflake")`, this typically only happens if: - `st.cache_resource.clear()` is called explicitly - TTL is set and expires - `max_entries` is set and exceeded ### Additional Consideration: Snowpark Sessions When users call `conn.session()`, they get a Snowpark Session that internally references `self._instance`. If the underlying connection is closed: ```python def session(self) -> Session: if running_in_sis(): return get_active_session() return Session.builder.configs({"connection": self._instance}).create() ``` Any Snowpark Sessions created from the connection will also fail because they hold a reference to the now-closed underlying connection object. ## Fix The fix is simple: reset `_raw_instance` to `None` after closing the connection: ```python def close(self) -> None: """Closes the underlying Snowflake connection.""" if self._raw_instance is not None: self._raw_instance.close() self._raw_instance = None # Added this line ``` This ensures that after `close()` is called, the next access to `_instance` will create a new connection instead of returning the closed one. ## Files Changed 1. **`lib/streamlit/connections/snowflake_connection.py`** - Fixed `close()` method to reset `_raw_instance = None` after closing 2. **`lib/tests/streamlit/connections/snowflake_connection_test.py`** - Added `TestSnowflakeConnectionClose` test class with: - `test_close_resets_raw_instance`: Verifies that `close()` closes the connection AND resets `_raw_instance` - `test_close_is_noop_when_not_connected`: Verifies that `close()` doesn't fail when `_raw_instance` is already `None` ## Testing ```bash PYTHONPATH=lib pytest lib/tests/streamlit/connections/snowflake_connection_test.py::TestSnowflakeConnectionClose -v ``` Output: ``` lib/tests/streamlit/connections/snowflake_connection_test.py::TestSnowflakeConnectionClose::test_close_resets_raw_instance PASSED lib/tests/streamlit/connections/snowflake_connection_test.py::TestSnowflakeConnectionClose::test_close_is_noop_when_not_connected PASSED ``` ## Recommendations for Users Until this fix is released, users experiencing this issue can: 1. **Avoid storing Snowpark Sessions long-term**: Instead of caching Snowpark Sessions, create them fresh when needed 2. **Check if using `st.cache_resource.clear()`**: If calling this anywhere in the app, it will close all cached connections 3. **Consider connection TTL settings**: If TTL is set on the connection, it may expire and close ## Impact - **Affected**: Users of `st.connection("snowflake")` and `st.connection("snowflake-callers-rights")` who experience cache eviction scenarios - **Severity**: Medium - The bug causes operations to fail with a confusing error message, but the workaround (restarting the app or avoiding cache clears) is available - **Scope**: Only affects `SnowflakeConnection` and its subclasses; other connection types (`SQLConnection`, `SnowparkConnection`) inherit the no-op `close()` from `BaseConnection` and are not affected </details> ## GitHub Issue Link (if applicable) ## Testing Plan - Added unit test. --- **Contribution License Agreement** By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license.

Add session scoping to caches.

159e283

Add a "scope" parameter through the caching stack. Create session-scoped caches as appropriate. Add methods to clear caches by session ID. Call those methods when sessions expire.

Copilot AI review requested due to automatic review settings December 29, 2025 20:50

Copilot started reviewing on behalf of sfc-gh-jkinkead December 29, 2025 20:51 View session

graphite-app bot reviewed Dec 29, 2025

View reviewed changes

lib/streamlit/runtime/caching/cache_data_api.py Outdated Show resolved Hide resolved

Copilot AI reviewed Dec 29, 2025

View reviewed changes

Fix typos; clarify comment.

19f5e15

graphite-app bot reviewed Dec 29, 2025

View reviewed changes

lib/tests/streamlit/runtime/caching/cache_utils_test.py Show resolved Hide resolved

Add pydoc to tests.

3816e0f

lukasmasuch added ai-review If applied to PR or issue will run AI review workflow security-assessment-completed Security assessment has been completed for PR change:feature PR contains new feature or enhancement implementation impact:users PR changes affect end users labels Dec 31, 2025

github-actions bot removed the ai-review If applied to PR or issue will run AI review workflow label Dec 31, 2025

cursor bot reviewed Dec 31, 2025

View reviewed changes

lukasmasuch reviewed Dec 31, 2025

View reviewed changes

Lift clear_session into module-level function.

204ab76

Clear backing store for session cache for the data cache.

sfc-gh-jkinkead commented Jan 5, 2026

View reviewed changes

sfc-gh-jkinkead requested a review from lukasmasuch January 5, 2026 20:05

sfc-gh-jkinkead added 2 commits January 5, 2026 12:43

Merge branch 'develop' into jkinkead-session-scoped-caches

07013b4

Update copyright in new file.

d8dd712

lukasmasuch approved these changes Jan 5, 2026

View reviewed changes

sfc-gh-jkinkead merged commit 9b42523 into develop Jan 5, 2026
42 checks passed

sfc-gh-jkinkead deleted the jkinkead-session-scoped-caches branch January 5, 2026 21:30

This was referenced Jan 6, 2026

Shut down the session after closing the browser tab by the user #8545

Open

st.cache_resource alternative for non-threadsafe object (that can't be serialized / used with st.cache_data) #6703

Closed

lukasmasuch mentioned this pull request Jan 21, 2026

Fix a regression with snowflake connection being closed #13665

Merged

Add session scoping to caches. #13482

Add session scoping to caches. #13482

Uh oh!

Conversation

sfc-gh-jkinkead commented Dec 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Describe your changes

GitHub Issue Link (if applicable)

Testing Plan

Uh oh!

snyk-io bot commented Dec 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ Snyk checks have passed. No issues have been found so far.

Uh oh!

github-actions bot commented Dec 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ PR preview is ready!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Dec 31, 2025

Summary

Code Quality

Test Coverage

Backwards Compatibility

Security & Risk

Recommendations

Verdict

Uh oh!

cursor bot Dec 31, 2025

Choose a reason for hiding this comment

Storage not closed when session caches are cleared

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-jkinkead left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lukasmasuch left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

sfc-gh-jkinkead commented Dec 29, 2025 •

edited

Loading

snyk-io bot commented Dec 29, 2025 •

edited

Loading

github-actions bot commented Dec 29, 2025 •

edited

Loading