Merge pull request #173 from latchbio/aidan/combine-signals-with-cells

Agent Context -- Combine Signals Into Cells

Author: AidanAbd

runtime/mount/agent.py

@@ -26,6 +26,8 @@
 
 cache_chunk_size = 20
 
+reactivity_ready_statuses = {"ran", "ok", "success", "error"}
+
 sandbox_root = os.environ.get("LATCH_SANDBOX_ROOT")
 if sandbox_root:
     import pathlib
@@ -2504,8 +2506,11 @@ async def get_global_info(args: dict) -> dict:
         })
         self.tool_map["get_global_info"] = get_global_info
 
-        async def refresh_cells_context(args: dict) -> dict:
-            context_result = await self.atomic_operation("get_context")
+        async def refresh_cells_context(_args: dict) -> dict:
+            context_result, reactivity_result = await asyncio.gather(
+                self.atomic_operation("get_context"),
+                self.atomic_operation("request_reactivity_summary"),
+            )
 
             if context_result.get("status") != "success":
                 return {
@@ -2514,6 +2519,14 @@ async def refresh_cells_context(args: dict) -> dict:
                     "summary": f"Failed to refresh cells: {context_result.get('error', 'Unknown error')}"
                 }
 
+            if reactivity_result.get("status") != "success":
+                print(f"[tool] refresh_cells_context: failed to get reactivity summary: {reactivity_result.get('error', 'Unknown error')}")
+                return {
+                    "tool_name": "refresh_cells_context",
+                    "success": False,
+                    "summary": f"Failed to get reactivity summary: {reactivity_result.get('error', 'Unknown error')}"
+                }
+
             context = context_result.get("context", {})
             self.latest_notebook_context = context
 
@@ -2533,6 +2546,8 @@ async def refresh_cells_context(args: dict) -> dict:
 
             current_tab_name = default_tab_name
 
+            cell_reactivity: dict[str, dict[str, object]] = reactivity_result.get("cell_reactivity", {})
+
             for cell in cells:
                 index = cell.get("index", "?")
                 cell_id = cell.get("cell_id", "?")
@@ -2578,65 +2593,61 @@ async def refresh_cells_context(args: dict) -> dict:
                         w_label = w.get("label", "")
                         cell_lines.append(f"- WIDGET: {w_type} | {w_label} | {w_key}")
 
-            context_dir = context_root / "notebook_context"
-            context_dir.mkdir(parents=True, exist_ok=True)
-            (context_dir / "cells.md").write_text("\n".join(cell_lines))
+                if tf_id is None:
+                    continue
 
-            return {
-                "tool_name": "refresh_cells_context",
-                "success": True,
-                "summary": f"Refreshed cells context for {cell_count} cells and stored result in notebook_context/cells.md",
-                "cell_count": cell_count,
-                "context_path": "notebook_context/cells.md"
-            }
+                reactivity_meta = cell_reactivity.get(str(tf_id))
+                is_reactivity_ready = status in reactivity_ready_statuses
 
-        async def refresh_reactivity_context(args: dict) -> dict:
-            reactivity_result = await self.atomic_operation("request_reactivity_summary")
+                cell_lines.append("\nREACTIVITY:")
 
-            if reactivity_result.get("status") != "success":
-                return {
-                    "tool_name": "refresh_reactivity_context",
-                    "success": False,
-                    "summary": f"Failed to refresh reactivity: {reactivity_result.get('error', 'Unknown error')}"
-                }
+                if not is_reactivity_ready:
+                    cell_lines.append("- Not available: run this cell to establish reactive dependencies.")
+                    continue
+
+                if reactivity_meta is None:
+                    cell_lines.append("- Reactivity data missing for this cell.")
+                    continue
 
-            reactivity_summary = reactivity_result.get("summary")
+                signals_defined = reactivity_meta.get("signals_defined", [])
+                depends_on_signals = reactivity_meta.get("depends_on_signals", [])
+                depends_on_cells = reactivity_meta.get("depends_on_cells", [])
+
+                cell_lines.append(  # noqa: FURB113
+                    "- Signals defined: "
+                    + (", ".join(signals_defined) if signals_defined else "None")
+                )
+                cell_lines.append(
+                    "- Depends on signals: "
+                    + (", ".join(depends_on_signals) if depends_on_signals else "None")
+                )
+                cell_lines.append(
+                    "- Depends on cells: "
+                    + (", ".join(depends_on_cells) if depends_on_cells else "None")
+                )
 
             context_dir = context_root / "notebook_context"
             context_dir.mkdir(parents=True, exist_ok=True)
