dify/e2e/features/step-definitions/agent-v2/tools.steps.ts

import type { AgentSoulDifyToolConfig } from '@dify/contracts/api/console/agent/types.gen'
import type { DifyWorld } from '../../support/world'
import { Given, Then, When } from '@cucumber/cucumber'
import { expect } from '@playwright/test'
import { sendAgentServiceApiChatMessage } from '../../agent-v2/support/access-point'
import { createConfiguredTestAgent } from '../../agent-v2/support/agent'
import {
  agentBuilderExpectedTokens,
  agentBuilderPreseededResources,
} from '../../agent-v2/support/agent-builder-resources'
import {
  createAgentSoulConfigWithDifyTool,
  createAgentSoulConfigWithModel,
  normalAgentSoulConfig,
} from '../../agent-v2/support/agent-soul'
import { asArray, asRecord, asString } from '../../agent-v2/support/fixtures/common'
import { hasToolEntry } from '../../agent-v2/support/fixtures/tools'
import { SERVICE_API_RUNTIME_STEP_TIMEOUT_MS } from '../../agent-v2/support/service-api-sse'
import { getPreseededToolContract } from '../../agent-v2/support/tools'
import { expectProviderToolActionVisible, getCurrentAgentId } from './configure-helpers'

const getToolsSection = (world: DifyWorld) => world.getPage().getByRole('region', { name: 'Tools' })

const getToolSelectorSearch = (world: DifyWorld) =>
  world.getPage().getByRole('searchbox', { name: 'Search integrations...' })

const jsonReplaceRuntimePrompt = [
  'You are a Dify Agent E2E JSON tool verifier.',
  'When the user asks to verify JSON Replace, call the JSON Replace tool exactly once.',
  `Use content {"marker":"${agentBuilderExpectedTokens.jsonToolBefore}","nested":{"status":"keep"}}.`,
  'Use query $.marker and replace_value E2E_AFTER.',
  'After the tool returns, answer with the resulting JSON and include E2E_AFTER.',
].join(' ')

const expectJsonReplaceToolDraft = async (world: DifyWorld) => {
  const agentId = getCurrentAgentId(world)
  const tool = getPreseededToolContract(world, agentBuilderPreseededResources.jsonReplaceTool)

  await expect
    .poll(
      async () => {
        const draft = await world
          .getConsoleClient()
          .agent.byAgentId.composer.get({ params: { agent_id: agentId } })
        const tools = asArray(asRecord(draft.agent_soul?.tools).dify_tools)

        return hasToolEntry(tools, tool)
      },
      { timeout: 30_000 },
    )
    .toBe(true)
}

const jsonReplaceToolConfig = (world: DifyWorld): AgentSoulDifyToolConfig => {
  const tool = getPreseededToolContract(world, agentBuilderPreseededResources.jsonReplaceTool)

  return {
    credential_type: 'unauthorized',
    enabled: true,
    provider: tool.providerDisplayName,
    provider_id: tool.providerName,
    provider_type: 'builtin',
    runtime_parameters: {
      ensure_ascii: true,
      replace_model: 'value',
      value_decode: false,
    },
    tool_name: tool.toolName,
  }
}

const getServiceApiSseEvents = (body: unknown) => asArray(asRecord(body).events).map(asRecord)

const getServiceApiEventData = (event: Record<string, unknown>) => asRecord(event.data)

const getServiceApiEventName = (event: Record<string, unknown>) => {
  const data = getServiceApiEventData(event)

  return asString(data.event) || asString(event.event)
}

const summarizeServiceApiRuntimeEvents = (body: unknown) =>
  getServiceApiSseEvents(body).map((event, index) => {
    const data = getServiceApiEventData(event)
    const observation = asString(data.observation)
    const answer = asString(data.answer)

    return {
      index,
      event: getServiceApiEventName(event),
      tool: asString(data.tool),
      hasToolInput: Boolean(asString(data.tool_input)),
      hasObservation: Boolean(observation),
      observation: observation.slice(0, 180),
      answer: answer.slice(0, 180),
    }
  })

const findJsonReplaceRuntimeThought = (body: unknown) =>
  getServiceApiSseEvents(body).find((event) => {
    const data = getServiceApiEventData(event)
    const toolInput = asString(data.tool_input)
    const observation = asString(data.observation)

    return (
      getServiceApiEventName(event) === 'agent_thought' &&
      asString(data.tool) === 'json_replace' &&
      toolInput.includes(agentBuilderExpectedTokens.jsonToolBefore) &&
      toolInput.includes('$.marker') &&
      observation.includes(agentBuilderExpectedTokens.jsonToolAfter)
    )
  })

