diff --git a/packages/narada-core/src/narada_core/actions/models.py b/packages/narada-core/src/narada_core/actions/models.py index 07d7176..4b22ad5 100644 --- a/packages/narada-core/src/narada_core/actions/models.py +++ b/packages/narada-core/src/narada_core/actions/models.py @@ -12,10 +12,7 @@ override, ) -from pydantic import ( - BaseModel, - Field, -) +from pydantic import BaseModel, ConfigDict, Field from narada_core.tracing import model as tracing_model @@ -50,6 +47,8 @@ class CriticResult(BaseModel): class AgentResponse(BaseModel, Generic[_StructuredOutputT]): + model_config = ConfigDict(populate_by_name=True) + request_id: str status: Literal["success", "error", "input-required"] text: str @@ -60,6 +59,7 @@ class AgentResponse(BaseModel, Generic[_StructuredOutputT]): ] usage: AgentUsage action_trace: tracing_model.ActionTrace | None = None + workflow_trace: dict[str, Any] | None = Field(default=None, alias="workflowTrace") critic_result: CriticResult | None = None @@ -458,3 +458,5 @@ class ExtensionActionResponse(BaseModel): status: Literal["success", "error", "aborted"] error: str | None = None data: str | None = None + action_trace: tracing_model.ActionTrace | None = None + workflowTrace: dict[str, Any] | None = None diff --git a/packages/narada-core/src/narada_core/models.py b/packages/narada-core/src/narada_core/models.py index 4685dc5..af468a2 100644 --- a/packages/narada-core/src/narada_core/models.py +++ b/packages/narada-core/src/narada_core/models.py @@ -1,7 +1,7 @@ from __future__ import annotations from enum import Enum, StrEnum -from typing import Annotated, Generic, Literal, NotRequired, TypedDict, TypeVar +from typing import Annotated, Any, Generic, Literal, NotRequired, TypedDict, TypeVar from pydantic import BaseModel, Field @@ -399,6 +399,7 @@ class ResponseContent(TypedDict, Generic[_MaybeStructuredOutput]): text: str structuredOutput: _MaybeStructuredOutput actionTrace: NotRequired[ActionTrace] + workflowTrace: NotRequired[dict[str, Any]] class Usage(TypedDict): diff --git a/packages/narada-core/src/narada_core/tracing/model.py b/packages/narada-core/src/narada_core/tracing/model.py index e4d2712..b425c32 100644 --- a/packages/narada-core/src/narada_core/tracing/model.py +++ b/packages/narada-core/src/narada_core/tracing/model.py @@ -189,6 +189,7 @@ class RunCustomAgentTrace(BaseModel): workflow_name: str status: Literal["success", "error"] error_message: str | None = None + subWorkflow: dict[str, Any] | None = None children: ActionTrace | None = None @@ -302,6 +303,11 @@ def _check_ts_ordering(self) -> PythonExtensionActionEvent: return self +class PythonSubWorkflowEvent(BaseModel): + kind: Literal["subWorkflow"] = "subWorkflow" + workflowTrace: dict[str, Any] + + class PythonSideEffectEvent(BaseModel): kind: Literal["sideEffect"] = "sideEffect" ts: int @@ -314,6 +320,7 @@ class PythonSideEffectEvent(BaseModel): | PythonStderrEvent | PythonSubAgentCallEvent | PythonExtensionActionEvent + | PythonSubWorkflowEvent | PythonSideEffectEvent, Field(discriminator="kind"), ] diff --git a/packages/narada-pyodide/src/narada/_trace.py b/packages/narada-pyodide/src/narada/_trace.py index 52be5bf..1c08718 100644 --- a/packages/narada-pyodide/src/narada/_trace.py +++ b/packages/narada-pyodide/src/narada/_trace.py @@ -133,6 +133,15 @@ def emit_extension_action( emit_trace_event(event) +def emit_sub_workflow(*, workflow_trace: dict[str, Any]) -> None: + emit_trace_event( + { + "kind": "subWorkflow", + "workflowTrace": workflow_trace, + } + ) + + def emit_side_effect(*, effect_type: SideEffectType, description: str) -> None: emit_trace_event( { diff --git a/packages/narada-pyodide/src/narada/window.py b/packages/narada-pyodide/src/narada/window.py index 6c2afb8..c69444b 100644 --- a/packages/narada-pyodide/src/narada/window.py +++ b/packages/narada-pyodide/src/narada/window.py @@ -1,4 +1,5 @@ import asyncio +import builtins import json import logging import os @@ -23,10 +24,10 @@ from narada_core.actions.critic import run_critic from narada_core.actions.models import ( DEFAULT_HITL_TIMEOUT_SECONDS, - AgenticMouseAction, - AgenticMouseActionRequest, AgenticMatchingSelectorsFinderRequest, AgenticMatchingSelectorsFinderResponse, + AgenticMouseAction, + AgenticMouseActionRequest, AgenticSelectorAction, AgenticSelectorRequest, AgenticSelectorResponse, @@ -103,10 +104,20 @@ def _parent_run_ids() -> list[str]: ) +def _parent_request_id() -> str | None: + parent_request_id = getattr(builtins, "_narada_request_id", None) + if isinstance(parent_request_id, str): + return parent_request_id + parent_request_id = globals().get("_narada_request_id") + return parent_request_id if isinstance(parent_request_id, str) else None + + if TYPE_CHECKING: # Magic function injected by the JavaScript harness to get the current user's ID token. async def _narada_get_id_token() -> str: ... + _narada_request_id: str | None + _StructuredOutput = TypeVar("_StructuredOutput", bound=BaseModel) @@ -215,6 +226,10 @@ def _current_parent_run_ids(self) -> list[str] | None: """ return None + def _current_parent_request_id(self) -> str | None: + """Returns the remote-dispatch request that owns the current Python execution.""" + return _parent_request_id() + async def _get_auth_headers(self) -> dict[str, str]: return await _build_auth_headers( api_key=self._api_key, @@ -387,6 +402,9 @@ async def dispatch_request( parent_run_ids = self._current_parent_run_ids() if parent_run_ids: body["parentRunIds"] = parent_run_ids + parent_request_id = self._current_parent_request_id() + if parent_request_id is not None: + body["parentRequestId"] = parent_request_id cloud_browser_session_id = self.cloud_browser_session_id if cloud_browser_session_id is not None: body["cloudBrowserSessionId"] = cloud_browser_session_id @@ -689,6 +707,12 @@ async def agent( if action_trace_raw is not None else None ) + workflow_trace = response_content.get("workflowTrace") + parent_request_id = self._current_parent_request_id() + # Preserve the response contract for direct callers, but avoid adding a second + # child node when the backend will stitch the child request into the parent row. + if workflow_trace is not None and parent_request_id is None: + _trace.emit_sub_workflow(workflow_trace=workflow_trace) critic_result: CriticResult | None = None if critic is not None: @@ -710,6 +734,7 @@ async def agent( structured_output=response_content.get("structuredOutput"), usage=AgentUsage.model_validate(remote_dispatch_response["usage"]), action_trace=action_trace, + workflow_trace=workflow_trace, critic_result=critic_result, ) @@ -1008,6 +1033,9 @@ async def _run_extension_action( parent_run_ids = self._current_parent_run_ids() if parent_run_ids: body["parentRunIds"] = parent_run_ids + parent_request_id = self._current_parent_request_id() + if parent_request_id is not None: + body["requestId"] = parent_request_id if timeout is not None: body["timeout"] = timeout @@ -1029,6 +1057,8 @@ async def _run_extension_action( resp_json = await fetch_response.json() response = ExtensionActionResponse.model_validate(resp_json) + if response.workflowTrace is not None: + _trace.emit_sub_workflow(workflow_trace=response.workflowTrace) if response.status == "error": raise NaradaError(response.error) if response.status == "aborted": diff --git a/packages/narada-pyodide/tests/test_cloud_browser.py b/packages/narada-pyodide/tests/test_cloud_browser.py index 75561d3..c5b4081 100644 --- a/packages/narada-pyodide/tests/test_cloud_browser.py +++ b/packages/narada-pyodide/tests/test_cloud_browser.py @@ -1,5 +1,6 @@ from __future__ import annotations +import builtins import json import sys from pathlib import Path @@ -85,6 +86,8 @@ def new() -> SimpleNamespace: client_module = importlib.import_module("narada.client") window_module = importlib.import_module("narada.window") window_module._narada_parent_run_ids = _FakeJsProxy([]) + window_module._narada_request_id = None + monkeypatch.setattr(builtins, "_narada_request_id", None, raising=False) window_module._narada_get_id_token = AsyncMock(return_value="frontend-id-token") return narada_pkg, client_module, window_module @@ -277,6 +280,123 @@ async def test_cloud_browser_window_dispatch_request_omits_parent_run_ids( assert "parentRunIds" not in payload +@pytest.mark.asyncio +async def test_cloud_browser_window_dispatch_request_keeps_parent_request_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pyfetch = AsyncMock( + side_effect=[ + _FakeResponse(json_data={"requestId": "child-request-123"}), + _FakeResponse(json_data={"status": "success", "response": None}), + ] + ) + _, _, window_module = _import_pyodide_narada(monkeypatch, pyfetch=pyfetch) + window_module._narada_parent_run_ids = _FakeJsProxy(["outer-run", "inner-run"]) + monkeypatch.setattr( + builtins, "_narada_request_id", "parent-request-123", raising=False + ) + + window = window_module.CloudBrowserWindow( + browser_window_id="browser-window-123", + session_id="session-123", + api_key="test-api-key", + ) + response = await window.dispatch_request(prompt="hello from cloud browser") + + assert response["status"] == "success" + post_call = pyfetch.await_args_list[0] + payload = json.loads(post_call.kwargs["body"]) + assert payload["parentRequestId"] == "parent-request-123" + assert "parentRunIds" not in payload + + +@pytest.mark.asyncio +async def test_window_agent_keeps_parent_request_id_from_injected_builtins( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pyfetch = AsyncMock( + side_effect=[ + _FakeResponse(json_data={"requestId": "child-request-123"}), + _FakeResponse( + json_data={ + "status": "success", + "response": { + "text": "done", + "output": {"type": "text", "content": "done"}, + }, + "usage": {"actions": 0, "credits": 0}, + } + ), + ] + ) + _, _, window_module = _import_pyodide_narada(monkeypatch, pyfetch=pyfetch) + window_module._narada_parent_run_ids = _FakeJsProxy(["outer-run", "inner-run"]) + monkeypatch.setattr( + builtins, "_narada_request_id", "parent-request-123", raising=False + ) + + window = window_module.CloudBrowserWindow( + browser_window_id="browser-window-123", + session_id="session-123", + api_key="test-api-key", + ) + response = await window.agent(prompt="run gui child", agent="/$USER/gui-child") + + assert response.status == "success" + post_call = pyfetch.await_args_list[0] + payload = json.loads(post_call.kwargs["body"]) + assert payload["parentRequestId"] == "parent-request-123" + + +@pytest.mark.asyncio +async def test_window_agent_exposes_workflow_trace_alias( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_trace = {"step_type": "workflow", "children": []} + pyfetch = AsyncMock( + side_effect=[ + _FakeResponse(json_data={"requestId": "child-request-123"}), + _FakeResponse( + json_data={ + "status": "success", + "response": { + "text": "done", + "output": {"type": "text", "content": "done"}, + "workflowTrace": workflow_trace, + }, + "usage": {"actions": 0, "credits": 0}, + } + ), + ] + ) + _, _, window_module = _import_pyodide_narada(monkeypatch, pyfetch=pyfetch) + emitted_events: list[str] = [] + monkeypatch.setattr( + sys.modules["narada._trace"], + "_narada_emit_trace_event", + emitted_events.append, + raising=False, + ) + + window = window_module.CloudBrowserWindow( + browser_window_id="browser-window-123", + session_id="session-123", + api_key="test-api-key", + ) + response = await window.agent(prompt="return a trace") + + assert response.workflow_trace == workflow_trace + assert response.model_dump(by_alias=True)["workflowTrace"] == workflow_trace + sub_workflow_events = [ + json.loads(event) + for event in emitted_events + if json.loads(event)["kind"] == "subWorkflow" + ] + assert sub_workflow_events == [ + {"kind": "subWorkflow", "workflowTrace": workflow_trace} + ] + + @pytest.mark.asyncio async def test_cloud_browser_window_dispatch_request_retries_poll_fetch_failures( monkeypatch: pytest.MonkeyPatch, @@ -624,6 +744,32 @@ async def test_remote_browser_window_without_cloud_session_keeps_extension_actio assert "parentRunIds" not in payload +@pytest.mark.asyncio +async def test_remote_browser_window_extension_action_keeps_parent_request_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pyfetch = AsyncMock( + return_value=_FakeResponse(json_data={"status": "success", "data": None}) + ) + _, _, window_module = _import_pyodide_narada(monkeypatch, pyfetch=pyfetch) + window_module._narada_parent_run_ids = _FakeJsProxy(["outer-run", "inner-run"]) + monkeypatch.setattr( + builtins, "_narada_request_id", "parent-request-123", raising=False + ) + + window = window_module.RemoteBrowserWindow( + browser_window_id="browser-window-123", + api_key="test-api-key", + ) + await window.close() + + call = pyfetch.await_args + assert call is not None + payload = json.loads(call.kwargs["body"]) + assert payload["requestId"] == "parent-request-123" + assert "parentRunIds" not in payload + + @pytest.mark.asyncio async def test_local_browser_window_dispatch_request_uses_latest_parent_run_ids( monkeypatch: pytest.MonkeyPatch, @@ -655,3 +801,53 @@ async def test_local_browser_window_dispatch_request_uses_latest_parent_run_ids( second_post = json.loads(pyfetch.await_args_list[2].kwargs["body"]) assert first_post["parentRunIds"] == ["run-a"] assert second_post["parentRunIds"] == ["run-b", "run-c"] + + +@pytest.mark.asyncio +async def test_local_browser_window_dispatch_request_includes_parent_request_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("NARADA_API_KEY", "test-api-key") + monkeypatch.setenv("NARADA_BROWSER_WINDOW_ID", "browser-window-123") + pyfetch = AsyncMock( + side_effect=[ + _FakeResponse(json_data={"requestId": "child-request-123"}), + _FakeResponse(json_data={"status": "success", "response": None}), + ] + ) + _, _, window_module = _import_pyodide_narada(monkeypatch, pyfetch=pyfetch) + window_module._narada_parent_run_ids = _FakeJsProxy(["run-a"]) + monkeypatch.setattr( + builtins, "_narada_request_id", "parent-request-123", raising=False + ) + + window = window_module.LocalBrowserWindow() + response = await window.dispatch_request(prompt="child prompt") + + assert response["status"] == "success" + post_payload = json.loads(pyfetch.await_args_list[0].kwargs["body"]) + assert post_payload["parentRequestId"] == "parent-request-123" + assert post_payload["parentRunIds"] == ["run-a"] + + +@pytest.mark.asyncio +async def test_local_browser_window_extension_action_includes_parent_request_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("NARADA_API_KEY", "test-api-key") + monkeypatch.setenv("NARADA_BROWSER_WINDOW_ID", "browser-window-123") + pyfetch = AsyncMock( + return_value=_FakeResponse(json_data={"status": "success", "data": None}) + ) + _, _, window_module = _import_pyodide_narada(monkeypatch, pyfetch=pyfetch) + window_module._narada_parent_run_ids = _FakeJsProxy(["run-a"]) + monkeypatch.setattr( + builtins, "_narada_request_id", "parent-request-123", raising=False + ) + + window = window_module.LocalBrowserWindow() + await window.close() + + post_payload = json.loads(pyfetch.await_args.kwargs["body"]) + assert post_payload["requestId"] == "parent-request-123" + assert post_payload["parentRunIds"] == ["run-a"] diff --git a/packages/narada/src/narada/window.py b/packages/narada/src/narada/window.py index 8df4bb3..494a976 100644 --- a/packages/narada/src/narada/window.py +++ b/packages/narada/src/narada/window.py @@ -643,6 +643,7 @@ async def agent( if action_trace_raw is not None else None ) + workflow_trace = response_content.get("workflowTrace") critic_result: CriticResult | None = None if critic is not None: @@ -664,6 +665,7 @@ async def agent( structured_output=response_content.get("structuredOutput"), usage=AgentUsage.model_validate(remote_dispatch_response["usage"]), action_trace=action_trace, + workflow_trace=workflow_trace, critic_result=critic_result, ) diff --git a/packages/narada/tests/test_cloud_browser.py b/packages/narada/tests/test_cloud_browser.py index db54062..9e5f6f5 100644 --- a/packages/narada/tests/test_cloud_browser.py +++ b/packages/narada/tests/test_cloud_browser.py @@ -4,6 +4,7 @@ import pytest from narada.client import Narada from narada.config import BrowserConfig +from narada.window import CloudBrowserWindow from narada_core.errors import NaradaTimeoutError @@ -171,3 +172,36 @@ async def test_initialize_cloud_browser_window_uses_domcontentloaded_for_retry_n ] assert wait_for_browser_window_id.await_count == 2 assert window.browser_window_id == "browser-window-123" + + +@pytest.mark.asyncio +async def test_window_agent_exposes_workflow_trace_alias( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_trace = {"step_type": "workflow", "children": []} + window = CloudBrowserWindow( + browser_window_id="browser-window-123", + session_id="session-123", + auth_headers={"x-api-key": "test-key"}, + ) + monkeypatch.setattr( + window, + "dispatch_request", + AsyncMock( + return_value={ + "requestId": "request-123", + "status": "success", + "response": { + "text": "done", + "output": {"type": "text", "content": "done"}, + "workflowTrace": workflow_trace, + }, + "usage": {"actions": 0, "credits": 0}, + } + ), + ) + + response = await window.agent(prompt="return a trace") + + assert response.workflow_trace == workflow_trace + assert response.model_dump(by_alias=True)["workflowTrace"] == workflow_trace