-
-            if reactivity_summary is not None:
-                (context_dir / "signals.md").write_text(reactivity_summary)
-            else:
-                (context_dir / "signals.md").write_text("# Reactivity Summary\n\nNo reactive dependencies in this notebook.\n")
+            (context_dir / "cells.md").write_text("\n".join(cell_lines))
 
             return {
-                "tool_name": "refresh_reactivity_context",
+                "tool_name": "refresh_cells_context",
                 "success": True,
-                "summary": "Refreshed reactivity context and stored result in notebook_context/signals.md",
-                "context_path": "notebook_context/signals.md",
+                "summary": f"Refreshed reactive notebook context for {cell_count} cells and stored result in notebook_context/cells.md",
+                "cell_count": cell_count,
+                "context_path": "notebook_context/cells.md"
             }
 
         self.tools.append({
             "name": "refresh_cells_context",
-            "description": "Refresh the cells.md context file with current notebook cell structure and contents.",
+            "description": "Refresh the cells.md context file with current reactive notebook structure and contents.",
             "input_schema": {
                 "type": "object",
                 "properties": {},
             },
         })
         self.tool_map["refresh_cells_context"] = refresh_cells_context
 
-        self.tools.append({
-            "name": "refresh_reactivity_context",
-            "description": "Refresh the signals.md context file with current reactive signal dependencies.",
-            "input_schema": {
-                "type": "object",
-                "properties": {},
-            },
-        })
-        self.tool_map["refresh_reactivity_context"] = refresh_reactivity_context
-
         if len(self.tools) > 0:
             self.tools[-1]["cache_control"] = {"type": "ephemeral"}
 
@@ -3198,7 +3209,6 @@ def read_context_file(filename: str) -> str:
             return f"# {filename}\n\nFile not yet generated."
 
         cells_content = read_context_file("cells.md")
-        signals_content = read_context_file("signals.md")
 
         def build_tree(path: Path, prefix: str = "") -> list[str]:
             lines = []
@@ -3224,7 +3234,6 @@ def build_tree(path: Path, prefix: str = "") -> list[str]:
             "truncated_messages": truncated_messages,
             "model": self.mode_config.get(self.mode, ("claude-sonnet-4-5-20250929", 1024))[0],
             "cells": cells_content,
-            "signals": signals_content,
             "tree": tree_content,
         }
 

runtime/mount/agent_config/README.md

@@ -23,10 +23,8 @@ agent_config/
 │   │       ├── reactivity.mdx
 │   │       └── widget-types.mdx
 │   ├── agent_scratch/            # Agent-created notes and state
-│   └── notebook_context/         # Auto-generated runtime state
-│       ├── cells.md              # (populated by the agent calling tool to refresh the context)
-│       ├── globals.md            # (populated by the agent calling tool to refresh the context)
-│       └── signals.md            # (populated by the agent calling tool to refresh the context)
+│   └── notebook_context/
+│       └── cells.md              # Populated by the agent via refresh_cells_context
 └── README.md                     # This file
 ```
 

runtime/mount/agent_config/system_prompt.md

@@ -726,10 +726,9 @@ See reactivity documentation in the `## Reactivity` section of `latch_api_docs/l
 
 ### Context Refresh Tools
 
-- `refresh_cells_context` - Update cells.md with current notebook structure
-- `refresh_reactivity_context` - Update signals.md with signal dependencies
+- `refresh_cells_context` - Update `cells.md` with current reactive notebook structure
 
-Use these tools to refresh notebook state files before reading them.
+Use this tool to refresh `cells.md` with the current reactive notebook state before reading it.
 
 ### Introspection Tools
 
@@ -744,7 +743,7 @@ Use these tools to refresh notebook state files before reading them.
 
 ## Context Files (Refresh On-Demand)
 