Given(
  'a runnable Agent v2 test agent with the JSON Replace tool has been created via API',
  async function (this: DifyWorld) {
    if (!this.agentBuilder.fixtures.stableModel)
      throw new Error('Create a JSON Replace runtime Agent after stable model fixture setup.')

    const agent = await createConfiguredTestAgent(this.getConsoleClient(), {
      agentSoul: createAgentSoulConfigWithDifyTool(
        createAgentSoulConfigWithModel(
          {
            ...normalAgentSoulConfig,
            prompt: {
              system_prompt: jsonReplaceRuntimePrompt,
            },
          },
          this.agentBuilder.fixtures.stableModel,
        ),
        jsonReplaceToolConfig(this),
      ),
    })

    this.createdAgentIds.push(agent.id)
    this.lastCreatedAgentName = agent.name
    this.lastCreatedAgentRole = agent.role ?? undefined
  },
)

When(
  'I add the Agent Builder JSON Replace tool from the Tools selector',
  async function (this: DifyWorld) {
    const page = this.getPage()
    const toolsSection = getToolsSection(this)

    await expect(toolsSection).toBeVisible({ timeout: 30_000 })
    await toolsSection.getByRole('button', { name: 'Add tool' }).click()

    const search = getToolSelectorSearch(this)
    await expect(search).toBeVisible()
    await search.fill('JSON Replace')

    await page.getByRole('button', { exact: true, name: 'JSON Replace' }).click()
    await expectProviderToolActionVisible(
      toolsSection,
      agentBuilderPreseededResources.jsonReplaceTool,
    )
  },
)

Then(
  'the Agent v2 JSON Replace tool should be saved in the Agent v2 draft',
  async function (this: DifyWorld) {
    await expectJsonReplaceToolDraft(this)
  },
)

Then(
  'I should see the Agent v2 JSON Replace tool in the Tools section',
  async function (this: DifyWorld) {
    await expectProviderToolActionVisible(
      getToolsSection(this),
      agentBuilderPreseededResources.jsonReplaceTool,
    )
    await expectJsonReplaceToolDraft(this)
  },
)

When(
  'I send the Agent v2 Backend service API JSON Replace request',
  { timeout: SERVICE_API_RUNTIME_STEP_TIMEOUT_MS },
  async function (this: DifyWorld) {
    const serviceApiBaseURL = this.agentBuilder.accessPoint.serviceApiBaseURL
    const apiKey = this.agentBuilder.accessPoint.generatedApiKey
    if (!serviceApiBaseURL)
      throw new Error('No Agent v2 service API endpoint found. Enable Backend service API first.')
    if (!apiKey)
      throw new Error('No Agent v2 API key found. Create a Backend service API key first.')

    this.agentBuilder.accessPoint.serviceApiResponse = await sendAgentServiceApiChatMessage({
      apiKey,
      query: [
        'Verify JSON Replace now.',
        `Replace ${agentBuilderExpectedTokens.jsonToolBefore} with ${agentBuilderExpectedTokens.jsonToolAfter}.`,
        `Return the JSON result and include ${agentBuilderExpectedTokens.jsonToolAfter}.`,
      ].join(' '),
      serviceApiBaseURL,
    })
  },
)

Then(
  'the Agent v2 Backend service API response should include the JSON Replace E2E marker',
  async function (this: DifyWorld) {
    const response = this.agentBuilder.accessPoint.serviceApiResponse
    if (!response) throw new Error('No Agent v2 Backend service API response was recorded.')
    if (!response.ok)
      throw new Error(
        `Agent v2 Backend service API JSON Replace request failed with ${response.status}: ${JSON.stringify(response.body)}`,
      )

    const jsonReplaceThought = findJsonReplaceRuntimeThought(response.body)
    if (!jsonReplaceThought) {
      throw new Error(
        [
          'Agent v2 Backend service API did not emit a JSON Replace agent_thought with matching tool input and observation.',
          `Received events: ${JSON.stringify(summarizeServiceApiRuntimeEvents(response.body))}`,
        ].join('\n'),
      )
    }

    const thought = getServiceApiEventData(jsonReplaceThought)
    expect(asString(thought.tool_input)).toContain(agentBuilderExpectedTokens.jsonToolBefore)
    expect(asString(thought.tool_input)).toContain('$.marker')
    expect(asString(thought.observation)).toContain(agentBuilderExpectedTokens.jsonToolAfter)
    expect(asString(asRecord(response.body).answer)).toContain(
      agentBuilderExpectedTokens.jsonToolAfter,
    )
  },
)