From 725cece5c10e88870878cadeb2cc33a8fc31dea5 Mon Sep 17 00:00:00 2001
From: Roberto Musso <roberto.musso@hpe.com>
Date: Fri, 20 Mar 2026 09:46:17 +0100
Subject: [PATCH] Add run_context to agent tool calls for FE run logging

- AgentTriggerRequest accepts optional agent_id (FE's stable electron-store UUID)
- _make_agent_executor injects run_context into every tool_call frame
  so Electron can attribute actions to the correct agent run
- run_local_agent accepts run_context and sends a run_complete WS frame
  when the run finishes so the FE can close the run record
- trigger_agent_run builds run_context with run_id=run_log.id and the
  stable agent_id, passes it through to run_local_agent

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 app/api/routes/agents.py | 13 +++++++++++--
 app/core/agent_runner.py | 21 ++++++++++++++++++++-
 app/schemas.py           |  1 +
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/app/api/routes/agents.py b/app/api/routes/agents.py
index fbb8cc0..53d0edd 100644
--- a/app/api/routes/agents.py
+++ b/app/api/routes/agents.py
@@ -190,8 +190,11 @@ async def trigger_agent_run(
         enabled=True,
     )
 
+    # Use the FE's stable agent_id if provided, fall back to the ephemeral config id.
+    stable_agent_id = body.agent_id or config.id
+
     run_log = AgentRunLog(
-        agent_id=config.id,
+        agent_id=stable_agent_id,
         agent_type="local",
         user_id=current_user.id,
         status="running",
@@ -200,8 +203,14 @@ async def trigger_agent_run(
     await db.commit()
     await db.refresh(run_log)
 
+    run_context = {
+        "type": "agent_batch",
+        "run_id": run_log.id,
+        "agent_id": stable_agent_id,
+    }
+
     asyncio.create_task(
-        run_local_agent(current_user.id, config, run_log, device_manager)
+        run_local_agent(current_user.id, config, run_log, device_manager, run_context)
     )
 
     return _to_run_log_response(run_log)
diff --git a/app/core/agent_runner.py b/app/core/agent_runner.py
index aaa8aef..4926a6d 100644
--- a/app/core/agent_runner.py
+++ b/app/core/agent_runner.py
@@ -188,12 +188,18 @@ def _is_overdue(schedule_cron: str, last_run_at: datetime | None) -> bool:
 def _make_agent_executor(
     user_id: str,
     device_mgr: DeviceConnectionManager,
+    run_context: dict | None = None,
 ) -> Any:
     """Create a WS callback for ``set_client_executor()`` so that all tools
     can use ``execute_on_client()`` during an agent run.
+
+    If *run_context* is provided it is attached to every ``tool_call`` frame
+    so the Electron client can attribute actions to the correct agent run.
     """
     async def _executor(payload: dict) -> dict:
         payload["type"] = "tool_call"
+        if run_context:
+            payload["run_context"] = run_context
         call_id = payload["id"]
         fut = device_mgr.create_pending_call(user_id, call_id)
         await device_mgr.send_frame(user_id, payload)
@@ -328,6 +334,7 @@ async def run_local_agent(
     config: LocalAgentConfig,
     run_log: AgentRunLog,
     device_mgr: DeviceConnectionManager,
+    run_context: dict | None = None,
 ) -> None:
     """Execute a local directory agent run using two-phase LLM-with-tools.
 
@@ -363,7 +370,7 @@ async def run_local_agent(
         return
 
     # ── Set up WS executor for tools ────────────────────────────────
-    executor = _make_agent_executor(user_id, device_mgr)
+    executor = _make_agent_executor(user_id, device_mgr, run_context)
     set_client_executor(executor)
 
     errors: list[str] = []
@@ -508,6 +515,18 @@ async def run_local_agent(
         len(errors),
     )
 
+    # Notify the Electron client that the run is complete so it can close
+    # the run record in its local SQLite.
+    if run_context and device_mgr.is_online(user_id):
+        try:
+            await device_mgr.send_frame(user_id, {
+                "type": "run_complete",
+                "run_context": run_context,
+                "status": final_status,
+            })
+        except Exception as exc:
+            logger.warning("agent_runner: run=%s failed to send run_complete: %s", run_id, exc)
+
 
 # ── Cloud agent runner ─────────────────────────────────────────────────────
 
diff --git a/app/schemas.py b/app/schemas.py
index 3e8a034..39143c4 100644
--- a/app/schemas.py
+++ b/app/schemas.py
@@ -295,6 +295,7 @@ class AgentCreationCheckResponse(BaseModel):
 class AgentTriggerRequest(BaseModel):
     directory: str = Field(min_length=1)
     device_id: str = Field(default="")
+    agent_id: str | None = None  # FE stable agent ID (electron-store UUID)
     what_to_extract: list[str] = Field(min_length=1)
     actions_by_type: dict[str, list[str]] | None = None
     batch_interval: str = Field(min_length=1)