Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion api/oss/src/apis/fastapi/sessions/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ class SessionStreamCommandRequestModel(BaseModel):
class SessionHeartbeatRequestModel(BaseModel):
# project scope comes from the caller's credential, never the body
session_id: str
replica_id: str
replica_id: str = Field(min_length=1)
turn_id: Optional[str] = None
is_running: bool = True

Expand Down Expand Up @@ -96,6 +96,13 @@ class SessionStateUpsertRequest(BaseModel):
default=None,
description="Remote sandbox id to record alongside the continuity state.",
)
sandbox_turn_index: Optional[int] = Field(
default=None,
description=(
"the writer's conversation turn index; the pointer write is applied only "
"when it is >= the row's data.latest_turn_index."
),
)


# ---------------------------------------------------------------------------
Expand Down
7 changes: 7 additions & 0 deletions api/oss/src/core/sessions/states/dtos.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,3 +85,10 @@ class SessionStateUpsert(BaseModel):
default=None,
description="Remote sandbox id to record alongside the continuity state.",
)
sandbox_turn_index: Optional[int] = Field(
default=None,
description=(
"the writer's conversation turn index; the pointer write is applied only "
"when it is >= the row's data.latest_turn_index."
),
)
20 changes: 18 additions & 2 deletions api/oss/src/dbs/postgres/sessions/states/dao.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

import uuid_utils.compat as uuid_utils

from sqlalchemy import select
from sqlalchemy import Integer, case, cast, func, select
from sqlalchemy.dialects.postgresql import insert

from oss.src.utils.logging import get_module_logger
Expand Down Expand Up @@ -94,7 +94,23 @@ async def set_session_state(
}
if "data" in upsert.model_fields_set:
update_values["data"] = stmt.excluded.data
if "sandbox_id" in upsert.model_fields_set:
guarded_pointer_write = (

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

🤖 The pointer guard: when the writer supplies sandbox_turn_index, sandbox_id and sandbox_fingerprint update only under this CASE compare-and-set against the row's latest_turn_index; a tokenless write behaves exactly as before, so existing callers are untouched. SQL note: inside ON CONFLICT DO UPDATE, the table-qualified column refers to the EXISTING row and excluded to the proposed one. The insert path (no existing row) applies unconditionally.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

we removed sandbox_fingerprint is this still relevant

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

The fingerprint is gone but this guard stays relevant: it is the turn-index compare-and-set on sandbox_id itself, the multi-runner safety mechanism. It uses only pre-existing columns. Commit 20dbabb removed the fingerprint branch below it.

"sandbox_id" in upsert.model_fields_set
and upsert.sandbox_turn_index is not None
)
if guarded_pointer_write:
pointer_write_allowed = (
func.coalesce(
cast(SessionStateDBE.data["latest_turn_index"].astext, Integer),
-1,
)
<= upsert.sandbox_turn_index
)
update_values["sandbox_id"] = case(
(pointer_write_allowed, stmt.excluded.sandbox_id),
else_=SessionStateDBE.sandbox_id,
)
elif "sandbox_id" in upsert.model_fields_set:
update_values["sandbox_id"] = stmt.excluded.sandbox_id

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Kept: the turn-index compare-and-set on sandbox_id (guarded_pointer_write above). This is the multi-runner safety mechanism and it uses only pre-existing columns. Only the fingerprint branch was removed here, so a stale write still cannot overwrite a newer pointer.


stmt = stmt.on_conflict_do_update(
Expand Down
26 changes: 25 additions & 1 deletion api/oss/src/dbs/postgres/sessions/states/mappings.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,37 @@
from typing import Optional

from pydantic import ValidationError

from oss.src.utils.logging import get_module_logger
from oss.src.core.sessions.states.dtos import (
SessionState,
SessionStateData,
SessionStateFlags,
)
from oss.src.dbs.postgres.sessions.states.dbes import SessionStateDBE

log = get_module_logger(__name__)


def _parse_state_data(dbe: SessionStateDBE) -> Optional[SessionStateData]:
if not dbe.data:
return None

try:
return SessionStateData.model_validate(dbe.data)
except ValidationError as e:
log.error(
"[SESSION_STATES] Corrupt session_states.data; degrading to None",
session_id=dbe.session_id,
project_id=dbe.project_id,
raw_data=dbe.data,
error=str(e),
)
return None


def dbe_to_dto(dbe: SessionStateDBE) -> SessionState:
data = SessionStateData.model_validate(dbe.data) if dbe.data else None
data = _parse_state_data(dbe)
return SessionState(
id=dbe.id,
project_id=dbe.project_id,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -202,3 +202,98 @@ def test_sandbox_id_accepts_runner_post(self, authed_api):
state = response.json()["session_state"]
assert state["session_id"] == session_id
assert state["sandbox_id"] == "sbx-runner"

def test_guarded_pointer_write_applies_at_latest_turn(self, authed_api):
session_id = str(uuid.uuid4())
authed_api(
"PUT",
"/sessions/states/",
params={"session_id": session_id},
json={
"data": {"latest_turn_index": 2},
"sandbox_id": "sbx-old",
},
)

response = authed_api(
"PUT",
"/sessions/states/",
params={"session_id": session_id},
json={
"sandbox_id": "sbx-new",
"sandbox_turn_index": 2,
},
)

assert response.status_code == 200
state = response.json()["session_state"]
assert state["sandbox_id"] == "sbx-new"

def test_stale_guarded_pointer_write_returns_unchanged_row(self, authed_api):
session_id = str(uuid.uuid4())
authed_api(
"PUT",
"/sessions/states/",
params={"session_id": session_id},
json={
"data": {"latest_turn_index": 3},
"sandbox_id": "sbx-current",
},
)

response = authed_api(
"PUT",
"/sessions/states/",
params={"session_id": session_id},
json={
"sandbox_id": "sbx-stale",
"sandbox_turn_index": 2,
},
)

assert response.status_code == 200
state = response.json()["session_state"]
assert state["sandbox_id"] == "sbx-current"
assert state["data"]["latest_turn_index"] == 3

def test_tokenless_pointer_write_remains_unconditional(self, authed_api):
session_id = str(uuid.uuid4())
authed_api(
"PUT",
"/sessions/states/",
params={"session_id": session_id},
json={
"data": {"latest_turn_index": 4},
"sandbox_id": "sbx-old",
},
)

response = authed_api(
"PUT",
"/sessions/states/",
params={"session_id": session_id},
json={
"sandbox_id": "sbx-tokenless",
},
)

assert response.status_code == 200
state = response.json()["session_state"]
assert state["sandbox_id"] == "sbx-tokenless"

Comment thread
coderabbitai[bot] marked this conversation as resolved.
def test_guarded_pointer_write_creates_missing_row(self, authed_api):
session_id = str(uuid.uuid4())
response = authed_api(
"PUT",
"/sessions/states/",
params={"session_id": session_id},
json={
"sandbox_id": "sbx-first",
"sandbox_turn_index": 7,
},
)

assert response.status_code == 200
state = response.json()["session_state"]
assert state["sandbox_id"] == "sbx-first"
assert state.get("data") is None
134 changes: 134 additions & 0 deletions docs/design/agent-workflows/projects/qa/scripts/warm_daytona_probe.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
# /// script
# requires-python = ">=3.10"
# dependencies = ["httpx>=0.27"]
# ///
"""Warm-daytona verification probe: timed turns against one session.

Drives the live agent service /invoke with sandbox=daytona and measures wall time per turn.
Turn 1 (no session header) is the cold create; turn 2 reuses the returned session id
immediately (park-to-running hit); a later rerun with --session <id> exercises the
stopped-restart path after the live window expires.

Usage:
uv run warm_daytona_probe.py # cold turn + immediate warm turn
uv run warm_daytona_probe.py --session <id> # one more turn on an existing session
uv run warm_daytona_probe.py --turns 1 # single cold turn only
"""

from __future__ import annotations

import argparse
import json
import os
import pathlib
import time

import httpx

BASE = os.environ.get("AGENTA_BASE", "http://localhost:8280")
PROJ = os.environ.get("AGENTA_PROJECT_ID", "019e8df5-2a58-7501-8fe2-56f7b332bd00")
MODEL = os.environ.get("AGENTA_QA_MODEL", "openai/gpt-4o-mini")


def api_key() -> str:
key = os.environ.get("AGENTA_API_KEY")
if key:
return key
repo = pathlib.Path(__file__).resolve().parents[6]
envf = repo / "examples/python/hotel_agent/draft/.env"
for line in envf.read_text().splitlines():
if line.startswith("AGENTA_API_KEY="):
return line.split("=", 1)[1].strip()
raise SystemExit("no AGENTA_API_KEY in env or hotel-agent .env")


def turn(
client: httpx.Client,
key: str,
history: list[dict],
message: str,
session_id: str | None,
) -> dict:
messages = [*history, {"role": "user", "content": message}]
body = {
"data": {
"inputs": {"messages": messages},
"parameters": {
"agent": {
"harness": {"kind": "pi_core"},
"sandbox": {"kind": "daytona"},
"llm": {"model": MODEL},
"instructions": {
"agents_md": "Reply with exactly the requested word and nothing else."
},
}
},
}
}
headers = {"Authorization": f"ApiKey {key}", "content-type": "application/json"}
if session_id:
headers["x-ag-session-id"] = session_id
started = time.monotonic()
resp = client.post(
f"{BASE}/services/agent/v0/invoke",
params={"project_id": PROJ},
headers=headers,
json=body,
timeout=240.0,
)
elapsed = time.monotonic() - started
out: dict = {"wall_s": round(elapsed, 2), "http": resp.status_code}
try:
payload = resp.json()
out["session_id"] = payload.get("session_id")
data = payload.get("data") or {}
outputs = data.get("outputs")
if isinstance(outputs, dict) and isinstance(outputs.get("messages"), list):
for msg in outputs["messages"]:
if msg.get("role") == "assistant":
out["reply"] = (msg.get("content") or "")[:120]
elif isinstance(outputs, dict):
out["reply"] = (outputs.get("content") or "")[:120]
out["status_message"] = ((payload.get("status") or {}).get("message") or "")[
:200
]
except Exception as exc: # noqa: BLE001
out["parse_error"] = str(exc)
out["raw"] = resp.text[:300]
return out


def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--session", default=None)
ap.add_argument("--turns", type=int, default=2)
ap.add_argument(
"--wait-before-last",
type=float,
default=0.0,
help="seconds to sleep before the final turn (to outlive the live window)",
)
args = ap.parse_args()

key = api_key()
session_id = args.session
history: list[dict] = []
with httpx.Client() as client:
for index in range(args.turns):
if args.wait_before_last and index == args.turns - 1:
print(json.dumps({"sleeping_s": args.wait_before_last}))
time.sleep(args.wait_before_last)
word = f"PING{index + 1}"
prompt = f"Reply with exactly the word {word}"
result = turn(client, key, history, prompt, session_id)
session_id = result.get("session_id") or session_id
history.append({"role": "user", "content": prompt})
history.append(
{"role": "assistant", "content": result.get("reply") or word}
)
label = "cold" if index == 0 and not args.session else "warm-candidate"
print(json.dumps({"turn": index + 1, "label": label, **result}))


if __name__ == "__main__":
main()
Loading
Loading