vikarti.anatra/Skyvern
1
0
Fork 0
mirror of https://github.com/Skyvern-AI/skyvern.git synced 2026-07-09 16:09:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

Fix post-navigation recording capture and interpretation resume (#6568)

Browse source
This commit is contained in:
Cindy Li 2026-06-14 23:50:12 -04:00 • committed by GitHub
parent cdacadb95d
commit 006bce80aa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 292 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

81
skyvern/forge/sdk/routes/streaming/channels/exfiltration.py
View file

@ -68,7 +68,8 @@ class ExfiltrationChannel(CdpChannel):
)
_binding_registered_pages: t.ClassVar[weakref.WeakSet[Page]] = weakref.WeakSet()
_adorn_init_script_pages: t.ClassVar[weakref.WeakSet[Page]] = weakref.WeakSet()
_rearm_in_flight_pages: t.ClassVar[weakref.WeakSet[Page]] = weakref.WeakSet()
_rearm_in_flight_pages: t.ClassVar[weakref.WeakKeyDictionary[Page, bool]] = weakref.WeakKeyDictionary()
_rearm_pending_full_nav_pages: t.ClassVar[weakref.WeakSet[Page]] = weakref.WeakSet()
def __init__(self, *, on_event: OnExfiltrationEvent, vnc_channel: VncChannel) -> None:
self.cdp_session: CDPSession | None = None
@ -339,21 +340,53 @@ class ExfiltrationChannel(CdpChannel):
self._emit_events(messages)
if event_name == "nav:frame_navigated":
if event_name in ("nav:frame_navigated", "nav:navigated_within_document"):
page = self.page
if not page or page.url.startswith("devtools:"):
return
if page and not page.url.startswith("devtools:"):
self._schedule_page_rearm(
page,
event_name=event_name,
wait_for_load=event_name == "nav:frame_navigated",
)
if page in self._rearm_in_flight_pages:
return
self._track_event_task(self._rearm_page_after_navigation(page, event_name=event_name))
async def _rearm_page_after_navigation(self, page: Page, *, event_name: str) -> None:
def _schedule_page_rearm(self, page: Page, *, event_name: str, wait_for_load: bool) -> None:
if page.url.startswith("devtools:"):
return
self._rearm_in_flight_pages.add(page)
in_flight_wait_for_load = self._rearm_in_flight_pages.get(page)
if in_flight_wait_for_load is not None:
if wait_for_load and not in_flight_wait_for_load:
self._rearm_pending_full_nav_pages.add(page)
return
self._track_event_task(
self._rearm_page_after_navigation(page, event_name=event_name, wait_for_load=wait_for_load)
)
async def rearm_all_pages(self) -> None:
browser_context = self.browser_context
if not browser_context:
page = self.page
if page and not page.url.startswith("devtools:"):
self._schedule_page_rearm(page, event_name="recording:rearm", wait_for_load=False)
return
for page in list(browser_context.pages):
if page.url.startswith("devtools:"):
continue
self._schedule_page_rearm(page, event_name="recording:rearm", wait_for_load=False)
async def _rearm_page_after_navigation(
self,
page: Page,
*,
event_name: str,
wait_for_load: bool = True,
) -> None:
if page.url.startswith("devtools:"):
return
self._rearm_in_flight_pages[page] = wait_for_load
try:
LOG.info(
"re-applying exfiltration and adornment after navigation",
@ -362,16 +395,17 @@ class ExfiltrationChannel(CdpChannel):
url=page.url,
)
try:
await page.wait_for_load_state("domcontentloaded", timeout=10_000)
except Exception:
LOG.warning(
"navigation re-arm timed out waiting for domcontentloaded",
class_name=self.class_name,
event_name=event_name,
url=page.url,
exc_info=True,
)
if wait_for_load:
try:
await page.wait_for_load_state("domcontentloaded", timeout=10_000)
except Exception:
LOG.warning(
"navigation re-arm timed out waiting for domcontentloaded",
class_name=self.class_name,
event_name=event_name,
url=page.url,
exc_info=True,
)
try:
await self._ensure_binding(page)
@ -386,7 +420,10 @@ class ExfiltrationChannel(CdpChannel):
exc_info=True,
)
finally:
self._rearm_in_flight_pages.discard(page)
self._rearm_in_flight_pages.pop(page, None)
if page in self._rearm_pending_full_nav_pages:
self._rearm_pending_full_nav_pages.discard(page)
self._schedule_page_rearm(page, event_name="nav:frame_navigated", wait_for_load=True)
async def adorn(self, page: Page) -> t.Self:
"""Add a mouse-following follower to the page."""