-The notebook state is available in two context files. These files are initialized when you start but are NOT automatically updated - you must explicitly refresh them when needed.
+The reactive notebook state is persisted in `cells.md`. You must explicitly refresh it when needed.
 
 ### cells.md
 
@@ -756,37 +755,20 @@ The notebook state is available in two context files. These files are initialize
 - After creating cells to see updated structure
 - Looking for specific code or widget locations
 - Checking cell execution status
+- Creating or reasoning about reactive relationships between cells
 
 **Refresh tool:** `refresh_cells_context`
 
 - Returns updated cell count
 - Returns context path to read result from
-- Writes latest cell structure to cells.md
+- Writes latest cell structure and, for every cell, the signals it defines and the signals/cells it depends on
 
 **Search with grep:**
 
 - Find by ID: `grep "CELL_ID: abc123" notebook_context/cells.md`
 - Find by code: `grep "import pandas" notebook_context/cells.md`
 
-**Format:** Cell metadata on separate lines (CELL_ID, CELL_INDEX, TYPE, STATUS), code between CODE_START/CODE_END markers.
-
-### signals.md
-
-**Location:** `notebook_context/signals.md`
-
-**Refresh when:**
-
-- Designing reactive workflows
-- Debugging reactivity issues
-- Before creating cross-cell dependencies
-
-**Refresh tool:** `refresh_reactivity_context`
-
-- Returns success status
-- Returns context path to read result from
-- Writes latest reactivity graph to signals.md
-
-**Contents:** Signal dependencies between cells, widget signals, global variable signals, subscription relationships.
+**Format:** Cell metadata on separate lines (CELL_ID, CELL_INDEX, TYPE, STATUS), code between `CODE_START/CODE_END` markers, followed by a `REACTIVITY` subsection summarizing which reactive signals this cell defines along with the signals and cells it depends on (will trigger this cell to re-run).
 
 ### Refresh Strategy
 
@@ -795,7 +777,6 @@ The notebook state is available in two context files. These files are initialize
 **Refresh selectively:**
 
 - Refresh cells before inspecting/modifying notebook structure
-- Refresh reactivity when working with signals
 
 ### Creating Custom Files
 
@@ -1223,7 +1204,7 @@ Which result would you like to proceed with?""",
 10. **Plots MUST render via `w_plot`** - Every figure requires the plot widget
 11. **Transformation cells MUST be self-contained** - Include all imports, definitions, and variable creation
 12. **Assay platform documentation MUST be read immediately upon identification and followed EXACTLY STEP BY STEP with ZERO deviation** - These workflows are authoritative and inflexible. Every action must be verified against the current step. Manual alternatives are forbidden when workflows are specified.
-13. **Refresh context files when needed** - Call refresh_cells_context or refresh_reactivity_context when you need current state (e.g., after cell executions, before verifying variables exist) and use the context_path returned by the tool to read the result using `read_file` tool.
+13. **Refresh context files when needed** - Call `refresh_cells_context` whenever you need the latest cell layout or reactivity summary (e.g., after cell executions, before verifying variables exist) and use the `context_path` returned by the tool to read the result using `read_file`.
 14. **Widget keys cannot be assumed** - If you are creating widget(s) and need the widget key(s), call refresh_cells_context after the cell with the widget(s) has run
 
 ## NEVER Do

runtime/mount/entrypoint.py

@@ -434,7 +434,7 @@ async def handle_kernel_messages(conn_k: SocketIo, auth: str) -> None:
                         "type": "agent_action_response",
                         "tx_id": tx_id,
                         "status": "success",
-                        "summary": msg.get("summary", "")
+                        "cell_reactivity": msg.get("cell_reactivity", {}),
                     })
                     continue
 
@@ -481,7 +481,9 @@ async def handle_agent_messages(conn_a: SocketIo) -> None:
     while True:
         msg = await conn_a.recv()
         msg_type = msg.get("type", "unknown")
-        print(f"[entrypoint] Agent > {msg_type}")
+
+        if msg_type != "agent_stream_delta":
+            print(f"[entrypoint] Agent > {msg_type}")
 
         if msg_type == "agent_action" and msg.get("action") == "request_reactivity_summary":
             if k_proc.conn_k is not None: