Add __enhanced_stack_trace query to workers

temporalio/worker/_workflow_instance.py

@@ -69,6 +69,36 @@
     WorkflowOutboundInterceptor,
 )
 
+from temporalio.service import __version__
+
+@dataclass
+class SDKInfo:
+    name    : str
+    version : str
+
+@dataclass
+class FileSlice:
+    content    : str
+    lineOffset : int
+
+@dataclass
+class FileLocation:
+    filePath     : str
+    line         : Optional[int] = -1
+    column       : Optional[int] = -1
+    functionName : Optional[str] = None
+
+@dataclass
+class StackTrace:
+    locations : list[FileLocation]
+
+@dataclass
+class EnhancedStackTrace:
+    sdk     : SDKInfo
+    sources : dict[str, FileSlice]
+    stacks  : list[StackTrace]
+
+
 logger = logging.getLogger(__name__)
 
 # Set to true to log all cases where we're ignoring things during delete
@@ -250,6 +280,14 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
             ret_type=str,
         )
 
+        self._queries["__enhanced_stack_trace"] = temporalio.workflow._QueryDefinition(
+            name="__enhanced_stack_trace",
+            fn=self._enhanced_stack_trace,
+            is_method=False,
+            arg_types=[],
+            ret_type=EnhancedStackTrace,
+        )
+
         # Maintain buffered signals for later-added dynamic handlers
         self._buffered_signals: Dict[
             str, List[temporalio.bridge.proto.workflow_activation.SignalWorkflow]
@@ -1790,6 +1828,38 @@ def _stack_trace(self) -> str:
                 + "\n".join(traceback.format_list(frames))
             )
         return "\n\n".join(stacks)
+
+    def _enhanced_stack_trace(self) -> EnhancedStackTrace:
+        sdk = SDKInfo("sdk-python", __version__)
+
+        sources = dict()
+        stacks = []
+
+        for task in list(self._tasks):
+            for frame in task.get_stack():
+                filename = frame.f_code.co_filename
+                line_number = frame.f_lineno
+                func_name = frame.f_code.co_name
+
+                try:
+                    source = inspect.getsourcelines(frame)
+                    code = ''.join(source[0])
+                    line_number = int(source[1])
+                except OSError as ose:
+                    code = "Cannot access code.\n---\n%s" % ose.strerror
+                    # TODO possibly include sentinel/property for success of src scrape? work out with ui
+                except Exception:
+                    code = "Generic Error.\n\n%s" % traceback.format_exc()
+
+                file_slice = FileSlice(code, line_number)
+                file_location = FileLocation(filename, line = line_number, functionName = func_name)
+
+                sources["%s %d" % (filename, line_number)] = file_slice
+                stacks.append(file_location)
+
+
+        est = EnhancedStackTrace(sdk, sources, stacks)
+        return est
 
     #### asyncio.AbstractEventLoop function impls ####
     # These are in the order defined in CPython's impl of the base class. Many

tests/worker/test_workflow.py

@@ -92,7 +92,7 @@
     Runtime,
     TelemetryConfig,
 )
-from temporalio.service import RPCError, RPCStatusCode
+from temporalio.service import RPCError, RPCStatusCode, __version__
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import (
     UnsandboxedWorkflowRunner,
@@ -494,7 +494,7 @@ async def test_workflow_signal_and_query_errors(client: Client):
             await handle.query("non-existent query")
         assert str(rpc_err.value) == (
             "Query handler for 'non-existent query' expected but not found,"
-            " known queries: [__stack_trace bad_query other_query]"
+            " known queries: [__enhanced_stack_trace __stack_trace bad_query other_query]"
         )
 
 
@@ -2097,6 +2097,30 @@ async def status() -> str:
         # TODO(cretz): Do more specific checks once we clean up traces
         assert "never_completing_coroutine" in trace
 
+async def test_workflow_enhanced_stack_trace(client: Client):
+    async with new_worker(
+        client, StackTraceWorkflow, LongSleepWorkflow, activities=[wait_cancel]
+    ) as worker:
+        handle = await client.start_workflow(
+            StackTraceWorkflow.run,
+            id=f"workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+
+        # Wait until waiting
+        async def status() -> str:
+            return await handle.query(StackTraceWorkflow.status)
+
+        await assert_eq_eventually("waiting", status)
+
+        # Send stack trace query
+        trace = await handle.query("__enhanced_stack_trace")
+
+        assert "never_completing_coroutine" in [ stack['functionName'] for stack in trace['stacks'] ]
+        # first line of never_completing_coroutine
+        assert 'self._status = "waiting"' in str(trace['sources'])
+        assert trace['sdk']['version'] == __version__
+
 
 @dataclass
 class MyDataClass: