383 lines
13 KiB
Python
383 lines
13 KiB
Python
"""Tests for prompt-scoped remote session state helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
|
|
def test_remote_session_store_round_trips_value_refs(session_state_module: Any) -> None:
|
|
"""Stored remote-only outputs should round-trip through opaque session refs."""
|
|
handle = session_state_module.RemoteSessionHandle(
|
|
session_id="session-1",
|
|
prompt_id="prompt-1",
|
|
owner_component_id="component-1",
|
|
)
|
|
store = session_state_module.InMemoryRemoteSessionStore()
|
|
|
|
ref = store.put_output(
|
|
handle,
|
|
node_id="node-7",
|
|
output_index=0,
|
|
value={"remote": "value"},
|
|
)
|
|
|
|
assert store.get_output(ref) == {"remote": "value"}
|
|
assert store.resolve_value(ref.to_payload()) == {"remote": "value"}
|
|
|
|
|
|
def test_remote_session_store_clear_session_invalidates_refs(session_state_module: Any) -> None:
|
|
"""Clearing a prompt-scoped session should invalidate every ref stored inside it."""
|
|
handle = session_state_module.RemoteSessionHandle(
|
|
session_id="session-2",
|
|
prompt_id="prompt-2",
|
|
owner_component_id="component-2",
|
|
)
|
|
store = session_state_module.InMemoryRemoteSessionStore()
|
|
ref = store.put_output(
|
|
handle,
|
|
node_id="node-9",
|
|
output_index=1,
|
|
value="secret",
|
|
)
|
|
|
|
store.clear_session(handle)
|
|
|
|
with pytest.raises(session_state_module.RemoteSessionStateError):
|
|
store.get_output(ref)
|
|
|
|
|
|
def test_remote_session_store_logs_reuse_resolution_and_cleanup(
|
|
session_state_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Session observability logs should cover reuse, ref resolution, and cleanup."""
|
|
handle = session_state_module.RemoteSessionHandle(
|
|
session_id="session-3",
|
|
prompt_id="prompt-3",
|
|
owner_component_id="component-3",
|
|
)
|
|
store = session_state_module.InMemoryRemoteSessionStore()
|
|
observed_messages: list[str] = []
|
|
|
|
def capture_log(message: str, *args: Any, **kwargs: Any) -> None:
|
|
"""Record one formatted session log line."""
|
|
del kwargs
|
|
observed_messages.append(message % args if args else message)
|
|
|
|
monkeypatch.setattr(session_state_module.logger, "info", capture_log)
|
|
|
|
store.ensure_session(handle)
|
|
store.ensure_session(handle)
|
|
ref = store.put_output(
|
|
handle,
|
|
node_id="node-11",
|
|
output_index=0,
|
|
value="payload",
|
|
)
|
|
assert store.resolve_value(ref.to_payload()) == "payload"
|
|
store.clear_session(handle)
|
|
|
|
assert any("Created remote session session_id=session-3" in message for message in observed_messages)
|
|
assert any("Reusing remote session session_id=session-3" in message for message in observed_messages)
|
|
assert any(
|
|
"Stored remote session value session_id=session-3 node_id=node-11 output_index=0" in message
|
|
for message in observed_messages
|
|
)
|
|
assert any(
|
|
"Resolving remote session ref session_id=session-3 node_id=node-11 output_index=0" in message
|
|
for message in observed_messages
|
|
)
|
|
assert any(
|
|
"Resolved remote session value session_id=session-3 node_id=node-11 output_index=0" in message
|
|
for message in observed_messages
|
|
)
|
|
assert any("Cleared remote session session_id=session-3" in message for message in observed_messages)
|
|
|
|
|
|
def test_stable_session_bridge_key_ignores_transient_session_ids(
|
|
session_state_module: Any,
|
|
) -> None:
|
|
"""Durable bridge keys should stay stable across prompt-scoped session id changes."""
|
|
first_key = session_state_module.stable_session_bridge_key(
|
|
producer_payload={"component_id": "1"},
|
|
producer_inputs={
|
|
"model": session_state_module.RemoteSessionBridgeRef(
|
|
bridge_key="RSB_upstream",
|
|
node_id="2",
|
|
output_index=0,
|
|
session_id="session-a",
|
|
).to_payload()
|
|
},
|
|
node_id="3",
|
|
output_index=0,
|
|
)
|
|
second_key = session_state_module.stable_session_bridge_key(
|
|
producer_payload={"component_id": "1"},
|
|
producer_inputs={
|
|
"model": session_state_module.RemoteSessionBridgeRef(
|
|
bridge_key="RSB_upstream",
|
|
node_id="2",
|
|
output_index=0,
|
|
session_id="session-b",
|
|
).to_payload()
|
|
},
|
|
node_id="3",
|
|
output_index=0,
|
|
)
|
|
|
|
assert first_key == second_key
|
|
|
|
|
|
def test_remote_session_bridge_record_round_trips_object_refs(
|
|
session_state_module: Any,
|
|
) -> None:
|
|
"""Bridge metadata should preserve object-backed producer inputs and outputs."""
|
|
producer_inputs_object = session_state_module.DurableObjectRef(
|
|
object_path="bridge_objects/aa/inputs.bin",
|
|
sha256="a" * 64,
|
|
size_bytes=10,
|
|
)
|
|
serialized_output_object = session_state_module.DurableObjectRef(
|
|
object_path="bridge_objects/bb/output.bin",
|
|
sha256="b" * 64,
|
|
size_bytes=20,
|
|
)
|
|
record = session_state_module.RemoteSessionBridgeRecord(
|
|
bridge_key="RSB_object_backed",
|
|
node_id="node-1",
|
|
output_index=0,
|
|
producer_payload={"component_id": "component-1"},
|
|
producer_inputs={},
|
|
recovery_kind=(
|
|
session_state_module.RemoteSessionBridgeRecoveryKind.PRODUCER_REPLAY
|
|
),
|
|
producer_inputs_retained=True,
|
|
producer_inputs_object=producer_inputs_object,
|
|
serialized_output_object=serialized_output_object,
|
|
serialized_output_io_type="LATENT",
|
|
)
|
|
|
|
restored = session_state_module.RemoteSessionBridgeRecord.from_payload(
|
|
record.to_payload()
|
|
)
|
|
|
|
assert restored == record
|
|
|
|
|
|
def test_remote_session_bridge_record_accepts_legacy_recovery_metadata(
|
|
session_state_module: Any,
|
|
) -> None:
|
|
"""Records created before recovery policies should remain readable."""
|
|
record = session_state_module.RemoteSessionBridgeRecord(
|
|
bridge_key="RSB_legacy",
|
|
node_id="node-1",
|
|
output_index=0,
|
|
producer_payload={"component_id": "component-1"},
|
|
producer_inputs={"value": 1},
|
|
)
|
|
payload = record.to_payload()
|
|
payload.pop("recovery_kind")
|
|
payload.pop("producer_inputs_retained")
|
|
|
|
restored = session_state_module.RemoteSessionBridgeRecord.from_payload(payload)
|
|
|
|
assert restored.recovery_kind is None
|
|
assert restored.producer_inputs_retained is None
|
|
assert restored.producer_inputs == {"value": 1}
|
|
|
|
|
|
def test_remote_session_store_resolves_bridge_refs_via_replay_callback(
|
|
session_state_module: Any,
|
|
) -> None:
|
|
"""Durable bridge refs should rehydrate into the target session when the source is gone."""
|
|
store = session_state_module.InMemoryRemoteSessionStore()
|
|
target_handle = session_state_module.RemoteSessionHandle(
|
|
session_id="session-target",
|
|
prompt_id="prompt-1",
|
|
owner_component_id="component-1",
|
|
)
|
|
bridge_ref = session_state_module.RemoteSessionBridgeRef(
|
|
bridge_key="RSB_bridge",
|
|
node_id="node-7",
|
|
output_index=0,
|
|
session_id="session-source",
|
|
)
|
|
observed_bridge_keys: list[str] = []
|
|
|
|
def bridge_resolver(ref: Any) -> str:
|
|
"""Populate the target session and return the replayed value."""
|
|
observed_bridge_keys.append(ref.bridge_key)
|
|
store.put_bridge_output(
|
|
target_handle,
|
|
bridge_key=ref.bridge_key,
|
|
node_id=ref.node_id,
|
|
output_index=ref.output_index,
|
|
value="rehydrated-value",
|
|
)
|
|
return "rehydrated-value"
|
|
|
|
assert (
|
|
store.resolve_value_with_bridges(
|
|
bridge_ref.to_payload(),
|
|
target_session_handle=target_handle,
|
|
bridge_resolver=bridge_resolver,
|
|
)
|
|
== "rehydrated-value"
|
|
)
|
|
assert observed_bridge_keys == ["RSB_bridge"]
|
|
assert (
|
|
store.resolve_value_with_bridges(
|
|
bridge_ref.to_payload(),
|
|
target_session_handle=target_handle,
|
|
bridge_resolver=bridge_resolver,
|
|
)
|
|
== "rehydrated-value"
|
|
)
|
|
assert observed_bridge_keys == ["RSB_bridge"]
|
|
|
|
|
|
def test_remote_session_store_resolves_same_output_bridge_refs_by_bridge_key(
|
|
session_state_module: Any,
|
|
) -> None:
|
|
"""Different bridge keys from the same node output must not collapse to the first value."""
|
|
store = session_state_module.InMemoryRemoteSessionStore()
|
|
target_handle = session_state_module.RemoteSessionHandle(
|
|
session_id="session-target",
|
|
prompt_id="prompt-1",
|
|
owner_component_id="component-1",
|
|
)
|
|
first_ref = session_state_module.RemoteSessionBridgeRef(
|
|
bridge_key="RSB_first",
|
|
node_id="508",
|
|
output_index=0,
|
|
session_id="session-source",
|
|
)
|
|
second_ref = session_state_module.RemoteSessionBridgeRef(
|
|
bridge_key="RSB_second",
|
|
node_id="508",
|
|
output_index=0,
|
|
session_id="session-source",
|
|
)
|
|
values_by_bridge_key = {
|
|
"RSB_first": "conditioning-a",
|
|
"RSB_second": "conditioning-b",
|
|
}
|
|
observed_bridge_keys: list[str] = []
|
|
|
|
def bridge_resolver(ref: Any) -> str:
|
|
"""Populate the target session using the bridge key, as remote rehydration does."""
|
|
observed_bridge_keys.append(ref.bridge_key)
|
|
value = values_by_bridge_key[ref.bridge_key]
|
|
store.put_bridge_output(
|
|
target_handle,
|
|
bridge_key=ref.bridge_key,
|
|
node_id=ref.node_id,
|
|
output_index=ref.output_index,
|
|
value=value,
|
|
)
|
|
return value
|
|
|
|
assert store.resolve_value_with_bridges(
|
|
[first_ref.to_payload(), second_ref.to_payload()],
|
|
target_session_handle=target_handle,
|
|
bridge_resolver=bridge_resolver,
|
|
) == ["conditioning-a", "conditioning-b"]
|
|
assert observed_bridge_keys == ["RSB_first", "RSB_second"]
|
|
|
|
assert store.resolve_value_with_bridges(
|
|
[first_ref.to_payload(), second_ref.to_payload()],
|
|
target_session_handle=target_handle,
|
|
bridge_resolver=bridge_resolver,
|
|
) == ["conditioning-a", "conditioning-b"]
|
|
assert observed_bridge_keys == ["RSB_first", "RSB_second"]
|
|
|
|
|
|
def test_remote_session_store_does_not_resolve_bridge_refs_by_ambiguous_node_slot(
|
|
session_state_module: Any,
|
|
) -> None:
|
|
"""Bridge refs from mapped same-node outputs must not use the overwritten node slot."""
|
|
store = session_state_module.InMemoryRemoteSessionStore()
|
|
source_handle = session_state_module.RemoteSessionHandle(
|
|
session_id="session-source",
|
|
prompt_id="prompt-1",
|
|
owner_component_id="component-1",
|
|
)
|
|
target_handle = session_state_module.RemoteSessionHandle(
|
|
session_id="session-target",
|
|
prompt_id="prompt-1",
|
|
owner_component_id="component-1",
|
|
)
|
|
first_ref = session_state_module.RemoteSessionBridgeRef(
|
|
bridge_key="RSB_first",
|
|
node_id="508",
|
|
output_index=0,
|
|
session_id=source_handle.session_id,
|
|
)
|
|
|
|
store.put_output(
|
|
source_handle,
|
|
node_id=first_ref.node_id,
|
|
output_index=first_ref.output_index,
|
|
value="conditioning-b",
|
|
)
|
|
observed_bridge_keys: list[str] = []
|
|
|
|
def bridge_resolver(ref: Any) -> str:
|
|
"""Return the durable bridge value instead of the ambiguous source node slot."""
|
|
observed_bridge_keys.append(ref.bridge_key)
|
|
return "conditioning-a"
|
|
|
|
assert store.resolve_value_with_bridges(
|
|
first_ref.to_payload(),
|
|
target_session_handle=target_handle,
|
|
bridge_resolver=bridge_resolver,
|
|
) == "conditioning-a"
|
|
assert observed_bridge_keys == ["RSB_first"]
|
|
|
|
|
|
def test_remote_session_store_resolves_nested_bridge_ref_lists(
|
|
session_state_module: Any,
|
|
) -> None:
|
|
"""Bridge refs nested inside list outputs should resolve item by item."""
|
|
store = session_state_module.InMemoryRemoteSessionStore()
|
|
source_handle = session_state_module.RemoteSessionHandle(
|
|
session_id="session-source",
|
|
prompt_id="prompt-1",
|
|
owner_component_id="component-1",
|
|
)
|
|
first_ref = store.put_bridge_output(
|
|
source_handle,
|
|
bridge_key="RSB_a",
|
|
node_id="node-7:item:0",
|
|
output_index=0,
|
|
value="conditioning-a",
|
|
)
|
|
second_ref = store.put_bridge_output(
|
|
source_handle,
|
|
bridge_key="RSB_b",
|
|
node_id="node-7:item:1",
|
|
output_index=0,
|
|
value="conditioning-b",
|
|
)
|
|
nested_refs = [
|
|
session_state_module.RemoteSessionBridgeRef(
|
|
bridge_key="RSB_a",
|
|
node_id=first_ref.node_id,
|
|
output_index=first_ref.output_index,
|
|
session_id=first_ref.session_id,
|
|
).to_payload(),
|
|
session_state_module.RemoteSessionBridgeRef(
|
|
bridge_key="RSB_b",
|
|
node_id=second_ref.node_id,
|
|
output_index=second_ref.output_index,
|
|
session_id=second_ref.session_id,
|
|
).to_payload(),
|
|
]
|
|
|
|
assert store.resolve_value_with_bridges(nested_refs) == [
|
|
"conditioning-a",
|
|
"conditioning-b",
|
|
]
|