ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: align research with shared chats

Browse source
This commit is contained in:
alkinun 2026-07-18 16:59:25 +03:00
commit 633211bd1d
6 changed files with 200 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

29
studio/backend/routes/research_runs.py
View file

@ -63,8 +63,8 @@ class ApprovePlan(BaseModel):
planHash: str = Field(min_length = 64, max_length = 64)
def _require_run(run_id: str, subject: str) -> dict:
run = db.get_run(run_id, subject)
def _require_run(run_id: str) -> dict:
run = db.get_run(run_id)
if run is None:
raise HTTPException(status_code = 404, detail = "Research run not found")
return run
@ -270,14 +270,14 @@ async def active_research_runs(
thread_id: str = Query(alias = "threadId"), current_subject: str = Depends(get_current_subject)
):
return {
"runs": db.list_active(current_subject, thread_id),
"runs": db.list_active(thread_id),
"hasRun": db.has_thread_claim(thread_id),
}
@router.get("/{run_id}")
async def get_research_run(run_id: str, current_subject: str = Depends(get_current_subject)):
return _require_run(run_id, current_subject)
return _require_run(run_id)
@router.put("/{run_id}/plan")
@ -286,12 +286,12 @@ async def update_research_plan(
payload: UpdatePlan,
current_subject: str = Depends(get_current_subject),
):
_require_run(run_id, current_subject)
_require_run(run_id)
try:
db.set_plan(run_id, payload.plan.model_dump(), payload.expectedRevision)
except (db.ResearchConflictError, KeyError) as exc:
raise HTTPException(status_code = 409, detail = str(exc)) from exc
run = _require_run(run_id, current_subject)
run = _require_run(run_id)
_sync_assistant(run)
return run
@ -303,7 +303,7 @@ async def approve_research_plan(
request: Request,
current_subject: str = Depends(get_current_subject),
):
_require_run(run_id, current_subject)
_require_run(run_id)
try:
db.approve(run_id, payload.planRevision, payload.planHash)
except (db.ResearchConflictError, KeyError) as exc:
@ -312,7 +312,7 @@ async def approve_research_plan(
if supervisor is not None:
supervisor.note_request_port(request)
supervisor.wake()
run = _require_run(run_id, current_subject)
run = _require_run(run_id)
_sync_assistant(run)
return run
@ -323,12 +323,12 @@ async def cancel_research_run(
request: Request,
current_subject: str = Depends(get_current_subject),
):
_require_run(run_id, current_subject)
_require_run(run_id)
status = db.request_cancel(run_id)
supervisor = getattr(request.app.state, "research_supervisor", None)
if supervisor is not None and status == "cancelling":
supervisor.cancel(run_id)
run = _require_run(run_id, current_subject)
run = _require_run(run_id)
_sync_assistant(run)
return run
@ -339,7 +339,7 @@ async def retry_research_run(
request: Request,
current_subject: str = Depends(get_current_subject),
):
_require_run(run_id, current_subject)
_require_run(run_id)
try:
db.retry(run_id)
except (db.ResearchConflictError, KeyError) as exc:
@ -348,7 +348,7 @@ async def retry_research_run(
if supervisor is not None:
supervisor.note_request_port(request)
supervisor.wake()
run = _require_run(run_id, current_subject)
run = _require_run(run_id)
_sync_assistant(run)
return run
@ -361,7 +361,7 @@ async def research_events(
last_event_id: str | None = Header(None, alias = "Last-Event-ID"),
current_subject: str = Depends(get_current_subject),
):
_require_run(run_id, current_subject)
_require_run(run_id)
header_after = int(last_event_id) if last_event_id and last_event_id.isdigit() else 0
cursor = max(after or 0, header_after)
@ -371,11 +371,10 @@ async def research_events(
events = await asyncio.to_thread(
db.wait_for_events,
run_id,
current_subject,
cursor,
15,
)
snapshot = await asyncio.to_thread(db.get_run, run_id, current_subject)
snapshot = await asyncio.to_thread(db.get_run, run_id)
if snapshot is None:
return
for event in events:

24
studio/backend/storage/research_runs_db.py
View file

@ -310,18 +310,18 @@ def get_run(run_id: str, owner_subject: str | None = None) -> dict | None:
conn.close()
def list_active(owner_subject: str, thread_id: str) -> list[dict]:
def list_active(thread_id: str) -> list[dict]:
conn = get_connection()
try:
placeholders = ",".join("?" for _ in ACTIVE_STATUSES)
rows = conn.execute(
f"SELECT id FROM research_runs WHERE owner_subject = ? AND thread_id = ? "
f"SELECT id FROM research_runs WHERE thread_id = ? "
f"AND status IN ({placeholders}) ORDER BY created_at",
(owner_subject, thread_id, *sorted(ACTIVE_STATUSES)),
(thread_id, *sorted(ACTIVE_STATUSES)),
).fetchall()
finally:
conn.close()
return [run for row in rows if (run := get_run(row["id"], owner_subject)) is not None]
return [run for row in rows if (run := get_run(row["id"])) is not None]
def has_thread_claim(thread_id: str) -> bool:
@ -1130,17 +1130,16 @@ def upsert_document_source(
def list_events(
run_id: str,
owner_subject: str,
after: int = 0,
limit: int = 1000,
) -> list[dict]:
conn = get_connection()
try:
rows = conn.execute(
"""SELECT e.seq, e.event_type, e.data_json, e.created_at
FROM research_events e JOIN research_runs r ON r.id=e.run_id
WHERE e.run_id=? AND r.owner_subject=? AND e.seq>? ORDER BY e.seq LIMIT ?""",
(run_id, owner_subject, after, limit),
"""SELECT seq, event_type, data_json, created_at
FROM research_events
WHERE run_id=? AND seq>? ORDER BY seq LIMIT ?""",
(run_id, after, limit),
).fetchall()
return [
{
@ -1157,22 +1156,21 @@ def list_events(
def wait_for_events(
run_id: str,
owner_subject: str,
after: int = 0,
timeout: float = 15,
) -> list[dict]:
"""Block until committed events are available or the keep-alive timeout expires."""
events = list_events(run_id, owner_subject, after)
events = list_events(run_id, after)
if events:
return events
with _EVENTS_CHANGED:
# Recheck under the condition lock so a commit cannot be missed between
# the initial query and waiting for its notification.
events = list_events(run_id, owner_subject, after)
events = list_events(run_id, after)
if events:
return events
_EVENTS_CHANGED.wait(timeout)
return list_events(run_id, owner_subject, after)
return list_events(run_id, after)
def recover_expired(now: int | None = None) -> int:

80
studio/backend/storage/studio_db.py
View file

@ -1766,6 +1766,55 @@ def sync_chat_messages(
conn.close()
_RESEARCH_LINK_KEYS = {
"researchRunId",
"researchRun",
"researchStatus",
"researchPlanRevision",
"serverManaged",
}
def _detach_research_message_json(
content_json: str, metadata_json: str | None
) -> tuple[str, str | None]:
content = _json_loads(content_json, [])
metadata = _json_loads(metadata_json, None)
custom = metadata.get("custom") if isinstance(metadata, dict) else None
linked = (
isinstance(metadata, dict)
and any(key in metadata for key in _RESEARCH_LINK_KEYS)
or isinstance(custom, dict)
and any(key in custom for key in _RESEARCH_LINK_KEYS)
or isinstance(content, list)
and any(
isinstance(part, dict) and any(key in part for key in _RESEARCH_LINK_KEYS)
for part in content
)
)
if not linked:
return content_json, metadata_json
if isinstance(content, list):
content = [
{key: value for key, value in part.items() if key not in _RESEARCH_LINK_KEYS}
if isinstance(part, dict)
else part
for part in content
]
if isinstance(metadata, dict):
metadata = {key: value for key, value in metadata.items() if key not in _RESEARCH_LINK_KEYS}
custom = metadata.get("custom")
if isinstance(custom, dict):
metadata["custom"] = {
key: value for key, value in custom.items() if key not in _RESEARCH_LINK_KEYS
}
return (
json.dumps(content, ensure_ascii = False),
json.dumps(metadata, ensure_ascii = False) if metadata is not None else None,
)
def fork_chat_thread(
source_thread_id: str,
branch_message_id: str,
@ -1838,6 +1887,23 @@ def fork_chat_thread(
branch_message_id,
),
)
fork_messages = []
for row in ancestry:
content_json, metadata_json = _detach_research_message_json(
row["content_json"], row["metadata_json"]
)
fork_messages.append(
(
id_map[row["id"]],
new_thread_id,
id_map.get(row["parent_id"]) if row["parent_id"] else None,
row["role"],
content_json,
row["attachments_json"],
metadata_json,
int(row["created_at"]),
)
)
conn.executemany(
"""
INSERT INTO chat_messages
@ -1845,19 +1911,7 @@ def fork_chat_thread(
metadata_json, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
[
(
id_map[row["id"]],
new_thread_id,
id_map.get(row["parent_id"]) if row["parent_id"] else None,
row["role"],
row["content_json"],
row["attachments_json"],
row["metadata_json"],
int(row["created_at"]),
)
for row in ancestry
],
fork_messages,
)
conn.commit()
thread_row = conn.execute(

67
studio/backend/tests/test_chat_history_storage.py
View file

@ -602,6 +602,73 @@ def test_fork_chat_thread_preserves_project_id(tmp_path, monkeypatch):
}
def test_fork_chat_thread_detaches_research_run_metadata(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread("src"))
studio_db.upsert_chat_message(_msg("user", None, 1))
studio_db.upsert_chat_message(
{
"id": "research-report",
"threadId": "src",
"parentId": "user",
"role": "assistant",
"content": [
{
"type": "text",
"text": "# Copied report",
"researchRunId": "run-source",
},
{
"type": "source",
"url": "https://example.com",
"title": "Example",
"researchStatus": "completed",
},
],
"metadata": {
"researchRunId": "run-source",
"researchStatus": "completed",
"researchPlanRevision": 1,
"serverManaged": True,
"model": "local-model",
},
"createdAt": 2,
}
)
studio_db.fork_chat_thread(
source_thread_id = "src",
branch_message_id = "research-report",
new_thread_id = "fork-1",
new_title = "fork",
created_at = 3,
id_factory = iter(("fork-user", "fork-report")).__next__,
)
report = next(
message
for message in studio_db.list_chat_messages("fork-1")
if message["role"] == "assistant"
)
assert report["content"][0]["text"] == "# Copied report"
assert report["content"][1]["url"] == "https://example.com"
assert all(
not ({"researchRunId", "researchStatus", "serverManaged"} & set(part))
for part in report["content"]
)
assert report["metadata"] == {"model": "local-model"}
def test_fork_detachment_detects_non_id_research_content_keys():
content_json, metadata_json = studio_db._detach_research_message_json(
'[{"type":"text","text":"Report","serverManaged":true}]',
'{"model":"local-model"}',
)
assert "serverManaged" not in content_json
assert metadata_json == '{"model": "local-model"}'
def test_fork_chat_thread_returns_none_for_missing_source(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
result = studio_db.fork_chat_thread(

55
studio/backend/tests/test_research_runs_storage.py
View file

@ -426,9 +426,9 @@ def test_revision_hash_conflicts_and_idempotent_approval(research_home):
research_db.approve("run-1", 1, "0" * 64)
assert research_db.approve("run-1", 1, first["planHash"]) == "queued"
event_count = len(research_db.list_events("run-1", "alice"))
event_count = len(research_db.list_events("run-1"))
assert research_db.approve("run-1", 1, first["planHash"]) == "queued"
assert len(research_db.list_events("run-1", "alice")) == event_count
assert len(research_db.list_events("run-1")) == event_count
def test_planner_cannot_finalize_after_its_lease_timestamp_expires(research_home):
@ -489,7 +489,7 @@ def test_expired_worker_cannot_write_progress_or_execution_state(research_home):
)
is False
)
events = research_db.list_events("run-1", "alice")
events = research_db.list_events("run-1")
assert all(event["type"] != "reasoning.updated" for event in events)
assert research_db.finish("run-1", "worker-1", "completed") is None
assert research_db.get_run("run-1")["status"] == "running"
@ -527,30 +527,29 @@ def test_cancel_is_durable_and_idempotent(research_home):
_create()
research_db.set_plan("run-1", _plan())
assert research_db.request_cancel("run-1") == "cancelled"
event_count = len(research_db.list_events("run-1", "alice"))
event_count = len(research_db.list_events("run-1"))
assert research_db.request_cancel("run-1") == "cancelled"
run = research_db.get_run("run-1")
assert run["cancelRequested"] is True
assert len(research_db.list_events("run-1", "alice")) == event_count
assert len(research_db.list_events("run-1")) == event_count
def test_repeated_running_cancel_does_not_emit_duplicate_event(research_home):
_create()
assert research_db.claim_next("worker-1") is not None
assert research_db.request_cancel("run-1") == "cancelling"
event_count = len(research_db.list_events("run-1", "alice"))
event_count = len(research_db.list_events("run-1"))
assert research_db.request_cancel("run-1") == "cancelling"
assert len(research_db.list_events("run-1", "alice")) == event_count
assert len(research_db.list_events("run-1")) == event_count
def test_event_replay_is_monotonic_and_owner_scoped(research_home):
def test_event_replay_is_monotonic_for_shared_run(research_home):
_create()
for number in range(4):
research_db.append_event("run-1", "progress", {"number": number})
events = research_db.list_events("run-1", "alice", after = 2)
events = research_db.list_events("run-1", after = 2)
assert [event["seq"] for event in events] == [3, 4, 5]
assert [event["data"]["number"] for event in events] == [1, 2, 3]
assert research_db.list_events("run-1", "bob") == []
@pytest.mark.parametrize("status", ["planning", "queued", "running"])
@ -652,9 +651,7 @@ def test_sources_are_normalized_by_url(research_home):
assert source["snippet"] == "two"
assert source["stepPosition"] == 1
source_events = [
event
for event in research_db.list_events("run-1", "alice")
if event["type"] == "source.added"
event for event in research_db.list_events("run-1") if event["type"] == "source.added"
]
assert source_events[-1]["data"]["snippet"] == "two"
assert source_events[-1]["data"]["stepPosition"] == 1
@ -673,7 +670,7 @@ def test_partial_report_is_persisted_and_emits_an_event(research_home):
run = research_db.get_run("run-1")
assert run["report"] == "Partial report"
assert run["lastEventSeq"] == before + 1
[event] = research_db.list_events("run-1", "alice", after = before)
[event] = research_db.list_events("run-1", after = before)
assert event["type"] == "report.updated"
assert event["data"] == {"length": 14, "delta": " report", "offset": 7, "attempt": 0}
@ -857,7 +854,7 @@ def test_retry_is_bounded_and_resumes_from_saved_plan(research_home):
assert retried["steps"] == []
assert retried["sources"] == []
assert research_db.get_reasoning_text("run-1") == ""
assert research_db.list_events("run-1", "alice")[-1]["data"]["attempt"] == 1
assert research_db.list_events("run-1")[-1]["data"]["attempt"] == 1
research_db.claim_next("worker-2")
research_db.finish("run-1", "worker-2", "failed", "again")
with pytest.raises(research_db.ResearchConflictError, match = "budget"):
@ -1275,12 +1272,36 @@ def test_research_claim_is_global_across_authenticated_subjects(research_home):
assert research_db.has_thread_claim("thread-1") is True
def test_shared_chat_subject_can_follow_and_cancel_research(research_home):
from routes.research_runs import (
active_research_runs,
cancel_research_run,
get_research_run,
)
_create()
visible = asyncio.run(get_research_run("run-1", current_subject = "bob"))
active = asyncio.run(active_research_runs("thread-1", current_subject = "bob"))
cancelled = asyncio.run(
cancel_research_run(
"run-1",
SimpleNamespace(app = SimpleNamespace(state = SimpleNamespace())),
current_subject = "bob",
)
)
assert visible["ownerSubject"] == "alice"
assert [run["id"] for run in active["runs"]] == ["run-1"]
assert active["hasRun"] is True
assert cancelled["status"] == "cancelling"
def test_list_active_returns_complete_snapshots(research_home):
_create()
research_db.set_plan("run-1", _plan())
research_db.upsert_source("run-1", 0, "https://example.com/source", "Source", "Evidence")
[run] = research_db.list_active("alice", "thread-1")
[run] = research_db.list_active("thread-1")
assert [step["title"] for step in run["steps"]] == ["First", "Second"]
assert run["sources"][0]["url"] == "https://example.com/source"
@ -1463,7 +1484,7 @@ def test_cancel_requested_wins_finish_cas(research_home, requested):
snapshot = research_db.get_run("run-1")
assert snapshot["status"] == "cancelled"
assert snapshot["report"] is None
terminal = research_db.list_events("run-1", "alice")[-1]
terminal = research_db.list_events("run-1")[-1]
assert terminal["type"] == "run.cancelled"
assert "report" not in terminal["data"]
assert terminal["data"]["error"] is None

4
tests/studio/test_deep_research_frontend_contract.py
View file

@ -142,6 +142,10 @@ def test_research_presentation_is_integrated() -> None:
assert "effectiveDeepResearchEnabled ||" in thread
assert "replayFrom: session?.lastAppliedSeq ?? 0" in coordinator
assert "loadBool(CHAT_DEEP_RESEARCH_ENABLED_KEY, false)" in store
checkpoint_update = store.split("setCheckpoint: (modelId, ggufVariant) =>", 1)[1].split(
"setActiveThreadId:", 1
)[0]
assert "saveBool(CHAT_DEEP_RESEARCH_ENABLED_KEY, false)" in checkpoint_update
assert "const permissionMode = loadPermissionMode();" in store
assert "permissionMode," in store