Use Port of Context (pctx) for code mode (#6765)

Signed-off-by: Elias Posen <elias@posen.ch>
Signed-off-by: Adrian Cole <adrian@tetrate.io>
Co-authored-by: Adrian Cole <adrian@tetrate.io>

ui/desktop/src/components/ToolCallWithResponse.tsx

@@ -545,7 +545,7 @@ function ToolCallView({
       case 'computer_control':
         return `poking around...`;
 
-      case 'execute_code': {
+      case 'execute': {
         const toolGraph = args.tool_graph as unknown as ToolGraphNode[] | undefined;
         if (toolGraph && Array.isArray(toolGraph) && toolGraph.length > 0) {
           if (toolGraph.length === 1) {
@@ -639,19 +639,16 @@ function ToolCallView({
       }
     >
       {(() => {
-        const toolName = toolCall.name.substring(toolCall.name.lastIndexOf('__') + 2);
-        const toolGraph = toolCall.arguments?.tool_graph as unknown as ToolGraphNode[] | undefined;
         const code = toolCall.arguments?.code as unknown as string | undefined;
-        const hasToolGraph =
-          toolName === 'execute_code' &&
-          toolGraph &&
-          Array.isArray(toolGraph) &&
-          toolGraph.length > 0;
+        const toolGraph = toolCall.arguments?.tool_graph as unknown as ToolGraphNode[] | undefined;
 
-        if (hasToolGraph) {
+        if (
+          toolCall.name === 'code_execution__execute' &&
+          (typeof code === 'string' || Array.isArray(toolGraph))
+        ) {
           return (
             <div className="border-t border-borderSubtle">
-              <ToolGraphView toolGraph={toolGraph} code={code} />
+              <CodeModeView toolGraph={toolGraph} code={code} />
             </div>
           );
         }
@@ -692,7 +689,7 @@ function ToolCallView({
         <>
           {toolResults.map((result, index) => (
             <div key={index} className={cn('border-t border-borderSubtle')}>
-              <ToolResultView result={result} isStartExpanded={false} />
+              <ToolResultView toolCall={toolCall} result={result} isStartExpanded={false} />
             </div>
           ))}
         </>
@@ -724,18 +721,19 @@ function ToolDetailsView({ toolCall, isStartExpanded }: ToolDetailsViewProps) {
   );
 }
 
-interface ToolGraphViewProps {
-  toolGraph: ToolGraphNode[];
+interface CodeModeViewProps {
+  toolGraph?: ToolGraphNode[];
   code?: string;
 }
 
-function ToolGraphView({ toolGraph, code }: ToolGraphViewProps) {
+function CodeModeView({ toolGraph, code }: CodeModeViewProps) {
   const renderGraph = () => {
-    if (toolGraph.length === 0) return null;
+    const graph = toolGraph ?? [];
+    if (graph.length === 0) return null;
 
     const lines: string[] = [];
 
-    toolGraph.forEach((node, index) => {
+    graph.forEach((node, index) => {
       const deps =
         node.depends_on.length > 0 ? ` (uses ${node.depends_on.map((d) => d + 1).join(', ')})` : '';
       lines.push(`${index + 1}. ${node.tool}: ${node.description}${deps}`);
@@ -746,16 +744,19 @@ function ToolGraphView({ toolGraph, code }: ToolGraphViewProps) {
 
   return (
     <div className="px-4 py-2">
-      <pre className="font-mono text-xs text-textSubtle whitespace-pre-wrap">{renderGraph()}</pre>
+      {toolGraph && (
+        <pre className="font-mono text-xs text-textSubtle whitespace-pre-wrap">{renderGraph()}</pre>
+      )}
       {code && (
         <div className="border-t border-borderSubtle -mx-4 mt-2">
           <ToolCallExpandable
             label={<span className="pl-4 font-sans text-sm">Code</span>}
             isStartExpanded={false}
           >
-            <pre className="font-mono text-xs text-textSubtle whitespace-pre-wrap overflow-x-auto px-4 py-2">
-              {code}
-            </pre>
+            <MarkdownContent
+              content={'```typescript\n' + code + '\n```'}
+              className="whitespace-pre-wrap max-w-full overflow-x-auto"
+            />
           </ToolCallExpandable>
         </div>
       )}
@@ -764,11 +765,15 @@ function ToolGraphView({ toolGraph, code }: ToolGraphViewProps) {
 }
 
 interface ToolResultViewProps {
+  toolCall: {
+    name: string;
+    arguments: Record<string, unknown>;
+  };
   result: Content;
   isStartExpanded: boolean;
 }
 
-function ToolResultView({ result, isStartExpanded }: ToolResultViewProps) {
+function ToolResultView({ toolCall, result, isStartExpanded }: ToolResultViewProps) {
   const hasText = (c: Content): c is Content & { text: string } =>
     'text' in c && typeof (c as Record<string, unknown>).text === 'string';
 
@@ -780,6 +785,18 @@ function ToolResultView({ result, isStartExpanded }: ToolResultViewProps) {
 
   const hasResource = (c: Content): c is Content & { resource: unknown } => 'resource' in c;
 
+  const wrapMarkdown = (text: string): string => {
+    if (
+      ['code_execution__list_functions', 'code_execution__get_function_details'].includes(
+        toolCall.name
+      )
+    ) {
+      return '```typescript\n' + text + '\n```';
+    } else {
+      return text;
+    }
+  };
+
   return (
     <ToolCallExpandable
       label={<span className="pl-4 py-1 font-sans text-sm">Output</span>}
@@ -788,7 +805,7 @@ function ToolResultView({ result, isStartExpanded }: ToolResultViewProps) {
       <div className="pl-4 pr-4 py-4">
         {hasText(result) && (
           <MarkdownContent
-            content={result.text}
+            content={wrapMarkdown(result.text)}
             className="whitespace-pre-wrap max-w-full overflow-x-auto"
           />
         )}