From 3d80f494aea098ae37481e07aa42a4a34657ef5f Mon Sep 17 00:00:00 2001
From: gtarpenning <griffin.tarpenning@wandb.com>
Date: Wed, 30 Oct 2024 14:13:04 -0700
Subject: [PATCH] object viewer support for deleted refs

---
 .../Home/Browse2/SmallRef.tsx                 | 13 +++++++++-
 .../Browse3/pages/CallPage/ObjectViewer.tsx   | 25 +++++++++++++++++--
 .../clickhouse_trace_server_batched.py        |  2 --
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse2/SmallRef.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse2/SmallRef.tsx
index 3ac3a170bcc..0e842dea81d 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse2/SmallRef.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse2/SmallRef.tsx
@@ -127,7 +127,9 @@ export const SmallRef: FC<{
   const objectVersion = useObjectVersion(objVersionKey);
   const opVersion = useOpVersion(opVersionKey);
 
-  const isDeleted = objectVersion?.error || opVersion?.error;
+  const isDeleted =
+    isObjDeleteError(objectVersion?.error) ||
+    isObjDeleteError(opVersion?.error);
 
   const versionIndex =
     objectVersion.result?.versionIndex ?? opVersion.result?.versionIndex;
@@ -212,3 +214,12 @@ export const parseRefMaybe = (s: string): ObjectRef | null => {
     return null;
   }
 };
+
+export const isObjDeleteError = (e: any): boolean => {
+  if (e == null) {
+    return false;
+  }
+  const errorStr = String(e);
+  const regex = /Obj .* was deleted at .*/;
+  return regex.test(errorStr);
+};
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx
index a21e47e3917..20c8c7e2309 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx
@@ -7,6 +7,7 @@ import {
   GridRowId,
 } from '@mui/x-data-grid-pro';
 import {Button} from '@wandb/weave/components/Button';
+import {parseRef} from '@wandb/weave/react';
 import _ from 'lodash';
 import React, {
   Dispatch,
@@ -20,7 +21,7 @@ import React, {
 
 import {LoadingDots} from '../../../../../LoadingDots';
 import {Browse2OpDefCode} from '../../../Browse2/Browse2OpDefCode';
-import {parseRefMaybe} from '../../../Browse2/SmallRef';
+import {objectRefDisplayName, parseRefMaybe} from '../../../Browse2/SmallRef';
 import {isWeaveRef} from '../../filters/common';
 import {StyledDataGrid} from '../../StyledDataGrid';
 import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types';
@@ -151,10 +152,17 @@ export const ObjectViewer = ({
 
     const refValues: RefValues = {};
     for (const [r, v] of _.zip(refs, resolvedRefData)) {
-      if (!r || !v) {
+      if (!r) {
         // Shouldn't be possible
         continue;
       }
+      if (!v) {
+        // Value for ref not found, probably deleted
+        refValues[r] = {
+          _weave_is_deleted_ref: objectRefDisplayName(parseRef(r)).label,
+        };
+        continue;
+      }
       let val = r;
       if (v == null) {
         console.error('Error resolving ref', r);
@@ -394,6 +402,19 @@ export const ObjectViewer = ({
           return null;
         }
 
+        // Hack to show the object name with a strikethrough if deleted
+        if (params.row.value?._weave_is_deleted_ref) {
+          return (
+            <Box
+              sx={{
+                textOverflow: 'ellipsis',
+                textDecoration: 'line-through',
+              }}>
+              {params.row.value?._weave_is_deleted_ref}
+            </Box>
+          );
+        }
+
         return (
           <ObjectViewerGroupingCell
             {...params}
diff --git a/weave/trace_server/clickhouse_trace_server_batched.py b/weave/trace_server/clickhouse_trace_server_batched.py
index 9ba56d368e6..5a51e4e3ed2 100644
--- a/weave/trace_server/clickhouse_trace_server_batched.py
+++ b/weave/trace_server/clickhouse_trace_server_batched.py
@@ -653,8 +653,6 @@ def _obj_read(
         if len(objs) == 0:
             raise NotFoundError(f"Obj {object_id}:{digest} not found")
 
-        print(objs)
-
         if objs[0].deleted_at is not None:
             raise ObjectDeletedError(
                 f"Obj {object_id}:v{objs[0].version_index} was deleted at {objs[0].deleted_at}"