diff --git a/studio/backend/routes/research_runs.py b/studio/backend/routes/research_runs.py index 479b634544..837ae7a370 100644 --- a/studio/backend/routes/research_runs.py +++ b/studio/backend/routes/research_runs.py @@ -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: diff --git a/studio/backend/storage/research_runs_db.py b/studio/backend/storage/research_runs_db.py index fcdf8a1d95..564510da7e 100644 --- a/studio/backend/storage/research_runs_db.py +++ b/studio/backend/storage/research_runs_db.py @@ -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: diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index d19270e23f..f870052b72 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -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( diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index 0239410734..c99c860cea 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -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( diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index 6668794a3e..4a3c981f08 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -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 diff --git a/tests/studio/test_deep_research_frontend_contract.py b/tests/studio/test_deep_research_frontend_contract.py index 55a33ff24e..8adec878a5 100644 --- a/tests/studio/test_deep_research_frontend_contract.py +++ b/tests/studio/test_deep_research_frontend_contract.py @@ -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