File size: 4,490 Bytes
b435029
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import { createStudioRender } from '../../domain/factories'
import type { StudioFileAttachment, StudioToolDefinition, StudioToolResult } from '../../domain/types'
import type { StudioRuntimeBackedToolContext } from '../../runtime/tools/tool-runtime-context'
import { isStudioRunCancelledError } from '../../runtime/execution/run-cancellation'
import {
  createUnconfiguredPlotRenderPort,
  type PlotRenderExecution,
  type PlotRenderPort
} from '../plot-render-port'
import { plotRenderToolParameters } from '../../tools/tool-parameters'
import { publishStudioRenderUpdated } from '../../render/render-events'
import { toWorkspaceRelativePath } from '../../tools/workspace-paths'

interface PlotRenderToolInput {
  concept: string
  code: string
}

export function createPlotStudioRenderTool(
  renderPort: PlotRenderPort = createUnconfiguredPlotRenderPort()
): StudioToolDefinition<PlotRenderToolInput> {
  return {
    name: 'render',
    parameters: plotRenderToolParameters,
    description: 'Execute matplotlib code and persist static plot outputs for preview.',
    allowedAgents: ['builder'],
    allowedStudioKinds: ['plot'],
    execute: async (input, context) => executePlotRenderTool(input, context as StudioRuntimeBackedToolContext, renderPort)
  }
}

async function executePlotRenderTool(
  input: PlotRenderToolInput,
  context: StudioRuntimeBackedToolContext,
  renderPort: PlotRenderPort
): Promise<StudioToolResult> {
  if (!input.concept?.trim() || !input.code?.trim()) {
    throw new Error('Render tool requires non-empty "concept" and "code"')
  }
  if (!context.renderStore) {
    throw new Error('Render store is unavailable')
  }

  const title = `Plot render: ${input.concept.slice(0, 80)}`
  const render = createStudioRender({
    ownerId: context.session.ownerId,
    sessionId: context.session.id,
    runId: context.run.id,
    kind: 'plot',
    title,
    concept: input.concept,
    outputMode: 'image',
    quality: 'medium',
    status: 'running',
    metadata: {
      studioKind: 'plot',
      outputMode: 'image'
    },
  })
  render.metadata = { ...(render.metadata ?? {}), renderId: render.id }
  const createdRender = await context.renderStore.create(render)
  publishStudioRenderUpdated(context.eventBus, createdRender)

  context.setToolMetadata?.({
    title,
    metadata: {
      renderId: render.id,
      studioKind: 'plot'
    }
  })

  try {
    const execution = await renderPort.execute({
      workspaceDirectory: context.session.directory,
      renderId: render.id,
      code: input.code,
      signal: context.abortSignal,
    })
    const attachments = buildAttachments(execution.imageDataUris, execution.imagePaths)
    const completed = await context.renderStore.update(context.session.ownerId, render.id, {
      status: 'completed',
      attachments,
      metadata: {
        ...(render.metadata ?? {}),
        imageCount: execution.imageDataUris.length,
        scriptPath: toWorkspaceRelativePath(context.session.directory, execution.scriptPath).replace(/\\/g, '/'),
        imagePaths: execution.imagePaths.map((imagePath) => toWorkspaceRelativePath(context.session.directory, imagePath).replace(/\\/g, '/')),
      },
    })
    if (!completed) {
      throw new Error('Render record disappeared before completion')
    }
    publishStudioRenderUpdated(context.eventBus, completed)

    return {
      title,
      output: `plot_render_id: ${render.id}`,
      attachments,
      metadata: {
        renderId: render.id,
        imageCount: execution.imageDataUris.length,
        scriptPath: completed.metadata?.scriptPath,
        imagePaths: completed.metadata?.imagePaths,
      },
    }
  } catch (error) {
    const failedRender = await context.renderStore.update(context.session.ownerId, render.id, {
      status: isStudioRunCancelledError(error) ? 'cancelled' : 'failed',
      error: error instanceof Error ? error.message : String(error),
    })
    if (failedRender) {
      publishStudioRenderUpdated(context.eventBus, failedRender)
    }
    throw error
  }
}

function buildAttachments(imageDataUris: string[], imagePaths: string[]): StudioFileAttachment[] {
  return imageDataUris.map((path, index) => ({
    kind: 'file',
    path,
    name: fileNameFromPath(imagePaths[index]) || `plot_${index + 1}.png`,
    mimeType: 'image/png'
  }))
}

function fileNameFromPath(path?: string): string {
  if (!path) {
    return ''
  }

  const parts = path.split(/[\\/]/)
  return parts[parts.length - 1] || path
}