Skip to content

update the task information to the first step - #31

Merged
BinHuangPJLAB merged 2 commits into
AI45Lab:v2from
BinHuangPJLAB:gateway-task-patch
Jul 30, 2026
Merged

update the task information to the first step#31
BinHuangPJLAB merged 2 commits into
AI45Lab:v2from
BinHuangPJLAB:gateway-task-patch

Conversation

@BinHuangPJLAB

@BinHuangPJLAB BinHuangPJLAB commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

Summary by CodeRabbit

  • New Features

    • Session recording can now persist and return associated dataset information, attached to the first step of each session.
    • Added the ability to list persisted session steps in trajectory order, including deterministic retrieval for cloud-backed storage.
  • Bug Fixes

    • Improved session environment binding by reusing already-resolved session context when binding details are present.
    • Added safer dataset-attachment coordination during batch step recording, including cleanup on cancellations and errors.

@coderabbitai

coderabbitai Bot commented Jul 30, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 098b7cad-a0c6-4d1f-90f3-0b4b17243947

📥 Commits

Reviewing files that changed from the base of the PR and between e254130 and a944c41.

📒 Files selected for processing (1)
  • core/data_manager/strategy/cloud_strategy_impl.py

📝 Walkthrough

Walkthrough

record_step now accepts and persists optional dataset metadata across gateway, cloud, and SQLite storage. DataManager and storage strategies also expose session-step listing, with cloud results parsed and returned in deterministic trajectory order.

Changes

Dataset Metadata and Session Step Retrieval

Layer / File(s) Summary
Record-step and retrieval contracts
core/data_manager/manager.py, core/data_manager/strategy/base_strategy.py
Recording accepts an optional dataset value, and DataManager delegates session-step listing to the storage strategy.
Gateway dataset extraction and first-step attachment
gateway/storage.py
Environment parsing carries dataset metadata, and batch inference persistence attaches it once to each session’s first step while clearing pending claims on cancellation or failure.
Strategy persistence and session listing
core/data_manager/strategy/cloud_strategy_impl.py, core/data_manager/strategy/sqlite_strategy_impl.py
Cloud records serialize dataset metadata and list parsed, deterministically ordered session steps; SQLite stores dataset in serialized environment state, and cloud queries use the RL dataset type.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant GatewayStorage
  participant DataManager
  participant StorageStrategy
  participant PersistedRecord
  GatewayStorage->>DataManager: record first step with dataset
  DataManager->>StorageStrategy: forward dataset
  StorageStrategy->>PersistedRecord: persist dataset metadata
  DataManager->>StorageStrategy: list_session_steps(session_id)
  StorageStrategy->>DataManager: return ordered session steps
Loading

Possibly related PRs

  • AI45Lab/SAfactory#18: Updates the same DataManager and storage interfaces for dataset-aware step recording and session-step listing.
  • AI45Lab/SAfactory#28: Updates the cloud dataset type constant and related cloud strategy handling.
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title matches the main change: moving task/dataset information onto the first step of a session.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)
core/data_manager/strategy/base_strategy.py (1)

168-176: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Implement session-step listing for SQLite.

SqliteStrategy inherits this empty default, so the new DataManager.list_session_steps() API always returns [] for the default SQLite backend. Add a SQLite override that flushes pending writes and orders SessionStep rows by step_id, id.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@core/data_manager/strategy/base_strategy.py` around lines 168 - 176,
Implement list_session_steps in SqliteStrategy to flush pending writes, query
the requested session’s SessionStep rows, and order them by step_id followed by
id. Return the persisted rows in the same List[Dict[str, Any]] shape expected by
DataManager.list_session_steps, while leaving the base strategy default
unchanged.
core/data_manager/strategy/cloud_strategy_impl.py (1)

793-799: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Expose persisted dataset metadata in retrieval results.

dataset is stored in meta_json at Line 674, but this method pops that payload and never maps meta["dataset"] onto the returned row. Cloud callers of list_session_steps() therefore cannot retrieve the dataset metadata this PR records.

Proposed fix
             row["env_state"] = _json_object(meta.get("env_state"))
             row["group_id"] = meta.get("group_id")
+            row["dataset"] = meta.get("dataset")
             if row.get("is_trainable") is None:
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@core/data_manager/strategy/cloud_strategy_impl.py` around lines 793 - 799,
Update the row-mapping logic in list_session_steps() to expose the persisted
dataset metadata by assigning row["dataset"] from meta.get("dataset") after
parsing meta_json. Preserve the existing metadata mappings and fallback behavior
for all other fields.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@gateway/storage.py`:
- Around line 44-45: The completed-session tracking in _dataset_written_sessions
grows indefinitely because entries are not removed when session state is
evicted. Update the successful close and _evict_expired() handling to remove
non-pending, completed session IDs from _dataset_written_sessions once their
corresponding cache entry is released, while preserving IDs still marked
pending.

---

Outside diff comments:
In `@core/data_manager/strategy/base_strategy.py`:
- Around line 168-176: Implement list_session_steps in SqliteStrategy to flush
pending writes, query the requested session’s SessionStep rows, and order them
by step_id followed by id. Return the persisted rows in the same List[Dict[str,
Any]] shape expected by DataManager.list_session_steps, while leaving the base
strategy default unchanged.

In `@core/data_manager/strategy/cloud_strategy_impl.py`:
- Around line 793-799: Update the row-mapping logic in list_session_steps() to
expose the persisted dataset metadata by assigning row["dataset"] from
meta.get("dataset") after parsing meta_json. Preserve the existing metadata
mappings and fallback behavior for all other fields.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: e4d59552-373f-4ba4-a6a4-2da8f11a9ade

📥 Commits

Reviewing files that changed from the base of the PR and between 3c0f634 and e254130.

📒 Files selected for processing (5)
  • core/data_manager/manager.py
  • core/data_manager/strategy/base_strategy.py
  • core/data_manager/strategy/cloud_strategy_impl.py
  • core/data_manager/strategy/sqlite_strategy_impl.py
  • gateway/storage.py

Comment thread gateway/storage.py
Comment on lines +44 to +45
self._dataset_pending_sessions: set[str] = set()
self._dataset_written_sessions: set[str] = set()

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Release completed dataset claims.

_dataset_written_sessions is never pruned, including when _evict_expired() removes the corresponding session cache entry. A long-running gateway accumulates one ID per completed session indefinitely. Clear completed, non-pending session IDs after a successful close (or use bounded TTL-based state).

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@gateway/storage.py` around lines 44 - 45, The completed-session tracking in
_dataset_written_sessions grows indefinitely because entries are not removed
when session state is evicted. Update the successful close and _evict_expired()
handling to remove non-pending, completed session IDs from
_dataset_written_sessions once their corresponding cache entry is released,
while preserving IDs still marked pending.

@BinHuangPJLAB
BinHuangPJLAB merged commit e65011c into AI45Lab:v2 Jul 30, 2026
1 check passed
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