63
skyvern/forge/sdk/routes/streaming/channels/message.py
View file

@ -69,6 +69,7 @@ class MessageKind(enum.StrEnum):
RELOAD = "reload"
RECORDING_CAPTURE_PAUSE = "recording-capture-pause"
RECORDING_CAPTURE_RESUME = "recording-capture-resume"
RECORDING_REARM_CAPTURE = "recording-rearm-capture"
RECORDING_INTERPRETATION_UPDATE = "recording-interpretation-update"
SCREENSHOT = "screenshot"
TAKE_CONTROL = "take-control"
@ -114,6 +115,7 @@ MessageKinds = t.Literal[
MessageKind.RELOAD,
MessageKind.RECORDING_CAPTURE_PAUSE,
MessageKind.RECORDING_CAPTURE_RESUME,
MessageKind.RECORDING_REARM_CAPTURE,
MessageKind.RECORDING_INTERPRETATION_UPDATE,
MessageKind.SCREENSHOT,
MessageKind.TAKE_CONTROL,
@ -148,6 +150,11 @@ class MessageInRecordingCaptureResume(Message):
kind: t.Literal[MessageKind.RECORDING_CAPTURE_RESUME] = MessageKind.RECORDING_CAPTURE_RESUME
@dataclasses.dataclass
class MessageInRecordingRearmCapture(Message):
kind: t.Literal[MessageKind.RECORDING_REARM_CAPTURE] = MessageKind.RECORDING_REARM_CAPTURE
@dataclasses.dataclass
class MessageInTakeControl(Message):
kind: t.Literal[MessageKind.TAKE_CONTROL] = MessageKind.TAKE_CONTROL
@ -287,6 +294,7 @@ MessageIn = (
| MessageInNavigate
| MessageInRecordingCapturePause
| MessageInRecordingCaptureResume
| MessageInRecordingRearmCapture
| MessageInReload
| MessageInTakeControl
| MessageInTakeScreenshot
@ -338,6 +346,8 @@ def reify_channel_message(data: dict) -> ChannelMessage:
return MessageInRecordingCapturePause()
case MessageKind.RECORDING_CAPTURE_RESUME:
return MessageInRecordingCaptureResume()
case MessageKind.RECORDING_REARM_CAPTURE:
return MessageInRecordingRearmCapture()
case MessageKind.GET_BROWSER_URL:
return MessageInGetBrowserUrl()
case MessageKind.GO_BACK:
@ -627,12 +637,6 @@ async def loop_stream_messages(message_channel: MessageChannel) -> None:
await send_error(message.kind, "Failed to paste into browser.")
case MessageKind.BEGIN_EXFILTRATION:
if exfiltration_channel is not None:
LOG.error(
"MessageChannel: cannot begin exfiltration: already active.", message_channel=message_channel
)
return
vnc_channel = get_vnc_channel(message_channel.client_id)
if not vnc_channel:
@ -649,20 +653,26 @@ async def loop_stream_messages(message_channel: MessageChannel) -> None:
else None
)
if browser_session_id and message.workflow_permanent_id and message.live_interpretation_enabled:
def on_interpretation_update(update: RecordingInterpretationUpdate) -> None:
message_channel.send_nowait(
messages=[
MessageOutRecordingInterpretationUpdate(
interpretation_session_id=update.interpretation_session_id,
session_revision=update.session_revision,
steps=update.steps,
pending=update.pending,
finalized=update.finalized,
)
]
)
def on_interpretation_update(update: RecordingInterpretationUpdate) -> None:
message_channel.send_nowait(
messages=[
MessageOutRecordingInterpretationUpdate(
interpretation_session_id=update.interpretation_session_id,
session_revision=update.session_revision,
steps=update.steps,
pending=update.pending,
finalized=update.finalized,
)
]
)
def ensure_live_interpretation() -> None:
nonlocal live_interpretation_browser_session_id
if not (
browser_session_id and message.workflow_permanent_id and message.live_interpretation_enabled
):
return
interpretation_registry.start_session(
browser_session_id=browser_session_id,
@ -672,6 +682,13 @@ async def loop_stream_messages(message_channel: MessageChannel) -> None:
)
live_interpretation_browser_session_id = browser_session_id
if exfiltration_channel is not None:
ensure_live_interpretation()
await exfiltration_channel.rearm_all_pages()
return
ensure_live_interpretation()
def on_event(events: list[ExfiltratedEvent]) -> None:
for event in events:
message_out_exfiltrated_event = MessageOutExfiltratedEvent(
@ -803,6 +820,10 @@ async def loop_stream_messages(message_channel: MessageChannel) -> None:
if live_interpretation_browser_session_id:
interpretation_registry.resume_capture(live_interpretation_browser_session_id)
case MessageKind.RECORDING_REARM_CAPTURE:
if exfiltration_channel is not None:
await exfiltration_channel.rearm_all_pages()
case MessageKind.ERROR:
await send(message)
@ -959,8 +980,8 @@ async def loop_stream_messages(message_channel: MessageChannel) -> None:
LOG.debug(f"{class_name} Closing the message channel stream.", **message_channel.identity)
if exfiltration_channel is not None:
await exfiltration_channel.stop()
if live_interpretation_browser_session_id:
await interpretation_registry.stop_session(live_interpretation_browser_session_id)
# Live interpretation is torn down only on explicit end-exfiltration so a
# disconnect→reconnect can resume the same in-flight recording session.
await message_channel.close(reason="loop-channel-closed")

21
skyvern/services/browser_recording/interpretation.py
View file

@ -483,6 +483,27 @@ class RecordingInterpretationSession:
suffix += 1
return f"{label}_{suffix}"
def emit_snapshot(self) -> None:
if self.session_revision == 0:
return
update = RecordingInterpretationUpdate(
interpretation_session_id=self.interpretation_session_id,
session_revision=self.session_revision,
steps=self.steps,
pending=self.pending,
finalized=self.finalized,
)
try:
self.on_update(update)
except Exception:
LOG.exception(
"Failed to emit recording interpretation snapshot",
browser_session_id=self.browser_session_id,
organization_id=self.organization_id,
)
def _emit_update(self) -> None:
self.session_revision += 1
update = RecordingInterpretationUpdate(

14
skyvern/services/browser_recording/session_registry.py
View file

@ -28,8 +28,20 @@ class RecordingInterpretationSessionRegistry:
workflow_permanent_id: str,
on_update: OnRecordingInterpretationUpdate,
) -> None:
self.discard_session(browser_session_id)
self._prune_expired_sessions()
existing = self._sessions.get(browser_session_id)
if (
existing is not None
and existing.workflow_permanent_id == workflow_permanent_id
and existing.organization_id == organization_id
and not existing.finalized
):
existing.on_update = on_update
self._last_seen[browser_session_id] = time.monotonic()
existing.emit_snapshot()
return
self.discard_session(browser_session_id)
self._sessions[browser_session_id] = RecordingInterpretationSession(
browser_session_id=browser_session_id,
organization_id=organization_id,

51
tests/unit/forge/sdk/routes/streaming/test_exfiltration_channel.py
View file

@ -69,12 +69,16 @@ def restore_exfiltration_channel_class_state() -> t.Iterator[None]:
active_binding_channels = weakref.WeakKeyDictionary(ExfiltrationChannel._active_binding_channels)
binding_registered_pages = weakref.WeakSet(ExfiltrationChannel._binding_registered_pages)
adorn_init_script_pages = weakref.WeakSet(ExfiltrationChannel._adorn_init_script_pages)
rearm_in_flight_pages = weakref.WeakKeyDictionary(ExfiltrationChannel._rearm_in_flight_pages)
rearm_pending_full_nav_pages = weakref.WeakSet(ExfiltrationChannel._rearm_pending_full_nav_pages)
yield
ExfiltrationChannel._active_binding_channels = active_binding_channels
ExfiltrationChannel._binding_registered_pages = binding_registered_pages
ExfiltrationChannel._adorn_init_script_pages = adorn_init_script_pages
ExfiltrationChannel._rearm_in_flight_pages = rearm_in_flight_pages
ExfiltrationChannel._rearm_pending_full_nav_pages = rearm_pending_full_nav_pages
class TestExfiltrationChannelEvents:
@ -314,10 +318,12 @@ class TestNavigationReExfiltration:
channel.adorn.assert_awaited_once_with(page)
@pytest.mark.asyncio
async def test_navigated_within_document_does_not_rearm(self) -> None:
async def test_navigated_within_document_rearms_without_load_wait(self) -> None:
channel, on_event = _make_channel()
page = _make_page("https://example.com/app#section")
page.wait_for_load_state = AsyncMock()
channel.page = page
channel._ensure_binding = AsyncMock()
channel.exfiltrate = AsyncMock(return_value=channel)
channel.adorn = AsyncMock(return_value=channel)
@ -327,8 +333,47 @@ class TestNavigationReExfiltration:
if channel._pending_event_tasks:
await asyncio.gather(*channel._pending_event_tasks)
channel.exfiltrate.assert_not_called()
channel.adorn.assert_not_called()
page.wait_for_load_state.assert_not_awaited()
channel._ensure_binding.assert_awaited_once_with(page)
channel.exfiltrate.assert_awaited_once_with(page)
channel.adorn.assert_awaited_once_with(page)
@pytest.mark.asyncio
async def test_frame_navigated_queues_full_rearm_when_same_document_rearm_in_flight(self) -> None:
channel, on_event = _make_channel()
page = _make_page("https://example.com/app")
channel.page = page
channel._rearm_in_flight_pages[page] = False
channel._handle_cdp_event("nav:frame_navigated", {"frame": {"url": "https://example.com/next"}})
on_event.assert_called_once()
assert page in channel._rearm_pending_full_nav_pages
assert not channel._pending_event_tasks
@pytest.mark.asyncio
async def test_pending_full_nav_rearm_runs_after_same_document_rearm_finishes(self) -> None:
channel, _on_event = _make_channel()
page = _make_page("https://example.com/next")
page.wait_for_load_state = AsyncMock()
channel._ensure_binding = AsyncMock()
channel.exfiltrate = AsyncMock(return_value=channel)
channel.adorn = AsyncMock(return_value=channel)
channel._rearm_pending_full_nav_pages.add(page)
await channel._rearm_page_after_navigation(
page,
event_name="nav:navigated_within_document",
wait_for_load=False,
)
if channel._pending_event_tasks:
await asyncio.gather(*channel._pending_event_tasks)
page.wait_for_load_state.assert_awaited_once_with("domcontentloaded", timeout=10_000)
channel._ensure_binding.assert_awaited()
assert channel._ensure_binding.await_count == 2
assert channel.exfiltrate.await_count == 2
assert channel.adorn.await_count == 2
def test_non_navigation_cdp_event_does_not_rearm(self) -> None:
channel, on_event = _make_channel()

7
tests/unit/forge/sdk/routes/streaming/test_message_reify.py
View file

@ -22,6 +22,7 @@ from skyvern.forge.sdk.routes.streaming.channels.message import (
MessageInGoBack,
MessageInGoForward,
MessageInNavigate,
MessageInRecordingRearmCapture,
MessageInReload,
MessageInTakeScreenshot,
MessageKind,
@ -103,6 +104,12 @@ class TestReifyChannelMessage:
assert msg.workflow_permanent_id == "wpid_123"
assert msg.live_interpretation_enabled is True
def test_recording_rearm_capture(self) -> None:
assert isinstance(
reify_channel_message({"kind": "recording-rearm-capture"}),
MessageInRecordingRearmCapture,
)
class TestMessageSerialization:
def test_recording_interpretation_update_serializes_pydantic_steps(self) -> None:

102
tests/unit/test_browser_recording_interpretation.py
View file

@ -350,3 +350,105 @@ async def test_recording_interpretation_session_advances_past_unhandled_actions(
assert session.emitted_action_count == 2
assert len(session.steps) == 1
def test_emit_snapshot_replays_current_revision_without_incrementing() -> None:
updates: list[int] = []
session = RecordingInterpretationSession(
browser_session_id=PBS_ID,
organization_id=ORG_ID,
workflow_permanent_id=WP_ID,
on_update=lambda update: updates.append(update.session_revision),
)
session.session_revision = 2
session.steps = [
RecordingDraftStep(
step_id="step-1",
action_kind=ActionKind.CLICK,
block_type="action",
label="click_submit",
title="Click submit",
navigation_goal="Click submit",
)
]
session.emit_snapshot()
assert updates == [2]
assert session.session_revision == 2
def test_start_session_resumes_existing_interpretation_session() -> None:
from skyvern.services.browser_recording.session_registry import RecordingInterpretationSessionRegistry
registry = RecordingInterpretationSessionRegistry()
first_updates: list[int] = []
second_updates: list[int] = []
registry.start_session(
browser_session_id=PBS_ID,
organization_id=ORG_ID,
workflow_permanent_id=WP_ID,
on_update=lambda update: first_updates.append(update.session_revision),
)
session = registry._sessions[PBS_ID]
session.session_revision = 3
session.steps = [
RecordingDraftStep(
step_id="step-1",
action_kind=ActionKind.CLICK,
block_type="action",
label="click_submit",
title="Click submit",
navigation_goal="Click submit",
)
]
registry.start_session(
browser_session_id=PBS_ID,
organization_id=ORG_ID,
workflow_permanent_id=WP_ID,
on_update=lambda update: second_updates.append(update.session_revision),
)
assert registry._sessions[PBS_ID] is session
assert first_updates == []
assert second_updates == [3]
def test_start_session_resumes_after_websocket_disconnect_without_stop() -> None:
from skyvern.services.browser_recording.session_registry import RecordingInterpretationSessionRegistry
registry = RecordingInterpretationSessionRegistry()
reconnect_updates: list[int] = []
registry.start_session(
browser_session_id=PBS_ID,
organization_id=ORG_ID,
workflow_permanent_id=WP_ID,
on_update=lambda _: None,
)
session = registry._sessions[PBS_ID]
session.session_revision = 4
session.steps = [
RecordingDraftStep(
step_id="step-1",
action_kind=ActionKind.CLICK,
block_type="action",
label="click_submit",
title="Click submit",
navigation_goal="Click submit",
)
]
# WebSocket loop teardown no longer calls stop_session; only end-exfiltration does.
registry.start_session(
browser_session_id=PBS_ID,
organization_id=ORG_ID,
workflow_permanent_id=WP_ID,
on_update=lambda update: reconnect_updates.append(update.session_revision),
)
assert registry._sessions[PBS_ID] is session
assert reconnect_updates == [4]