[
  {
    "path": ".github/.codedev.yml",
    "content": "coverage:\n  status:\n    project: #add everything under here, more options at https://docs.codecov.com/docs/commit-status\n      default: # default is the status check's name, not default settings\n        target: auto #default\n        threshold: 1% #allow coverage to drop by 1%\n        base: auto\n        if_ci_failed: error #success, failure, error, ignore\n    patch:\n      default:\n        target: 82% #default\n        threshold: 1% #allow coverage to drop by 1%\n        base: auto\n        if_ci_failed: error #success, failure, error, ignore\n\ncomment: #this is a top-level key\n  layout: \" diff, flags, files\"\n  behavior: default\n  require_changes: false  # if true: only post the comment if coverage changes\n  require_base: false        # [true :: must have a base report to post]\n  require_head: true       # [true :: must have a head report to post]\n  hide_project_coverage: false # [true :: only show coverage on the git diff aka patch coverage]\n\n# sample regex patterns\nignore:\n  - \"tests\"\n  - \"examples/\"\n  - \"mock/\"\n  - \"callbacks/interface.go\"\n  - \"utils/safe\"\n  - \"components/tool/utils/create_options.go\""
  },
  {
    "path": ".github/.commit-rules.json",
    "content": "{\n  \"allowedTypes\": [\n    \"feat\",\n    \"fix\",\n    \"docs\",\n    \"style\",\n    \"refactor\",\n    \"perf\",\n    \"test\",\n    \"build\",\n    \"ci\",\n    \"chore\",\n    \"revert\"\n  ],\n  \"allowedScopes\": [\n    \"adk\",\n    \"adk/filesystem\",\n    \"callbacks\",\n    \"components\",\n    \"compose\",\n    \"deep\",\n    \"dynamictool\",\n    \"filesystem\",\n    \"flow\",\n    \"internal\",\n    \"middlewares\",\n    \"planexecute\",\n    \"plantask\",\n    \"prebuilt\",\n    \"reduction\",\n    \"schema\",\n    \"skill\",\n    \"summarization\",\n    \"supervisor\",\n    \"toolsearch\",\n    \"utils\",\n    \"docs\",\n    \"ci\",\n    \"serialization\"\n  ]\n}\n"
  },
  {
    "path": ".github/ISSUE_TEMPLATE/bug_report.md",
    "content": "---\nname: Bug report\nabout: Create a report to help us improve\ntitle: ''\nlabels: ''\nassignees: ''\n\n---\n\n**Describe the bug**\n\nA clear and concise description of what the bug is.\n\n**To Reproduce**\n\nSteps to reproduce the behavior:\n1. Go to '...'\n2. Click on '....'\n3. Scroll down to '....'\n4. See error\n\n**Expected behavior**\n\nA clear and concise description of what you expected to happen.\n\n**Screenshots**\n\nIf applicable, add screenshots to help explain your problem.\n\n**Version:**\n\nPlease provide the version of {project_name} you are using.\n\n**Environment:**\n\nThe output of `go env`.\n\n**Additional context**\n\nAdd any other context about the problem here.\n"
  },
  {
    "path": ".github/ISSUE_TEMPLATE/feature_request.md",
    "content": "---\nname: Feature request\nabout: Suggest an idea for this project\ntitle: ''\nlabels: ''\nassignees: ''\n\n---\n\n**Is your feature request related to a problem? Please describe.**\n\nA clear and concise description of what the problem is. Ex. I'm always frustrated when [...]\n\n**Describe the solution you'd like**\n\nA clear and concise description of what you want to happen.\n\n**Describe alternatives you've considered**\n\nA clear and concise description of any alternative solutions or features you've considered.\n\n**Additional context**\n\nAdd any other context or screenshots about the feature request here.\n"
  },
  {
    "path": ".github/PULL_REQUEST_TEMPLATE.md",
    "content": "#### What type of PR is this?\n<!--\nAdd one of the following kinds:\n\nbuild: Changes that affect the build system or external dependencies (example scopes: gulp, broccoli, npm)\nci: Changes to our CI configuration files and scripts (example scopes: Travis, Circle, BrowserStack, SauceLabs)\ndocs: Documentation only changes\nfeat: A new feature\noptimize: A new optimization\nfix: A bug fix\nperf: A code change that improves performance\nrefactor: A code change that neither fixes a bug nor adds a feature\nstyle: Changes that do not affect the meaning of the code (white space, formatting, missing semi-colons, etc)\ntest: Adding missing tests or correcting existing tests\nchore: Changes to the build process or auxiliary tools and libraries such as documentation generation\n-->\n\n#### Check the PR title.\n<!--\nThe description of the title will be attached in Release Notes, \nso please describe it from user-oriented, what this PR does / why we need it.\nPlease check your PR title with the below requirements:\n-->\n- [ ] This PR title match the format: \\<type\\>(optional scope): \\<description\\>\n- [ ] The description of this PR title is user-oriented and clear enough for others to understand.\n- [ ] Attach the PR updating the user documentation if the current PR requires user awareness at the usage level. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)\n\n\n#### (Optional) Translate the PR title into Chinese.\n\n\n#### (Optional) More detailed description for this PR(en: English/zh: Chinese).\n<!--\nProvide more detailed info for review(e.g., it's recommended to provide perf data if this is a perf type PR).\n-->\nen:\nzh(optional): \n\n\n#### (Optional) Which issue(s) this PR fixes:\n<!--\nAutomatically closes linked issue when PR is merged.\nEg: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.\n-->\n\n#### (optional) The PR that updates user documentation:\n<!--\nIf the current PR requires user awareness at the usage level, please submit a PR to update user docs. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)\n-->\n"
  },
  {
    "path": ".github/workflows/pr-check.yml",
    "content": "name: Pull Request Check\n\non: [ pull_request ]\n\njobs:\n  compliant:\n    runs-on: ubuntu-latest\n    steps:\n      - uses: actions/checkout@v4\n\n      - name: Check License Header\n        uses: apache/skywalking-eyes/header@v0.4.0\n        env:\n          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n\n      - name: Check Spell\n        uses: crate-ci/typos@v1.42.3\n\n  golangci-lint:\n    runs-on: ubuntu-latest\n    permissions:\n      contents: write\n      pull-requests: write\n      repository-projects: write\n    steps:\n      - uses: actions/checkout@v4\n      - name: Set up Go\n        uses: actions/setup-go@v5\n        with:\n          go-version: 1.18\n          # for self-hosted, the cache path is shared across projects,\n          # and it works well without the cache of GitHub actions\n          # Enable it if we're going to use GitHub only\n          cache: true\n\n      - name: Golang CI Lint\n        # https://golangci-lint.run/\n        uses: golangci/golangci-lint-action@v9.2.0\n        with:\n          version: v2.8.0\n          args: --timeout 5m\n\n  commit-msg-check:\n    name: Commit Message Check\n    runs-on: ubuntu-latest\n    permissions:\n      contents: read\n      pull-requests: read\n    steps:\n      - uses: actions/checkout@v4\n      - name: Validate commit messages format and scope\n        uses: actions/github-script@v7\n        with:\n          github-token: ${{ secrets.GITHUB_TOKEN }}\n          script: |\n            const fs = require('fs');\n            const pr = context.payload.pull_request;\n            if (!pr) {\n              core.setFailed('This workflow must run on pull_request events.');\n              return;\n            }\n\n            let allowedTypes = [];\n            let allowedScopes = [];\n            try {\n              const raw = fs.readFileSync('.github/.commit-rules.json', 'utf8');\n              const cfg = JSON.parse(raw);\n              allowedTypes = Array.isArray(cfg.allowedTypes) ? cfg.allowedTypes : [];\n              allowedScopes = Array.isArray(cfg.allowedScopes) ? cfg.allowedScopes : [];\n            } catch (e) {\n              core.setFailed('Cannot read .github/.commit-rules.json: ' + e.message);\n              return;\n            }\n            if (!allowedTypes.length) {\n              core.setFailed('allowedTypes is empty in .github/.commit-rules.json');\n              return;\n            }\n\n            const { owner, repo } = context.repo;\n            const pull_number = pr.number;\n            const commits = await github.paginate(\n              github.rest.pulls.listCommits,\n              { owner, repo, pull_number, per_page: 100 }\n            );\n\n            let errors = [];\n\n            for (const c of commits) {\n              const sha = c.sha.slice(0, 7);\n              const subject = (c.commit.message || '').split('\\n')[0];\n              const m = subject.match(/^([a-z]+)(\\(([a-z0-9\\-\\/]+)\\))?:\\s(.+)$/);\n              if (!m) {\n                errors.push(`(${sha}) invalid format: \"${subject}\"`);\n                continue;\n              }\n\n              const type = m[1];\n              const scope = m[3]; // may be undefined\n              const desc = m[4];\n\n              if (!allowedTypes.includes(type)) {\n                errors.push(`(${sha}) invalid type \"${type}\"`);\n              }\n\n              if (!desc || !desc.trim()) {\n                errors.push(`(${sha}) description must be non-empty`);\n              }\n\n              if (scope) {\n                const topScope = scope.split('/')[0];\n                if (allowedScopes.length && !allowedScopes.includes(topScope)) {\n                  errors.push(`(${sha}) invalid scope \"${scope}\"`);\n                }\n              }\n            }\n\n            if (errors.length) {\n              core.setFailed('Commit message check failed:\\n' + errors.join('\\n'));\n            } else {\n              core.info('All commit messages conform to \"<type>(optional scope): <description>\" and scope rules.');\n            }\n\n  pr-title-check:\n    name: PR Title Check\n    runs-on: ubuntu-latest\n    permissions:\n      contents: read\n      pull-requests: read\n    steps:\n      - uses: actions/checkout@v4\n      - name: Read commit rules\n        id: rules\n        uses: actions/github-script@v7\n        with:\n          github-token: ${{ secrets.GITHUB_TOKEN }}\n          script: |\n            const fs = require('fs');\n            let cfg;\n            try {\n              const raw = fs.readFileSync('.github/.commit-rules.json', 'utf8');\n              cfg = JSON.parse(raw);\n            } catch (e) {\n              core.setFailed('Cannot read .github/.commit-rules.json: ' + e.message);\n              return;\n            }\n            const toMultiline = (list) => Array.isArray(list) ? list.join('\\n') : '';\n            core.setOutput('types', toMultiline(cfg.allowedTypes));\n            core.setOutput('scopes', toMultiline(cfg.allowedScopes));\n      - name: Validate PR title\n        uses: amannn/action-semantic-pull-request@v6.1.1\n        env:\n          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n        with:\n          types: ${{ steps.rules.outputs.types }}\n          scopes: ${{ steps.rules.outputs.scopes }}\n          requireScope: false\n"
  },
  {
    "path": ".github/workflows/tag-notification.yml",
    "content": "name: Tag Notification\n\non:\n  push:\n    tags:\n      - 'v*'\n\njobs:\n  notify:\n    runs-on: ubuntu-latest\n    steps:\n      - uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          fetch-tags: true\n\n      - name: Fetch tag info\n        run: |\n          git fetch --tags -f\n          \n      - name: Get tag info and send notification\n        run: |\n          # Get the tag name\n          TAG_NAME=\"${{ github.ref_name }}\"\n          echo \"Processing tag: $TAG_NAME\"\n          \n          # Get tag message\n          echo \"Getting tag message...\"\n          TAG_MESSAGE=$(git tag -l --format='%(contents)' \"$TAG_NAME\")\n          echo \"Tag message:\"\n          echo \"$TAG_MESSAGE\"\n          echo \"---\"\n          \n          # Create base content parts\n          HEADER=\"### 🏷️ Eino New Tag Created: \\`$TAG_NAME\\`\"\n          VERSION_INFO=\"📦 Version: \\`$TAG_NAME\\`\"\n          \n          # Prepare the message parts for jq\n          if [ ! -z \"$TAG_MESSAGE\" ]; then\n            # Pass all parts to jq and let it handle the formatting\n            jq -n \\\n              --arg header \"$HEADER\" \\\n              --arg version \"$VERSION_INFO\" \\\n              --arg notes \"$TAG_MESSAGE\" \\\n              --arg repo_url \"https://github.com/${{ github.repository }}/releases/tag/$TAG_NAME\" \\\n              '{\n                \"msg_type\": \"interactive\",\n                \"card\": {\n                  \"elements\": [\n                    {\n                      \"tag\": \"markdown\",\n                      \"content\": ($header + \"\\n\\n\" + $version + \"\\n\\n### 📝 Release Notes:\\n\" + $notes)\n                    },\n                    {\n                      \"tag\": \"action\",\n                      \"actions\": [\n                        {\n                          \"tag\": \"button\",\n                          \"text\": {\n                            \"tag\": \"plain_text\",\n                            \"content\": \"🔗 View Tag\"\n                          },\n                          \"url\": $repo_url,\n                          \"type\": \"default\"\n                        }\n                      ]\n                    }\n                  ],\n                  \"header\": {\n                    \"title\": {\n                      \"tag\": \"plain_text\",\n                      \"content\": \"🏷️ Eino New Tag Created\"\n                    }\n                  }\n                }\n              }' > webhook_payload.json\n          else\n            # Without release notes\n            jq -n \\\n              --arg header \"$HEADER\" \\\n              --arg version \"$VERSION_INFO\" \\\n              --arg repo_url \"https://github.com/${{ github.repository }}/releases/tag/$TAG_NAME\" \\\n              '{\n                \"msg_type\": \"interactive\",\n                \"card\": {\n                  \"elements\": [\n                    {\n                      \"tag\": \"markdown\",\n                      \"content\": ($header + \"\\n\\n\" + $version)\n                    },\n                    {\n                      \"tag\": \"action\",\n                      \"actions\": [\n                        {\n                          \"tag\": \"button\",\n                          \"text\": {\n                            \"tag\": \"plain_text\",\n                            \"content\": \"🔗 View Tag\"\n                          },\n                          \"url\": $repo_url,\n                          \"type\": \"default\"\n                        }\n                      ]\n                    }\n                  ],\n                  \"header\": {\n                    \"title\": {\n                      \"tag\": \"plain_text\",\n                      \"content\": \"🏷️ Eino New Tag Created\"\n                    }\n                  }\n                }\n              }' > webhook_payload.json\n          fi\n          \n          # Send webhook\n          curl -X POST \\\n               -H \"Content-Type: application/json\" \\\n               -d @webhook_payload.json \\\n               \"${{ secrets.FEISHU_WEBHOOK_URL }}\" "
  },
  {
    "path": ".github/workflows/tests.yml",
    "content": "name: Eino Tests\n\non:\n  pull_request:\n  push:\n    branches:\n      - main\n\nenv:\n  DEFAULT_GO_VERSION: \"1.18\"\n\njobs:\n  unit-test:\n    name: eino-unit-test\n    runs-on: ubuntu-latest\n    permissions:\n      contents: write\n      pull-requests: write\n      repository-projects: write\n    env:\n      COVERAGE_FILE: coverage.out\n      BREAKDOWN_FILE: main.breakdown\n      \n    steps:\n      - uses: actions/checkout@v4\n      - name: Set up Go\n        uses: actions/setup-go@v5\n        with:\n          go-version: ${{ env.DEFAULT_GO_VERSION }}\n          \n      - name: Exec Go Test\n        run: |\n          modules=`find . -name \"go.mod\" -exec dirname {} \\;`\n          echo $modules\n          list=\"\"\n          coverpkg=\"\"\n          if [[ ! -f \"go.work\" ]];then go work init;fi\n          for module in $modules; do go work use $module; list=$module\"/... \"$list; coverpkg=$module\"/...,\"$coverpkg; done\n          go work sync\n          go test -race -v -coverprofile=${{ env.COVERAGE_FILE }} -gcflags=\"all=-l -N\" -coverpkg=$coverpkg $list\n          \n      - name: Upload coverage to Codecov\n        uses: codecov/codecov-action@v5\n        with:\n          name: eino-unit-test\n          env_vars: GOLANG,EINO\n          files: ${{ env.COVERAGE_FILE }}\n          token: ${{ secrets.CODECOV_TOKEN }}\n          codecov_yml_path: ./github/.codecov.yml\n\n  benchmark-test:\n    runs-on: ubuntu-latest\n    permissions:\n      contents: write\n      pull-requests: write\n      repository-projects: write\n    steps:\n      - uses: actions/checkout@v4\n      - name: Set up Go\n        uses: actions/setup-go@v5\n        with:\n          go-version: ${{ env.DEFAULT_GO_VERSION }}\n\n      - name: Run Benchmark Tests\n        run: go test -bench=. -benchmem -run=none ./...\n\n  compatibility-test:\n    strategy:\n      matrix:\n        go: [ \"1.19\", \"1.20\", \"1.21\", \"1.22\", \"1.23\", \"1.24\" ]\n    runs-on: ubuntu-latest\n    permissions:\n      contents: write\n      pull-requests: write\n      repository-projects: write\n    steps:\n      - uses: actions/checkout@v4\n      - name: Set up Go\n        uses: actions/setup-go@v5\n        with:\n          go-version: ${{ matrix.go }}\n          cache: true\n\n      - name: Compatibility Test\n        run: |\n          # just basic unit test, no coverage report\n          go test -race ./...\n\n  api-compatibility:\n    name: api-compatibility-check\n    runs-on: ubuntu-latest\n    permissions:\n      contents: write\n      pull-requests: write\n      repository-projects: write\n    if: github.event_name == 'pull_request'\n    \n    steps:\n      - uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          \n      - name: Set up Go\n        uses: actions/setup-go@v5\n        with:\n          go-version: \"1.22\"\n          \n      - name: Install go-apidiff\n        run: go install github.com/joelanford/go-apidiff@v0.8.2\n          \n      - name: Check API compatibility\n        id: apidiff\n        run: |\n          BASE_SHA=${{ github.event.pull_request.base.sha }}\n          HEAD_SHA=${{ github.event.pull_request.head.sha }}\n          \n          echo \"Checking API compatibility between $BASE_SHA and $HEAD_SHA\"\n          \n          go mod tidy\n          \n          if ! DIFF_OUTPUT=$(go-apidiff $BASE_SHA $HEAD_SHA 2>&1); then\n            echo \"go-apidiff output: $DIFF_OUTPUT\"\n          fi\n\n          echo \"diff_output<<EOF\" >> $GITHUB_ENV\n          echo \"$DIFF_OUTPUT\" >> $GITHUB_ENV\n          echo \"EOF\" >> $GITHUB_ENV\n          \n          if echo \"$DIFF_OUTPUT\" | grep -q \"Incompatible changes:\"; then\n            echo \"has_breaking_changes=true\" >> $GITHUB_OUTPUT\n          else\n            echo \"has_breaking_changes=false\" >> $GITHUB_OUTPUT\n          fi\n\n      - name: Create Review Thread\n        if: steps.apidiff.outputs.has_breaking_changes == 'true'\n        continue-on-error: true\n        uses: actions/github-script@v7\n        with:\n          github-token: ${{ secrets.GITHUB_TOKEN }}\n          script: |\n            const reviewComments = await github.rest.pulls.listReviewComments({\n              owner: context.repo.owner,\n              repo: context.repo.repo,\n              pull_number: context.issue.number\n            });\n\n            const existingPackageComments = new Map();\n            \n            for (const comment of reviewComments.data) {\n              if (comment.body.includes('Breaking API Changes Detected')) {\n                const packageMatch = comment.body.match(/Package: `([^`]+)`/);\n                if (packageMatch) {\n                  const pkg = packageMatch[1];\n                  if (!existingPackageComments.has(pkg)) {\n                    existingPackageComments.set(pkg, new Set());\n                  }\n                  existingPackageComments.get(pkg).add(comment.path);\n                }\n              }\n            }\n            \n            const files = await github.rest.pulls.listFiles({\n              owner: context.repo.owner,\n              repo: context.repo.repo,\n              pull_number: context.issue.number\n            });\n            \n            const diffOutput = process.env.diff_output || '';\n            \n            const breakingChanges = new Map();\n            \n            let currentPackage = '';\n            let isInIncompatibleSection = false;\n            const lines = diffOutput.split('\\n');\n            \n            for (let i = 0; i < lines.length; i++) {\n              const line = lines[i].trim();\n              \n              if (line.startsWith('github.com/')) {\n                currentPackage = line;\n                if (!breakingChanges.has(currentPackage)) {\n                  breakingChanges.set(currentPackage, []);\n                }\n                continue;\n              }\n              \n              if (line === 'Incompatible changes:') {\n                isInIncompatibleSection = true;\n                continue;\n              }\n              \n              if (line === '') {\n                isInIncompatibleSection = false;\n                continue;\n              }\n              \n              if (isInIncompatibleSection && line.startsWith('- ')) {\n                const change = line.substring(2);\n                if (currentPackage) {\n                  breakingChanges.get(currentPackage).push(change);\n                }\n              }\n            }\n            \n            const changedFiles = files.data;\n            \n            for (const [pkg, changes] of breakingChanges) {\n              if (changes.length === 0) continue; \n              \n              const pkgPath = pkg.split('/').slice(3).join('/');\n              const matchingFile = changedFiles.find(file => \n                file.filename.includes(pkgPath)\n              ) || changedFiles[0];\n              \n              const hasCommentForPackage = existingPackageComments.has(pkg) && \n                existingPackageComments.get(pkg).has(matchingFile.filename);\n              \n              if (matchingFile && !hasCommentForPackage) {\n                const changesList = changes.map(change => {\n                  const [name, desc] = change.split(':').map(s => s.trim());\n                  return `- **${name}:** ${desc}`;\n                }).join('\\n');\n                \n                const commentBody = [\n                  '🚨 **Breaking API Changes Detected**',\n                  '',\n                  `Package: \\`${pkg}\\``,\n                  '',\n                  'Incompatible changes:',\n                  changesList,\n                  '',\n                  '<details>',\n                  '<summary>Review Guidelines</summary>',\n                  '',\n                  'Please ensure that:',\n                  '- The changes are absolutely necessary',\n                  '- They are properly documented',\n                  '- Migration guides are provided if needed',\n                  '</details>',\n                  '',\n                  '⚠️ Please resolve this thread after reviewing the breaking changes.'\n                ].join('\\n');\n                \n                await github.rest.pulls.createReview({\n                  owner: context.repo.owner,\n                  repo: context.repo.repo,\n                  pull_number: context.issue.number,\n                  event: 'COMMENT',\n                  comments: [{\n                    path: matchingFile.filename,\n                    position: matchingFile.patch ? matchingFile.patch.split('\\n').findIndex(line => line.startsWith('+')) + 1 : 1,\n                    body: commentBody\n                  }]\n                });\n                \n                if (!existingPackageComments.has(pkg)) {\n                  existingPackageComments.set(pkg, new Set());\n                }\n                existingPackageComments.get(pkg).add(matchingFile.filename);\n              }\n            }"
  },
  {
    "path": ".gitignore",
    "content": "# Binaries for programs and plugins\n*.exe\n*.exe~\n*.dll\n*.so\n*.dylib\n\n# Test binary, built with `go test -c`\n*.test\n\n# Output of the go coverage tool, specifically when used with LiteIDE\n*.out\n\n# Dependency directories (remove the comment below to include it)\n# vendor/\n\n# Go workspace file\ngo.work\ngo.work.sum\n\n# env file\n.env\n\n# the result of the go build\noutput*\noutput/*\n\n# Files generated by IDEs\n.idea/\n*.iml\n\n# Vim swap files\n*.swp\n\n# Vscode files\n.vscode\n\n/patches\n\n/vendor\n\n# Trae files\n.trae\n\n# Specs files (internal documentation)\n**/specs/\n\n# Reports (generated analysis files)\nreports/\n\n.DS_Store\n*.log\nCLAUDE.md\n\n# Specs directories\n*/specs\n/todos\n/.claude/\n\n# Internal dev setup (not for public repo)\n/scripts/dev_setup_internal.sh\n"
  },
  {
    "path": ".golangci.yaml",
    "content": "# output configuration options\nversion: \"2\"\n# All available settings of specific linters.\n# Refer to https://golangci-lint.run/usage/linters\nlinters:\n  default: standard\n  enable:\n    - revive\n    - godoclint\n    - funlen\n    - cyclop\n  disable:\n    - errcheck\n    - staticcheck\n    - unused\n    - ineffassign\n  exclusions:\n    generated: lax\n    paths:\n      - \".*_test.go\"\n      - \".*_mock.go\"\n    rules:\n      - path: \"^internal/.*\"\n        linters:\n          - revive\n      - text: \"var-naming: don't use underscores in Go names\"\n        linters:\n          - revive\n      - path: \"/utils/\"\n        text: \"var-naming: avoid meaningless package names\"\n        linters:\n          - revive\n      - text: \"exported: type name will be used as agent.AgentOption by other packages\"\n        linters:\n          - revive\n      - path: \"adk/prebuilt/deep/task_tool.go\"\n        text: \"argument-limit: maximum number of arguments per function exceeded\"\n        linters:\n          - revive\n      - path: \"compose/component_to_graph_node.go\"\n        text: \"argument-limit: maximum number of arguments per function exceeded\"\n        linters:\n          - revive\n      - path: \"compose/graph_run.go\"\n        text: \"argument-limit: maximum number of arguments per function exceeded\"\n        linters:\n          - revive\n      - path: \"adk/workflow.go\"\n        text: \"argument-limit: maximum number of arguments per function exceeded\"\n        linters:\n          - revive\n      - path: \"compose/graph.go\"\n        linters:\n          - cyclop\n        text: \"calculated cyclomatic complexity for function compile\"\n      - path: \"schema/message.go\"\n        linters:\n          - cyclop\n        text: \"calculated cyclomatic complexity for function ConcatMessages\"\n      - path: \"compose/graph_run.go\"\n        linters:\n          - cyclop\n        text: \"calculated cyclomatic complexity for function run\"\n      - path: \"compose/graph.go\"\n        linters:\n          - funlen\n        text: \"Function 'compile' is too long\"\n      - path: \"compose/graph_run.go\"\n        linters:\n          - funlen\n        text: \"Function 'run' is too long\"\n  settings:\n    govet:\n      enable-all: true\n      # Disable analyzers by name.\n      # Run `go tool vet help` to see all analyzers.\n      disable:\n        - fieldalignment\n    revive:\n      # Sets the default failure confidence.\n      # This means that linting errors with less than 0.8 confidence will be ignored.\n      # Default: 0.8\n      confidence: 0.8\n      rules:\n        # Exported function and methods should have comments. \n        - name: exported\n          severity: error\n          exclude:\n            - \"^internal/.*\"\n          arguments:\n            - \"disable-checks-on-constants\"\n            - \"disable-checks-on-variables\"\n            - \"disable-checks-on-types\"\n            - \"disable-checks-on-methods\"\n        - name: package-comments\n          disabled: false\n        - name: var-naming\n          disabled: false\n          arguments:\n            # AllowList\n            - [ \"utils\", \"s_\", \"err_\", \"err__\", \"plan_\", \"userInput_\", \"executedSteps_\", \"executedStep_\", \"iterator_\", \"in_\", \"out_\" ]\n            # DenyList\n            - [ ]\n            - - extra-bad-package-names:\n                  - helpers\n                  - models\n        - name: argument-limit\n          arguments: [ 6 ]\n        - name: function-length\n          arguments: [ 120, 0 ]\n    godoclint:\n      check-exported: true\n      require-package-documentation: true\n    funlen:\n      lines: 200\n      statements: 120\n    cyclop:\n      max-complexity: 40\n      package-average: 20\n\nformatters:\n  enable:\n    - gci\n    - gofmt\n  settings:\n    gofmt:\n      # Simplify code: gofmt with `-s` option.\n      # Default: true\n      simplify: true\n      # Apply the rewrite rules to the source before reformatting.\n      # https://pkg.go.dev/cmd/gofmt\n      # Default: []\n      rewrite-rules:\n        - pattern: 'interface{}'\n          replacement: 'any'\n        - pattern: 'a[b:len(a)]'\n          replacement: 'a[b:]'\n    gci:\n      # Section configuration to compare against.\n      # Section names are case-insensitive and may contain parameters in ().\n      # The default order of sections is `standard > default > custom > blank > dot > alias > localmodule`.\n      # If `custom-order` is `true`, it follows the order of `sections` option.\n      # Default: [\"standard\", \"default\"]\n      sections:\n        - standard\n        - default\n        - localmodule\n      custom-order: true\n"
  },
  {
    "path": ".licenserc.yaml",
    "content": "header:\n  license:\n    spdx-id: Apache-2.0\n    copyright-owner: CloudWeGo Authors\n\n  template: |\n    /*\n     * Copyright {{ .Year }} CloudWeGo Authors\n     *\n     * Licensed under the Apache License, Version 2.0 (the \"License\");\n     * you may not use this file except in compliance with the License.\n     * You may obtain a copy of the License at\n     *\n     *     https://www.apache.org/licenses/LICENSE-2.0\n     *\n     * Unless required by applicable law or agreed to in writing, software\n     * distributed under the License is distributed on an \"AS IS\" BASIS,\n     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n     * See the License for the specific language governing permissions and\n     * limitations under the License.\n     */\n\n  paths:\n    - '**/*.go'\n    - '**/*.s'\n\n  comment: on-failure"
  },
  {
    "path": "CODE_OF_CONDUCT.md",
    "content": "# Contributor Covenant Code of Conduct\n\n## Our Pledge\n\nWe as members, contributors, and leaders pledge to make participation in our\ncommunity a harassment-free experience for everyone, regardless of age, body\nsize, visible or invisible disability, ethnicity, sex characteristics, gender\nidentity and expression, level of experience, education, socio-economic status,\nnationality, personal appearance, race, religion, or sexual identity\nand orientation.\n\nWe pledge to act and interact in ways that contribute to an open, welcoming,\ndiverse, inclusive, and healthy community.\n\n## Our Standards\n\nExamples of behavior that contributes to a positive environment for our\ncommunity include:\n\n* Demonstrating empathy and kindness toward other people\n* Being respectful of differing opinions, viewpoints, and experiences\n* Giving and gracefully accepting constructive feedback\n* Accepting responsibility and apologizing to those affected by our mistakes,\n  and learning from the experience\n* Focusing on what is best not just for us as individuals, but for the\n  overall community\n\nExamples of unacceptable behavior include:\n\n* The use of sexualized language or imagery, and sexual attention or\n  advances of any kind\n* Trolling, insulting or derogatory comments, and personal or political attacks\n* Public or private harassment\n* Publishing others' private information, such as a physical or email\n  address, without their explicit permission\n* Other conduct which could reasonably be considered inappropriate in a\n  professional setting\n\n## Enforcement Responsibilities\n\nCommunity leaders are responsible for clarifying and enforcing our standards of\nacceptable behavior and will take appropriate and fair corrective action in\nresponse to any behavior that they deem inappropriate, threatening, offensive,\nor harmful.\n\nCommunity leaders have the right and responsibility to remove, edit, or reject\ncomments, commits, code, wiki edits, issues, and other contributions that are\nnot aligned to this Code of Conduct, and will communicate reasons for moderation\ndecisions when appropriate.\n\n## Scope\n\nThis Code of Conduct applies within all community spaces, and also applies when\nan individual is officially representing the community in public spaces.\nExamples of representing our community include using an official e-mail address,\nposting via an official social media account, or acting as an appointed\nrepresentative at an online or offline event.\n\n## Enforcement\n\nInstances of abusive, harassing, or otherwise unacceptable behavior may be\nreported to the community leaders responsible for enforcement at\nconduct@cloudwego.io.\nAll complaints will be reviewed and investigated promptly and fairly.\n\nAll community leaders are obligated to respect the privacy and security of the\nreporter of any incident.\n\n## Enforcement Guidelines\n\nCommunity leaders will follow these Community Impact Guidelines in determining\nthe consequences for any action they deem in violation of this Code of Conduct:\n\n### 1. Correction\n\n**Community Impact**: Use of inappropriate language or other behavior deemed\nunprofessional or unwelcome in the community.\n\n**Consequence**: A private, written warning from community leaders, providing\nclarity around the nature of the violation and an explanation of why the\nbehavior was inappropriate. A public apology may be requested.\n\n### 2. Warning\n\n**Community Impact**: A violation through a single incident or series\nof actions.\n\n**Consequence**: A warning with consequences for continued behavior. No\ninteraction with the people involved, including unsolicited interaction with\nthose enforcing the Code of Conduct, for a specified period of time. This\nincludes avoiding interactions in community spaces as well as external channels\nlike social media. Violating these terms may lead to a temporary or\npermanent ban.\n\n### 3. Temporary Ban\n\n**Community Impact**: A serious violation of community standards, including\nsustained inappropriate behavior.\n\n**Consequence**: A temporary ban from any sort of interaction or public\ncommunication with the community for a specified period of time. No public or\nprivate interaction with the people involved, including unsolicited interaction\nwith those enforcing the Code of Conduct, is allowed during this period.\nViolating these terms may lead to a permanent ban.\n\n### 4. Permanent Ban\n\n**Community Impact**: Demonstrating a pattern of violation of community\nstandards, including sustained inappropriate behavior,  harassment of an\nindividual, or aggression toward or disparagement of classes of individuals.\n\n**Consequence**: A permanent ban from any sort of public interaction within\nthe community.\n\n## Attribution\n\nThis Code of Conduct is adapted from the [Contributor Covenant][homepage],\nversion 2.0, available at\nhttps://www.contributor-covenant.org/version/2/0/code_of_conduct.html.\n\nCommunity Impact Guidelines were inspired by [Mozilla's code of conduct\nenforcement ladder](https://github.com/mozilla/diversity).\n\n[homepage]: https://www.contributor-covenant.org\n\nFor answers to common questions about this code of conduct, see the FAQ at\nhttps://www.contributor-covenant.org/faq. Translations are available at\nhttps://www.contributor-covenant.org/translations.\n"
  },
  {
    "path": "CONTRIBUTING.md",
    "content": "# How to Contribute\n\n## Your First Pull Request\nWe use GitHub for our codebase. You can start by reading [How To Pull Request](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/about-pull-requests).\n\n## Branch Organization\nWe use [git-flow](https://nvie.com/posts/a-successful-git-branching-model/) as our branch organization, as known as [FDD](https://en.wikipedia.org/wiki/Feature-driven_development)\n\n## Bugs\n### 1. How to Find Known Issues\nWe are using [Github Issues](https://github.com/cloudwego/{project_name}/issues) for our public bugs. We keep a close eye on this and try to make it clear when we have an internal fix in progress. Before filing a new task, try to make sure your problem doesn’t already exist.\n\n### 2. Reporting New Issues\nProviding a reduced test code is a recommended way for reporting issues. Then can place in:\n- Just in issues\n- [Golang Playground](https://play.golang.org/)\n\n### 3. Security Bugs\nPlease do not report the safe disclosure of bugs to public issues. Contact us by [Support Email](mailto:conduct@cloudwego.io)\n\n## How to Get in Touch\n- [Email](mailto:conduct@cloudwego.io)\n\n## Submit a Pull Request\nBefore you submit your Pull Request (PR) consider the following guidelines:\n1. Search [GitHub](https://github.com/cloudwego/{project_name}/pulls) for an open or closed PR that relates to your submission. You don't want to duplicate existing efforts.\n2. Be sure that an issue describes the problem you're fixing, or documents the design for the feature you'd like to add. Discussing the design upfront helps to ensure that we're ready to accept your work.\n3. [Fork](https://docs.github.com/en/github/getting-started-with-github/fork-a-repo) the cloudwego {project_name} repo.\n4. In your forked repository, make your changes in a new git branch:\n    ```\n    git checkout -b my-fix-branch develop\n    ```\n5. Create your patch, including appropriate test cases.\n6. Follow our [Style Guides](#code-style-guides).\n7. Commit your changes using a descriptive commit message that follows [AngularJS Git Commit Message Conventions](https://docs.google.com/document/d/1QrDFcIiPjSLDn3EL15IJygNPiHORgU1_OOAqWjiDU5Y/edit).\n   Adherence to these conventions is necessary because release notes are automatically generated from these messages.\n8. Push your branch to GitHub:\n    ```\n    git push origin my-fix-branch\n    ```\n9. In GitHub, send a pull request to `{project_name}:develop`\n\n## Contribution Prerequisites\n- Our development environment keeps up with [Go Official](https://golang.org/project/).\n- You need fully checking with lint tools before submit your pull request. [gofmt](https://golang.org/pkg/cmd/gofmt/) and [golangci-lint](https://github.com/golangci/golangci-lint)\n- You are familiar with [GitHub](https://github.com)\n- Maybe you need familiar with [Actions](https://github.com/features/actions)(our default workflow tool).\n\n## Code Style Guides\n- [Effective Go](https://golang.org/doc/effective_go)\n- [Go Code Review Comments](https://github.com/golang/go/wiki/CodeReviewComments)\n"
  },
  {
    "path": "LICENSE-APACHE",
    "content": "                                 Apache License\n                           Version 2.0, January 2004\n                        http://www.apache.org/licenses/\n\n   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n   1. Definitions.\n\n      \"License\" shall mean the terms and conditions for use, reproduction,\n      and distribution as defined by Sections 1 through 9 of this document.\n\n      \"Licensor\" shall mean the copyright owner or entity authorized by\n      the copyright owner that is granting the License.\n\n      \"Legal Entity\" shall mean the union of the acting entity and all\n      other entities that control, are controlled by, or are under common\n      control with that entity. For the purposes of this definition,\n      \"control\" means (i) the power, direct or indirect, to cause the\n      direction or management of such entity, whether by contract or\n      otherwise, or (ii) ownership of fifty percent (50%) or more of the\n      outstanding shares, or (iii) beneficial ownership of such entity.\n\n      \"You\" (or \"Your\") shall mean an individual or Legal Entity\n      exercising permissions granted by this License.\n\n      \"Source\" form shall mean the preferred form for making modifications,\n      including but not limited to software source code, documentation\n      source, and configuration files.\n\n      \"Object\" form shall mean any form resulting from mechanical\n      transformation or translation of a Source form, including but\n      not limited to compiled object code, generated documentation,\n      and conversions to other media types.\n\n      \"Work\" shall mean the work of authorship, whether in Source or\n      Object form, made available under the License, as indicated by a\n      copyright notice that is included in or attached to the work\n      (an example is provided in the Appendix below).\n\n      \"Derivative Works\" shall mean any work, whether in Source or Object\n      form, that is based on (or derived from) the Work and for which the\n      editorial revisions, annotations, elaborations, or other modifications\n      represent, as a whole, an original work of authorship. For the purposes\n      of this License, Derivative Works shall not include works that remain\n      separable from, or merely link (or bind by name) to the interfaces of,\n      the Work and Derivative Works thereof.\n\n      \"Contribution\" shall mean any work of authorship, including\n      the original version of the Work and any modifications or additions\n      to that Work or Derivative Works thereof, that is intentionally\n      submitted to Licensor for inclusion in the Work by the copyright owner\n      or by an individual or Legal Entity authorized to submit on behalf of\n      the copyright owner. For the purposes of this definition, \"submitted\"\n      means any form of electronic, verbal, or written communication sent\n      to the Licensor or its representatives, including but not limited to\n      communication on electronic mailing lists, source code control systems,\n      and issue tracking systems that are managed by, or on behalf of, the\n      Licensor for the purpose of discussing and improving the Work, but\n      excluding communication that is conspicuously marked or otherwise\n      designated in writing by the copyright owner as \"Not a Contribution.\"\n\n      \"Contributor\" shall mean Licensor and any individual or Legal Entity\n      on behalf of whom a Contribution has been received by Licensor and\n      subsequently incorporated within the Work.\n\n   2. Grant of Copyright License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      copyright license to reproduce, prepare Derivative Works of,\n      publicly display, publicly perform, sublicense, and distribute the\n      Work and such Derivative Works in Source or Object form.\n\n   3. Grant of Patent License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      (except as stated in this section) patent license to make, have made,\n      use, offer to sell, sell, import, and otherwise transfer the Work,\n      where such license applies only to those patent claims licensable\n      by such Contributor that are necessarily infringed by their\n      Contribution(s) alone or by combination of their Contribution(s)\n      with the Work to which such Contribution(s) was submitted. If You\n      institute patent litigation against any entity (including a\n      cross-claim or counterclaim in a lawsuit) alleging that the Work\n      or a Contribution incorporated within the Work constitutes direct\n      or contributory patent infringement, then any patent licenses\n      granted to You under this License for that Work shall terminate\n      as of the date such litigation is filed.\n\n   4. Redistribution. You may reproduce and distribute copies of the\n      Work or Derivative Works thereof in any medium, with or without\n      modifications, and in Source or Object form, provided that You\n      meet the following conditions:\n\n      (a) You must give any other recipients of the Work or\n          Derivative Works a copy of this License; and\n\n      (b) You must cause any modified files to carry prominent notices\n          stating that You changed the files; and\n\n      (c) You must retain, in the Source form of any Derivative Works\n          that You distribute, all copyright, patent, trademark, and\n          attribution notices from the Source form of the Work,\n          excluding those notices that do not pertain to any part of\n          the Derivative Works; and\n\n      (d) If the Work includes a \"NOTICE\" text file as part of its\n          distribution, then any Derivative Works that You distribute must\n          include a readable copy of the attribution notices contained\n          within such NOTICE file, excluding those notices that do not\n          pertain to any part of the Derivative Works, in at least one\n          of the following places: within a NOTICE text file distributed\n          as part of the Derivative Works; within the Source form or\n          documentation, if provided along with the Derivative Works; or,\n          within a display generated by the Derivative Works, if and\n          wherever such third-party notices normally appear. The contents\n          of the NOTICE file are for informational purposes only and\n          do not modify the License. You may add Your own attribution\n          notices within Derivative Works that You distribute, alongside\n          or as an addendum to the NOTICE text from the Work, provided\n          that such additional attribution notices cannot be construed\n          as modifying the License.\n\n      You may add Your own copyright statement to Your modifications and\n      may provide additional or different license terms and conditions\n      for use, reproduction, or distribution of Your modifications, or\n      for any such Derivative Works as a whole, provided Your use,\n      reproduction, and distribution of the Work otherwise complies with\n      the conditions stated in this License.\n\n   5. Submission of Contributions. Unless You explicitly state otherwise,\n      any Contribution intentionally submitted for inclusion in the Work\n      by You to the Licensor shall be under the terms and conditions of\n      this License, without any additional terms or conditions.\n      Notwithstanding the above, nothing herein shall supersede or modify\n      the terms of any separate license agreement you may have executed\n      with Licensor regarding such Contributions.\n\n   6. Trademarks. This License does not grant permission to use the trade\n      names, trademarks, service marks, or product names of the Licensor,\n      except as required for reasonable and customary use in describing the\n      origin of the Work and reproducing the content of the NOTICE file.\n\n   7. Disclaimer of Warranty. Unless required by applicable law or\n      agreed to in writing, Licensor provides the Work (and each\n      Contributor provides its Contributions) on an \"AS IS\" BASIS,\n      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n      implied, including, without limitation, any warranties or conditions\n      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n      PARTICULAR PURPOSE. You are solely responsible for determining the\n      appropriateness of using or redistributing the Work and assume any\n      risks associated with Your exercise of permissions under this License.\n\n   8. Limitation of Liability. In no event and under no legal theory,\n      whether in tort (including negligence), contract, or otherwise,\n      unless required by applicable law (such as deliberate and grossly\n      negligent acts) or agreed to in writing, shall any Contributor be\n      liable to You for damages, including any direct, indirect, special,\n      incidental, or consequential damages of any character arising as a\n      result of this License or out of the use or inability to use the\n      Work (including but not limited to damages for loss of goodwill,\n      work stoppage, computer failure or malfunction, or any and all\n      other commercial damages or losses), even if such Contributor\n      has been advised of the possibility of such damages.\n\n   9. Accepting Warranty or Additional Liability. While redistributing\n      the Work or Derivative Works thereof, You may choose to offer,\n      and charge a fee for, acceptance of support, warranty, indemnity,\n      or other liability obligations and/or rights consistent with this\n      License. However, in accepting such obligations, You may act only\n      on Your own behalf and on Your sole responsibility, not on behalf\n      of any other Contributor, and only if You agree to indemnify,\n      defend, and hold each Contributor harmless for any liability\n      incurred by, or claims asserted against, such Contributor by reason\n      of your accepting any such warranty or additional liability.\n\n   END OF TERMS AND CONDITIONS\n\n"
  },
  {
    "path": "README.md",
    "content": "# Eino\n\n![coverage](https://raw.githubusercontent.com/cloudwego/eino/badges/.badges/main/coverage.svg)\n[![Release](https://img.shields.io/github/v/release/cloudwego/eino)](https://github.com/cloudwego/eino/releases)\n[![WebSite](https://img.shields.io/website?up_message=cloudwego&url=https%3A%2F%2Fwww.cloudwego.io%2F)](https://www.cloudwego.io/)\n[![License](https://img.shields.io/github/license/cloudwego/eino)](https://github.com/cloudwego/eino/blob/main/LICENSE)\n[![Go Report Card](https://goreportcard.com/badge/github.com/cloudwego/eino)](https://goreportcard.com/report/github.com/cloudwego/eino)\n[![OpenIssue](https://img.shields.io/github/issues/cloudwego/eino)](https://github.com/cloudwego/kitex/eino)\n[![ClosedIssue](https://img.shields.io/github/issues-closed/cloudwego/eino)](https://github.com/cloudwego/eino/issues?q=is%3Aissue+is%3Aclosed)\n![Stars](https://img.shields.io/github/stars/cloudwego/eino)\n![Forks](https://img.shields.io/github/forks/cloudwego/eino)\n\nEnglish | [中文](README.zh_CN.md)\n\n# Overview\n\n**Eino['aino]** is an LLM application development framework in Golang. It draws from LangChain, Google ADK, and other open-source frameworks, and is designed to follow Golang conventions.\n\nEino provides:\n- **[Components](https://github.com/cloudwego/eino-ext)**: reusable building blocks like `ChatModel`, `Tool`, `Retriever`, and `ChatTemplate`, with official implementations for OpenAI, Ollama, and more.\n- **Agent Development Kit (ADK)**: build AI agents with tool use, multi-agent coordination, context management, interrupt/resume for human-in-the-loop, and ready-to-use agent patterns.\n- **Composition**: connect components into graphs and workflows that can run standalone or be exposed as tools for agents.\n- **[Examples](https://github.com/cloudwego/eino-examples)**: working code for common patterns and real-world use cases.\n\n![](.github/static/img/eino/eino_concept.jpeg)\n\n# Quick Start\n\n## ChatModelAgent\n\nConfigure a ChatModel, optionally add tools, and you have a working agent:\n\n```Go\nchatModel, _ := openai.NewChatModel(ctx, &openai.ChatModelConfig{\n    Model:  \"gpt-4o\",\n    APIKey: os.Getenv(\"OPENAI_API_KEY\"),\n})\n\nagent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n    Model: chatModel,\n})\n\nrunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent})\niter := runner.Query(ctx, \"Hello, who are you?\")\nfor {\n    event, ok := iter.Next()\n    if !ok {\n        break\n    }\n    fmt.Println(event.Message.Content)\n}\n```\n\nAdd tools to give the agent capabilities:\n\n```Go\nagent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n    Model: chatModel,\n    ToolsConfig: adk.ToolsConfig{\n        ToolsNodeConfig: compose.ToolsNodeConfig{\n            Tools: []tool.BaseTool{weatherTool, calculatorTool},\n        },\n    },\n})\n```\n\nThe agent handles the ReAct loop internally — it decides when to call tools and when to respond.\n\n→ [ChatModelAgent examples](https://github.com/cloudwego/eino-examples/tree/main/adk/intro) · [docs](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/chat_model/)\n\n## DeepAgent\n\nFor complex tasks, use DeepAgent. It breaks down problems into steps, delegates to sub-agents, and tracks progress:\n\n```Go\ndeepAgent, _ := deep.New(ctx, &deep.Config{\n    ChatModel: chatModel,\n    SubAgents: []adk.Agent{researchAgent, codeAgent},\n    ToolsConfig: adk.ToolsConfig{\n        ToolsNodeConfig: compose.ToolsNodeConfig{\n            Tools: []tool.BaseTool{shellTool, pythonTool, webSearchTool},\n        },\n    },\n})\n\nrunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: deepAgent})\niter := runner.Query(ctx, \"Analyze the sales data in report.csv and generate a summary chart\")\n```\n\nDeepAgent can be configured to coordinate multiple specialized agents, run shell commands, execute Python code, and search the web.\n\n→ [DeepAgent example](https://github.com/cloudwego/eino-examples/tree/main/adk/multiagent/deep) · [docs](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/deepagents/)\n\n## Composition\n\nWhen you need precise control over execution flow, use `compose` to build graphs and workflows:\n\n```Go\ngraph := compose.NewGraph[*Input, *Output]()\ngraph.AddLambdaNode(\"validate\", validateFn)\ngraph.AddChatModelNode(\"generate\", chatModel)\ngraph.AddLambdaNode(\"format\", formatFn)\n\ngraph.AddEdge(compose.START, \"validate\")\ngraph.AddEdge(\"validate\", \"generate\")\ngraph.AddEdge(\"generate\", \"format\")\ngraph.AddEdge(\"format\", compose.END)\n\nrunnable, _ := graph.Compile(ctx)\nresult, _ := runnable.Invoke(ctx, input)\n```\n\nCompositions can be exposed as tools for agents, bridging deterministic workflows with autonomous behavior:\n\n```Go\ntool, _ := graphtool.NewInvokableGraphTool(graph, \"data_pipeline\", \"Process and validate data\")\n\nagent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n    Model: chatModel,\n    ToolsConfig: adk.ToolsConfig{\n        ToolsNodeConfig: compose.ToolsNodeConfig{\n            Tools: []tool.BaseTool{tool},\n        },\n    },\n})\n```\n\nThis lets you build domain-specific pipelines with exact control, then let agents decide when to use them.\n\n→ [GraphTool examples](https://github.com/cloudwego/eino-examples/tree/main/adk/common/tool/graphtool) · [compose docs](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/)\n\n# Key Features\n\n## Component Ecosystem\n\nEino defines component abstractions (ChatModel, Tool, Retriever, Embedding, etc.) with official implementations for OpenAI, Claude, Gemini, Ark, Ollama, Elasticsearch, and more.\n\n→ [eino-ext](https://github.com/cloudwego/eino-ext)\n\n## Stream Processing\n\nEino automatically handles streaming throughout orchestration: concatenating, boxing, merging, and copying streams as data flows between nodes. Components only implement the streaming paradigms that make sense for them; the framework handles the rest.\n\n→ [docs](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/)\n\n## Callback Aspects\n\nInject logging, tracing, and metrics at fixed points (OnStart, OnEnd, OnError, OnStartWithStreamInput, OnEndWithStreamOutput) across components, graphs, and agents.\n\n→ [docs](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/callback_manual/)\n\n## Interrupt/Resume\n\nAny agent or tool can pause execution for human input and resume from checkpoint. The framework handles state persistence and routing.\n\n→ [docs](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_hitl/) · [examples](https://github.com/cloudwego/eino-examples/tree/main/adk/human-in-the-loop)\n\n# Framework Structure\n\n![](.github/static/img/eino/eino_framework.jpeg)\n\nThe Eino framework consists of:\n\n- Eino (this repo): Type definitions, streaming mechanism, component abstractions, orchestration, agent implementations, aspect mechanisms\n\n- [EinoExt](https://github.com/cloudwego/eino-ext): Component implementations, callback handlers, usage examples, evaluators, prompt optimizers\n\n- [Eino Devops](https://github.com/cloudwego/eino-ext/tree/main/devops): Visualized development and debugging\n\n- [EinoExamples](https://github.com/cloudwego/eino-examples): Example applications and best practices\n\n## Documentation\n\n- [Eino User Manual](https://www.cloudwego.io/zh/docs/eino/)\n- [Eino: Quick Start](https://www.cloudwego.io/zh/docs/eino/quick_start/)\n\n## Dependencies\n- Go 1.18 and above.\n\n## Code Style\n\nThis repo uses `golangci-lint`. Check locally with:\n\n```bash\ngolangci-lint run ./...\n```\n\nRules enforced:\n- Exported functions, interfaces, packages, etc. should have GoDoc comments\n- Code should be formatted with `gofmt -s`\n- Import order should follow `goimports` (std -> third party -> local)\n\n## Security\n\nIf you discover a potential security issue, notify Bytedance Security via the [security center](https://security.bytedance.com/src) or [vulnerability reporting email](sec@bytedance.com).\n\nDo **not** create a public GitHub issue.\n\n## Contact\n- Membership: [COMMUNITY MEMBERSHIP](https://github.com/cloudwego/community/blob/main/COMMUNITY_MEMBERSHIP.md)\n- Issues: [Issues](https://github.com/cloudwego/eino/issues)\n- Lark: Scan the QR code below with [Feishu](https://www.feishu.cn/en/) to join the CloudWeGo/eino user group.\n\n&ensp;&ensp;&ensp; <img src=\".github/static/img/eino/lark_group_zh.png\" alt=\"LarkGroup\" width=\"200\"/>\n\n## License\n\nThis project is licensed under the [Apache-2.0 License](LICENSE-APACHE).\n"
  },
  {
    "path": "README.zh_CN.md",
    "content": "# Eino\n\n![coverage](https://raw.githubusercontent.com/cloudwego/eino/badges/.badges/main/coverage.svg)\n[![Release](https://img.shields.io/github/v/release/cloudwego/eino)](https://github.com/cloudwego/eino/releases)\n[![WebSite](https://img.shields.io/website?up_message=cloudwego&url=https%3A%2F%2Fwww.cloudwego.io%2F)](https://www.cloudwego.io/)\n[![License](https://img.shields.io/github/license/cloudwego/eino)](https://github.com/cloudwego/eino/blob/main/LICENSE)\n[![Go Report Card](https://goreportcard.com/badge/github.com/cloudwego/eino)](https://goreportcard.com/report/github.com/cloudwego/eino)\n[![OpenIssue](https://img.shields.io/github/issues/cloudwego/eino)](https://github.com/cloudwego/kitex/eino)\n[![ClosedIssue](https://img.shields.io/github/issues-closed/cloudwego/eino)](https://github.com/cloudwego/eino/issues?q=is%3Aissue+is%3Aclosed)\n![Stars](https://img.shields.io/github/stars/cloudwego/eino)\n![Forks](https://img.shields.io/github/forks/cloudwego/eino)\n\n[English](README.md) | 中文\n\n# 简介\n\n**Eino['aino]** 是一个 Go 语言的 LLM 应用开发框架，借鉴了 LangChain、Google ADK 等开源项目，按照 Go 的惯例设计。\n\nEino 提供：\n- **[组件](https://github.com/cloudwego/eino-ext)**：`ChatModel`、`Tool`、`Retriever`、`ChatTemplate` 等可复用模块，官方实现覆盖 OpenAI、Ollama 等\n- **智能体开发套件（ADK）**：支持工具调用、多智能体协同、上下文管理、中断/恢复等人机交互，以及开箱即用的智能体模式\n- **编排**：把组件组装成图或工作流，既能独立运行，也能作为工具给智能体调用\n- **[示例](https://github.com/cloudwego/eino-examples)**：常见模式和实际场景的可运行代码\n\n![](.github/static/img/eino/eino_concept.jpeg)\n\n# 快速上手\n\n## ChatModelAgent\n\n配置好 ChatModel，加上工具（可选），就能跑起来：\n\n```Go\nchatModel, _ := openai.NewChatModel(ctx, &openai.ChatModelConfig{\n    Model:  \"gpt-4o\",\n    APIKey: os.Getenv(\"OPENAI_API_KEY\"),\n})\n\nagent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n    Model: chatModel,\n})\n\nrunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent})\niter := runner.Query(ctx, \"Hello, who are you?\")\nfor {\n    event, ok := iter.Next()\n    if !ok {\n        break\n    }\n    fmt.Println(event.Message.Content)\n}\n```\n\n加工具让智能体有更多能力：\n\n```Go\nagent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n    Model: chatModel,\n    ToolsConfig: adk.ToolsConfig{\n        ToolsNodeConfig: compose.ToolsNodeConfig{\n            Tools: []tool.BaseTool{weatherTool, calculatorTool},\n        },\n    },\n})\n```\n\n智能体内部自动处理 ReAct 循环，自己判断什么时候调工具、什么时候回复。\n\n→ [ChatModelAgent 示例](https://github.com/cloudwego/eino-examples/tree/main/adk/intro) · [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/eino_adk/agent_implementation/chat_model/)\n\n## DeepAgent\n\n复杂任务用 DeepAgent，它会把问题拆成步骤，分派给子智能体，并追踪进度：\n\n```Go\ndeepAgent, _ := deep.New(ctx, &deep.Config{\n    ChatModel: chatModel,\n    SubAgents: []adk.Agent{researchAgent, codeAgent},\n    ToolsConfig: adk.ToolsConfig{\n        ToolsNodeConfig: compose.ToolsNodeConfig{\n            Tools: []tool.BaseTool{shellTool, pythonTool, webSearchTool},\n        },\n    },\n})\n\nrunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: deepAgent})\niter := runner.Query(ctx, \"Analyze the sales data in report.csv and generate a summary chart\")\n```\n\nDeepAgent 可以配置成：协调多个专业智能体、跑 shell 命令、执行 Python、搜索网络。\n\n→ [DeepAgent 示例](https://github.com/cloudwego/eino-examples/tree/main/adk/multiagent/deep) · [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/eino_adk/agent_implementation/deepagents/)\n\n## 编排\n\n需要精确控制执行流程时，用 `compose` 搭图或工作流：\n\n```Go\ngraph := compose.NewGraph[*Input, *Output]()\ngraph.AddLambdaNode(\"validate\", validateFn)\ngraph.AddChatModelNode(\"generate\", chatModel)\ngraph.AddLambdaNode(\"format\", formatFn)\n\ngraph.AddEdge(compose.START, \"validate\")\ngraph.AddEdge(\"validate\", \"generate\")\ngraph.AddEdge(\"generate\", \"format\")\ngraph.AddEdge(\"format\", compose.END)\n\nrunnable, _ := graph.Compile(ctx)\nresult, _ := runnable.Invoke(ctx, input)\n```\n\n编排出来的流程可以包装成工具给智能体用，把确定性流程和自主决策结合起来：\n\n```Go\ntool, _ := graphtool.NewInvokableGraphTool(graph, \"data_pipeline\", \"Process and validate data\")\n\nagent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n    Model: chatModel,\n    ToolsConfig: adk.ToolsConfig{\n        ToolsNodeConfig: compose.ToolsNodeConfig{\n            Tools: []tool.BaseTool{tool},\n        },\n    },\n})\n```\n\n这样你可以写出精确可控的业务流程，再让智能体决定什么时候调用。\n\n→ [GraphTool 示例](https://github.com/cloudwego/eino-examples/tree/main/adk/common/tool/graphtool) · [编排文档](https://www.cloudwego.io/zh/docs/eino/core_modules/chain_and_graph_orchestration/)\n\n# 主要特性\n\n## 组件生态\n\nEino 定义了组件抽象（ChatModel、Tool、Retriever、Embedding 等），官方实现覆盖 OpenAI、Claude、Gemini、Ark、Ollama、Elasticsearch 等。\n\n→ [eino-ext](https://github.com/cloudwego/eino-ext)\n\n## 流式处理\n\nEino 在编排中自动处理流式：拼接、装箱、合并、复制。组件只需实现有业务意义的流式范式，框架处理剩下的。\n\n→ [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/)\n\n## 回调切面\n\n在固定切点（OnStart、OnEnd、OnError、OnStartWithStreamInput、OnEndWithStreamOutput）注入日志、追踪、指标，适用于组件、图、智能体。\n\n→ [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/chain_and_graph_orchestration/callback_manual/)\n\n## 中断/恢复\n\n任何智能体或工具都能暂停等待人工输入，从检查点恢复。框架处理状态持久化和路由。\n\n→ [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/eino_adk/agent_hitl/) · [示例](https://github.com/cloudwego/eino-examples/tree/main/adk/human-in-the-loop)\n\n# 框架结构\n\n![](.github/static/img/eino/eino_framework.jpeg)\n\nEino 框架包含：\n\n- Eino（本仓库）：类型定义、流处理机制、组件抽象、编排、智能体实现、切面机制\n\n- [EinoExt](https://github.com/cloudwego/eino-ext)：组件实现、回调处理器、使用示例、评估器、提示优化器\n\n- [Eino Devops](https://github.com/cloudwego/eino-ext/tree/main/devops)：可视化开发和调试\n\n- [EinoExamples](https://github.com/cloudwego/eino-examples)：示例应用和最佳实践\n\n## 文档\n\n- [Eino 用户手册](https://www.cloudwego.io/zh/docs/eino/)\n- [Eino: 快速开始](https://www.cloudwego.io/zh/docs/eino/quick_start/)\n\n## 依赖\n- Go 1.18 及以上\n\n## 代码规范\n\n本仓库使用 `golangci-lint`，本地检查：\n\n```bash\ngolangci-lint run ./...\n```\n\n规则：\n- 导出的函数、接口、package 等需要 GoDoc 注释\n- 代码格式符合 `gofmt -s`\n- import 顺序符合 `goimports`（std -> third party -> local）\n\n## 安全\n\n发现安全问题请通过[安全中心](https://security.bytedance.com/src)或[漏洞报告邮箱](sec@bytedance.com)联系字节跳动安全团队。\n\n请**不要**创建公开的 GitHub Issue。\n\n## 联系我们\n- 成为 member：[COMMUNITY MEMBERSHIP](https://github.com/cloudwego/community/blob/main/COMMUNITY_MEMBERSHIP.md)\n- Issues：[Issues](https://github.com/cloudwego/eino/issues)\n- 飞书：扫码加入 CloudWeGo/eino 用户群\n\n&ensp;&ensp;&ensp; <img src=\".github/static/img/eino/lark_group_zh.png\" alt=\"LarkGroup\" width=\"200\"/>\n\n## 开源许可证\n\n本项目基于 [Apache-2.0 许可证](LICENSE-APACHE) 开源。\n"
  },
  {
    "path": "_typos.toml",
    "content": "# Typo check: https://github.com/crate-ci/typos\n[default]\n\n[default.extend-words]\nInvokable = \"Invokable\"\ninvokable = \"invokable\"\nInvokableLambda = \"InvokableLambda\"\nInvokableRun = \"InvokableRun\"\ntyp = \"typ\"\nbyted = \"byted\"\ncpy = \"cpy\"\nmak = \"mak\"\n[files]\nextend-exclude = [\"go.mod\", \"go.sum\", \"check_branch_name.sh\"]\n"
  },
  {
    "path": "adk/agent_tool.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package adk provides core agent development kit utilities and types.\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nvar (\n\tdefaultAgentToolParam = schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\"request\": {\n\t\t\tDesc:     \"request to be processed\",\n\t\t\tRequired: true,\n\t\t\tType:     schema.String,\n\t\t},\n\t})\n)\n\ntype AgentToolOptions struct {\n\tfullChatHistoryAsInput bool\n\tagentInputSchema       *schema.ParamsOneOf\n}\n\ntype AgentToolOption func(*AgentToolOptions)\n\n// WithFullChatHistoryAsInput enables using the full chat history as input.\nfunc WithFullChatHistoryAsInput() AgentToolOption {\n\treturn func(options *AgentToolOptions) {\n\t\toptions.fullChatHistoryAsInput = true\n\t}\n}\n\n// WithAgentInputSchema sets a custom input schema for the agent tool.\nfunc WithAgentInputSchema(schema *schema.ParamsOneOf) AgentToolOption {\n\treturn func(options *AgentToolOptions) {\n\t\toptions.agentInputSchema = schema\n\t}\n}\n\nfunc withAgentToolEnableStreaming(enabled bool) tool.Option {\n\treturn tool.WrapImplSpecificOptFn(func(opt *agentToolOptions) {\n\t\topt.enableStreaming = enabled\n\t})\n}\n\n// NewAgentTool creates a tool that wraps an agent for invocation.\n//\n// Event Streaming:\n// When EmitInternalEvents is enabled in ToolsConfig, the agent tool will emit AgentEvent\n// from the inner agent to the parent agent's AsyncGenerator, allowing real-time streaming\n// of the inner agent's output to the end-user via Runner.\n//\n// Note that these forwarded events are NOT recorded in the parent agent's runSession.\n// They are only emitted to the end-user and have no effect on the parent agent's state\n// or checkpoint. The only exception is Interrupted action, which is propagated via\n// CompositeInterrupt to enable proper interrupt/resume across agent boundaries.\n//\n// Action Scoping:\n// Actions emitted by the inner agent are scoped to the agent tool boundary:\n//   - Interrupted: Propagated via CompositeInterrupt to allow proper interrupt/resume across boundaries\n//   - Exit, TransferToAgent, BreakLoop: Ignored outside the agent tool; these actions only affect\n//     the inner agent's execution and do not propagate to the parent agent\n//\n// This scoping ensures that nested agents cannot unexpectedly terminate or transfer control\n// of their parent agent's execution flow.\nfunc NewAgentTool(_ context.Context, agent Agent, options ...AgentToolOption) tool.BaseTool {\n\topts := &AgentToolOptions{}\n\tfor _, opt := range options {\n\t\topt(opts)\n\t}\n\n\treturn &agentTool{\n\t\tagent:                  agent,\n\t\tfullChatHistoryAsInput: opts.fullChatHistoryAsInput,\n\t\tinputSchema:            opts.agentInputSchema,\n\t}\n}\n\ntype agentTool struct {\n\tagent Agent\n\n\tfullChatHistoryAsInput bool\n\tinputSchema            *schema.ParamsOneOf\n}\n\nfunc (at *agentTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\tparam := at.inputSchema\n\tif param == nil {\n\t\tparam = defaultAgentToolParam\n\t}\n\n\treturn &schema.ToolInfo{\n\t\tName:        at.agent.Name(ctx),\n\t\tDesc:        at.agent.Description(ctx),\n\t\tParamsOneOf: param,\n\t}, nil\n}\n\nfunc (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tgen, enableStreaming := getEmitGeneratorAndEnableStreaming(opts)\n\tvar ms *bridgeStore\n\tvar iter *AsyncIterator[*AgentEvent]\n\tvar err error\n\n\twasInterrupted, hasState, state := tool.GetInterruptState[[]byte](ctx)\n\tif !wasInterrupted {\n\t\tms = newBridgeStore()\n\t\tvar input []Message\n\t\tif at.fullChatHistoryAsInput {\n\t\t\tinput, err = getReactChatHistory(ctx, at.agent.Name(ctx))\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", err\n\t\t\t}\n\t\t} else {\n\t\t\tif at.inputSchema == nil {\n\t\t\t\t// default input schema\n\t\t\t\ttype request struct {\n\t\t\t\t\tRequest string `json:\"request\"`\n\t\t\t\t}\n\n\t\t\t\treq := &request{}\n\t\t\t\terr = sonic.UnmarshalString(argumentsInJSON, req)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn \"\", err\n\t\t\t\t}\n\t\t\t\targumentsInJSON = req.Request\n\t\t\t}\n\t\t\tinput = []Message{\n\t\t\t\tschema.UserMessage(argumentsInJSON),\n\t\t\t}\n\t\t}\n\n\t\titer = newInvokableAgentToolRunner(at.agent, ms, enableStreaming).Run(ctx, input,\n\t\t\tappend(getOptionsByAgentName(at.agent.Name(ctx), opts), WithCheckPointID(bridgeCheckpointID), withSharedParentSession())...)\n\t} else {\n\t\tif !hasState {\n\t\t\treturn \"\", fmt.Errorf(\"agent tool '%s' interrupt has happened, but cannot find interrupt state\", at.agent.Name(ctx))\n\t\t}\n\n\t\tms = newResumeBridgeStore(state)\n\n\t\titer, err = newInvokableAgentToolRunner(at.agent, ms, enableStreaming).\n\t\t\tResume(ctx, bridgeCheckpointID, append(getOptionsByAgentName(at.agent.Name(ctx), opts), withSharedParentSession())...)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t}\n\n\tvar lastEvent *AgentEvent\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\n\t\tif lastEvent != nil &&\n\t\t\tlastEvent.Output != nil &&\n\t\t\tlastEvent.Output.MessageOutput != nil &&\n\t\t\tlastEvent.Output.MessageOutput.MessageStream != nil {\n\t\t\tlastEvent.Output.MessageOutput.MessageStream.Close()\n\t\t}\n\n\t\tif event.Err != nil {\n\t\t\treturn \"\", event.Err\n\t\t}\n\n\t\tif gen != nil {\n\t\t\tif event.Action == nil || event.Action.Interrupted == nil {\n\t\t\t\tif parentRunCtx := getRunCtx(ctx); parentRunCtx != nil && len(parentRunCtx.RunPath) > 0 {\n\t\t\t\t\trp := make([]RunStep, 0, len(parentRunCtx.RunPath)+len(event.RunPath))\n\t\t\t\t\trp = append(rp, parentRunCtx.RunPath...)\n\t\t\t\t\trp = append(rp, event.RunPath...)\n\t\t\t\t\tevent.RunPath = rp\n\t\t\t\t}\n\t\t\t\ttmp := copyAgentEvent(event)\n\t\t\t\tgen.Send(event)\n\t\t\t\tevent = tmp\n\t\t\t}\n\t\t}\n\n\t\tlastEvent = event\n\t}\n\n\tif lastEvent != nil && lastEvent.Action != nil && lastEvent.Action.Interrupted != nil {\n\t\tdata, existed, err_ := ms.Get(ctx, bridgeCheckpointID)\n\t\tif err_ != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to get interrupt info: %w\", err_)\n\t\t}\n\t\tif !existed {\n\t\t\treturn \"\", fmt.Errorf(\"interrupt has happened, but cannot find interrupt info\")\n\t\t}\n\n\t\treturn \"\", tool.CompositeInterrupt(ctx, \"agent tool interrupt\", data,\n\t\t\tlastEvent.Action.internalInterrupted)\n\t}\n\n\tif lastEvent == nil {\n\t\treturn \"\", errors.New(\"no event returned\")\n\t}\n\n\tvar ret string\n\tif lastEvent.Output != nil {\n\t\tif output := lastEvent.Output.MessageOutput; output != nil {\n\t\t\tmsg, err := output.GetMessage()\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", err\n\t\t\t}\n\t\t\tret = msg.Content\n\t\t}\n\t}\n\n\treturn ret, nil\n}\n\n// agentToolOptions is a wrapper structure used to convert AgentRunOption slices to tool.Option.\n// It stores the agent name and corresponding run options for tool-specific processing.\ntype agentToolOptions struct {\n\tagentName       string\n\topts            []AgentRunOption\n\tgenerator       *AsyncGenerator[*AgentEvent]\n\tenableStreaming bool\n}\n\nfunc withAgentToolOptions(agentName string, opts []AgentRunOption) tool.Option {\n\treturn tool.WrapImplSpecificOptFn(func(opt *agentToolOptions) {\n\t\topt.agentName = agentName\n\t\topt.opts = opts\n\t})\n}\n\nfunc withAgentToolEventGenerator(gen *AsyncGenerator[*AgentEvent]) tool.Option {\n\treturn tool.WrapImplSpecificOptFn(func(o *agentToolOptions) {\n\t\to.generator = gen\n\t})\n}\n\nfunc getOptionsByAgentName(agentName string, opts []tool.Option) []AgentRunOption {\n\tvar ret []AgentRunOption\n\tfor _, opt := range opts {\n\t\to := tool.GetImplSpecificOptions[agentToolOptions](nil, opt)\n\t\tif o != nil && o.agentName == agentName {\n\t\t\tret = append(ret, o.opts...)\n\t\t}\n\t}\n\treturn ret\n}\n\nfunc getEmitGeneratorAndEnableStreaming(opts []tool.Option) (*AsyncGenerator[*AgentEvent], bool) {\n\to := tool.GetImplSpecificOptions[agentToolOptions](nil, opts...)\n\tif o == nil {\n\t\treturn nil, false\n\t}\n\n\treturn o.generator, o.enableStreaming\n}\n\nfunc getReactChatHistory(ctx context.Context, destAgentName string) ([]Message, error) {\n\tvar messages []Message\n\terr := compose.ProcessState(ctx, func(ctx context.Context, st *State) error {\n\t\tmessages = make([]Message, len(st.Messages)-1)\n\t\tcopy(messages, st.Messages[:len(st.Messages)-1]) // remove the last assistant message, which is the tool call message\n\t\treturn nil\n\t})\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to get chat history from state: %w\", err)\n\t}\n\n\tvar agentName string\n\tif runCtx := getRunCtx(ctx); runCtx != nil && len(runCtx.RunPath) > 0 {\n\t\tagentName = runCtx.RunPath[len(runCtx.RunPath)-1].agentName\n\t}\n\n\ta, t := GenTransferMessages(ctx, destAgentName)\n\tmessages = append(messages, a, t)\n\thistory := make([]Message, 0, len(messages))\n\tfor _, msg := range messages {\n\t\tif msg.Role == schema.System {\n\t\t\tcontinue\n\t\t}\n\n\t\tif msg.Role == schema.Assistant || msg.Role == schema.Tool {\n\t\t\tmsg = rewriteMessage(msg, agentName)\n\t\t}\n\n\t\thistory = append(history, msg)\n\t}\n\n\treturn history, nil\n}\n\nfunc newInvokableAgentToolRunner(agent Agent, store compose.CheckPointStore, enableStreaming bool) *Runner {\n\treturn &Runner{\n\t\ta:               agent,\n\t\tenableStreaming: enableStreaming,\n\t\tstore:           store,\n\t}\n}\n"
  },
  {
    "path": "adk/agent_tool_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// mockAgent implements the Agent interface for testing\ntype mockAgentForTool struct {\n\tname        string\n\tdescription string\n\tresponses   []*AgentEvent\n}\n\nfunc (a *mockAgentForTool) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *mockAgentForTool) Description(_ context.Context) string {\n\treturn a.description\n}\n\nfunc (a *mockAgentForTool) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo func() {\n\t\tdefer generator.Close()\n\n\t\tfor _, event := range a.responses {\n\t\t\tgenerator.Send(event)\n\n\t\t\t// If the event has an Exit action, stop sending events\n\t\t\tif event.Action != nil && event.Action.Exit {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}()\n\n\treturn iterator\n}\n\nfunc newMockAgentForTool(name, description string, responses []*AgentEvent) *mockAgentForTool {\n\treturn &mockAgentForTool{\n\t\tname:        name,\n\t\tdescription: description,\n\t\tresponses:   responses,\n\t}\n}\n\nfunc TestAgentTool_Info(t *testing.T) {\n\t// Create a mock agent\n\tmockAgent_ := newMockAgentForTool(\"TestAgent\", \"Test agent description\", nil)\n\n\t// Create an agentTool with the mock agent\n\tagentTool_ := NewAgentTool(context.Background(), mockAgent_)\n\n\t// Test the Info method\n\tctx := context.Background()\n\tinfo, err := agentTool_.Info(ctx)\n\n\t// Verify results\n\tassert.NoError(t, err)\n\tassert.NotNil(t, info)\n\tassert.Equal(t, \"TestAgent\", info.Name)\n\tassert.Equal(t, \"Test agent description\", info.Desc)\n\tassert.NotNil(t, info.ParamsOneOf)\n}\n\nfunc TestAgentTool_SharedParentSessionValues(t *testing.T) {\n\tctx := context.Background()\n\n\tinner := &sessionValuesAgent{name: \"inner\"}\n\tinnerTool := NewAgentTool(ctx, inner).(tool.InvokableTool)\n\n\tinput := &AgentInput{Messages: []Message{schema.UserMessage(\"q\")}}\n\tctx, _ = initRunCtx(ctx, \"outer\", input)\n\tAddSessionValue(ctx, \"parent_key\", \"parent_val\")\n\tparentSession := getRunCtx(ctx).Session\n\n\t_, err := innerTool.InvokableRun(ctx, `{\"request\":\"hello\"}`)\n\tassert.NoError(t, err)\n\n\tassert.Equal(t, \"parent_val\", inner.seenParentValue)\n\tassert.NotNil(t, inner.capturedSession)\n\tassert.NotSame(t, parentSession, inner.capturedSession)\n\tassert.NotNil(t, parentSession.valuesMtx)\n\tassert.Same(t, parentSession.valuesMtx, inner.capturedSession.valuesMtx)\n\n\tmtx := parentSession.valuesMtx\n\tmtx.Lock()\n\tinner.capturedSession.Values[\"direct_child_key\"] = \"direct_child_val\"\n\tmtx.Unlock()\n\n\tmtx.Lock()\n\tv2, ok2 := parentSession.Values[\"direct_child_key\"]\n\tmtx.Unlock()\n\tassert.True(t, ok2)\n\tassert.Equal(t, \"direct_child_val\", v2)\n\n\tmtx.Lock()\n\tparentSession.Values[\"direct_parent_key\"] = \"direct_parent_val\"\n\tmtx.Unlock()\n\n\tmtx.Lock()\n\tv3, ok3 := inner.capturedSession.Values[\"direct_parent_key\"]\n\tmtx.Unlock()\n\tassert.True(t, ok3)\n\tassert.Equal(t, \"direct_parent_val\", v3)\n\n\tv, ok := GetSessionValue(ctx, \"child_key\")\n\tassert.True(t, ok)\n\tassert.Equal(t, \"child_val\", v)\n}\n\ntype sessionValuesAgent struct {\n\tname            string\n\tseenParentValue any\n\tcapturedSession *runSession\n}\n\nfunc (a *sessionValuesAgent) Name(context.Context) string        { return a.name }\nfunc (a *sessionValuesAgent) Description(context.Context) string { return \"test\" }\nfunc (a *sessionValuesAgent) Run(ctx context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tif rc := getRunCtx(ctx); rc != nil {\n\t\ta.capturedSession = rc.Session\n\t}\n\ta.seenParentValue, _ = GetSessionValue(ctx, \"parent_key\")\n\tAddSessionValue(ctx, \"child_key\", \"child_val\")\n\n\tit, gen := NewAsyncIteratorPair[*AgentEvent]()\n\tgen.Send(&AgentEvent{\n\t\tAgentName: a.name,\n\t\tOutput: &AgentOutput{\n\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\tIsStreaming: false,\n\t\t\t\tMessage:     schema.AssistantMessage(\"ok\", nil),\n\t\t\t\tRole:        schema.Assistant,\n\t\t\t},\n\t\t},\n\t})\n\tgen.Close()\n\treturn it\n}\n\nfunc TestAgentTool_InvokableRun(t *testing.T) {\n\t// Create a context\n\tctx := context.Background()\n\n\t// Test cases\n\ttests := []struct {\n\t\tname           string\n\t\tagentResponses []*AgentEvent\n\t\trequest        string\n\t\texpectedOutput string\n\t\texpectError    bool\n\t}{\n\t\t{\n\t\t\tname: \"successful model response\",\n\t\t\tagentResponses: []*AgentEvent{\n\t\t\t\t{\n\t\t\t\t\tAgentName: \"TestAgent\",\n\t\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\t\t\tMessage:     schema.AssistantMessage(\"Test response\", nil),\n\t\t\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\trequest:        `{\"request\":\"Test request\"}`,\n\t\t\texpectedOutput: \"Test response\",\n\t\t\texpectError:    false,\n\t\t},\n\t\t{\n\t\t\tname: \"successful tool call response\",\n\t\t\tagentResponses: []*AgentEvent{\n\t\t\t\t{\n\t\t\t\t\tAgentName: \"TestAgent\",\n\t\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\t\t\tMessage:     schema.ToolMessage(\"Tool response\", \"test-id\"),\n\t\t\t\t\t\t\tRole:        schema.Tool,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\trequest:        `{\"request\":\"Test tool request\"}`,\n\t\t\texpectedOutput: \"Tool response\",\n\t\t\texpectError:    false,\n\t\t},\n\t\t{\n\t\t\tname:           \"invalid request JSON\",\n\t\t\tagentResponses: nil,\n\t\t\trequest:        `invalid json`,\n\t\t\texpectedOutput: \"\",\n\t\t\texpectError:    true,\n\t\t},\n\t\t{\n\t\t\tname:           \"no events returned\",\n\t\t\tagentResponses: []*AgentEvent{},\n\t\t\trequest:        `{\"request\":\"Test request\"}`,\n\t\t\texpectedOutput: \"\",\n\t\t\texpectError:    true,\n\t\t},\n\t\t{\n\t\t\tname: \"error in event\",\n\t\t\tagentResponses: []*AgentEvent{\n\t\t\t\t{\n\t\t\t\t\tAgentName: \"TestAgent\",\n\t\t\t\t\tErr:       assert.AnError,\n\t\t\t\t},\n\t\t\t},\n\t\t\trequest:        `{\"request\":\"Test request\"}`,\n\t\t\texpectedOutput: \"\",\n\t\t\texpectError:    true,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\t// Create a mock agent with the test responses\n\t\t\tmockAgent_ := newMockAgentForTool(\"TestAgent\", \"Test agent description\", tt.agentResponses)\n\n\t\t\t// Create an agentTool with the mock agent\n\t\t\tagentTool_ := NewAgentTool(ctx, mockAgent_)\n\n\t\t\t// Call InvokableRun\n\t\t\toutput, err := agentTool_.(tool.InvokableTool).InvokableRun(ctx, tt.request)\n\n\t\t\t// Verify results\n\t\t\tif tt.expectError {\n\t\t\t\tassert.Error(t, err)\n\t\t\t} else {\n\t\t\t\tassert.NoError(t, err)\n\t\t\t\tassert.Equal(t, tt.expectedOutput, output)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestGetReactHistory(t *testing.T) {\n\tg := compose.NewGraph[string, []Message](compose.WithGenLocalState(func(ctx context.Context) (state *State) {\n\t\treturn &State{\n\t\t\tMessages: []Message{\n\t\t\t\tschema.UserMessage(\"user query\"),\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{{ID: \"tool call id 1\", Function: schema.FunctionCall{Name: \"tool1\", Arguments: \"arguments1\"}}}),\n\t\t\t\tschema.ToolMessage(\"tool result 1\", \"tool call id 1\", schema.WithToolName(\"tool1\")),\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{{ID: \"tool call id 2\", Function: schema.FunctionCall{Name: \"tool2\", Arguments: \"arguments2\"}}}),\n\t\t\t},\n\t\t}\n\t}))\n\tassert.NoError(t, g.AddLambdaNode(\"1\", compose.InvokableLambda(func(ctx context.Context, input string) (output []Message, err error) {\n\t\treturn getReactChatHistory(ctx, \"DestAgentName\")\n\t})))\n\tassert.NoError(t, g.AddEdge(compose.START, \"1\"))\n\tassert.NoError(t, g.AddEdge(\"1\", compose.END))\n\n\tctx := context.Background()\n\tctx, _ = initRunCtx(ctx, \"MyAgent\", nil)\n\trunner, err := g.Compile(ctx)\n\tassert.NoError(t, err)\n\tresult, err := runner.Invoke(ctx, \"\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, []Message{\n\t\tschema.UserMessage(\"user query\"),\n\t\tschema.UserMessage(\"For context: [MyAgent] called tool: `tool1` with arguments: arguments1.\"),\n\t\tschema.UserMessage(\"For context: [MyAgent] `tool1` tool returned result: tool result 1.\"),\n\t\tschema.UserMessage(\"For context: [MyAgent] called tool: `transfer_to_agent` with arguments: DestAgentName.\"),\n\t\tschema.UserMessage(\"For context: [MyAgent] `transfer_to_agent` tool returned result: successfully transferred to agent [DestAgentName].\"),\n\t}, result)\n}\n\n// mockAgentWithInputCapture implements the Agent interface for testing and captures the input it receives\ntype mockAgentWithInputCapture struct {\n\tname          string\n\tdescription   string\n\tcapturedInput []Message\n\tresponses     []*AgentEvent\n}\n\nfunc (a *mockAgentWithInputCapture) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *mockAgentWithInputCapture) Description(_ context.Context) string {\n\treturn a.description\n}\n\nfunc (a *mockAgentWithInputCapture) Run(_ context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\ta.capturedInput = input.Messages\n\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo func() {\n\t\tdefer generator.Close()\n\n\t\tfor _, event := range a.responses {\n\t\t\tgenerator.Send(event)\n\n\t\t\t// If the event has an Exit action, stop sending events\n\t\t\tif event.Action != nil && event.Action.Exit {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}()\n\n\treturn iterator\n}\n\nfunc newMockAgentWithInputCapture(name, description string, responses []*AgentEvent) *mockAgentWithInputCapture {\n\treturn &mockAgentWithInputCapture{\n\t\tname:        name,\n\t\tdescription: description,\n\t\tresponses:   responses,\n\t}\n}\n\nfunc TestAgentToolWithOptions(t *testing.T) {\n\t// Test Case 1: WithFullChatHistoryAsInput\n\tt.Run(\"WithFullChatHistoryAsInput\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// 1. Set up a mock agent that will capture the input it receives\n\t\tmockAgent := newMockAgentWithInputCapture(\"test-agent\", \"a test agent\", []*AgentEvent{\n\t\t\t{\n\t\t\t\tAgentName: \"test-agent\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\t\tMessage:     schema.AssistantMessage(\"done\", nil),\n\t\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\t// 2. Create an agentTool with the option\n\t\tagentTool := NewAgentTool(ctx, mockAgent, WithFullChatHistoryAsInput())\n\n\t\t// 3. Set up a context with a chat history using a graph\n\t\thistory := []Message{\n\t\t\tschema.UserMessage(\"first user message\"),\n\t\t\tschema.AssistantMessage(\"first assistant response\", nil),\n\t\t}\n\n\t\tg := compose.NewGraph[string, string](compose.WithGenLocalState(func(ctx context.Context) (state *State) {\n\t\t\treturn &State{\n\t\t\t\tMessages: append(history, schema.AssistantMessage(\"tool call msg\", nil)),\n\t\t\t}\n\t\t}))\n\n\t\tassert.NoError(t, g.AddLambdaNode(\"1\", compose.InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\t// Run the tool within the graph context that has the state\n\t\t\t_, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{\"request\":\"some ignored input\"}`)\n\t\t\treturn \"done\", err\n\t\t})))\n\t\tassert.NoError(t, g.AddEdge(compose.START, \"1\"))\n\t\tassert.NoError(t, g.AddEdge(\"1\", compose.END))\n\n\t\tctx, _ = initRunCtx(ctx, \"react-agent\", nil)\n\t\trunner, err := g.Compile(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// 4. Run the graph which will execute the tool with the state\n\t\t_, err = runner.Invoke(ctx, \"\")\n\t\tassert.NoError(t, err)\n\n\t\t// 5. Assert that the agent received the full history\n\t\t// The agent should receive: history (minus last assistant message) + transfer messages\n\t\tassert.Len(t, mockAgent.capturedInput, 4) // 2 from history + 2 transfer messages\n\t\tassert.Equal(t, \"first user message\", mockAgent.capturedInput[0].Content)\n\t\tassert.Equal(t, \"For context: [react-agent] said: first assistant response.\", mockAgent.capturedInput[1].Content)\n\t\tassert.Equal(t, \"For context: [react-agent] called tool: `transfer_to_agent` with arguments: test-agent.\", mockAgent.capturedInput[2].Content)\n\t\tassert.Equal(t, \"For context: [react-agent] `transfer_to_agent` tool returned result: successfully transferred to agent [test-agent].\", mockAgent.capturedInput[3].Content)\n\t})\n\n\t// Test Case 2: WithAgentInputSchema\n\tt.Run(\"WithAgentInputSchema\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// 1. Define a custom schema\n\t\tcustomSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"custom_arg\": {\n\t\t\t\tDesc:     \"a custom argument\",\n\t\t\t\tRequired: true,\n\t\t\t\tType:     schema.String,\n\t\t\t},\n\t\t})\n\n\t\t// 2. Set up a mock agent to capture input\n\t\tmockAgent := newMockAgentWithInputCapture(\"schema-agent\", \"agent with custom schema\", []*AgentEvent{\n\t\t\t{\n\t\t\t\tAgentName: \"schema-agent\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\t\tMessage:     schema.AssistantMessage(\"schema processed\", nil),\n\t\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\t// 3. Create agentTool with the custom schema option\n\t\tagentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(customSchema))\n\n\t\t// 4. Verify the Info() method returns the custom schema\n\t\tinfo, err := agentTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, customSchema, info.ParamsOneOf)\n\n\t\t// 5. Run the tool with arguments matching the custom schema\n\t\t_, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{\"custom_arg\":\"hello world\"}`)\n\t\tassert.NoError(t, err)\n\n\t\t// 6. Assert that the agent received the correctly parsed argument\n\t\t// With custom schema, the agent should receive the raw JSON as input\n\t\tassert.Len(t, mockAgent.capturedInput, 1)\n\t\tassert.Equal(t, `{\"custom_arg\":\"hello world\"}`, mockAgent.capturedInput[0].Content)\n\t})\n\n\t// Test Case 3: WithAgentInputSchema with complex schema\n\tt.Run(\"WithAgentInputSchema_ComplexSchema\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// 1. Define a complex custom schema with multiple parameters\n\t\tcomplexSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"name\": {\n\t\t\t\tDesc:     \"user name\",\n\t\t\t\tRequired: true,\n\t\t\t\tType:     schema.String,\n\t\t\t},\n\t\t\t\"age\": {\n\t\t\t\tDesc:     \"user age\",\n\t\t\t\tRequired: false,\n\t\t\t\tType:     schema.Integer,\n\t\t\t},\n\t\t\t\"active\": {\n\t\t\t\tDesc:     \"user status\",\n\t\t\t\tRequired: false,\n\t\t\t\tType:     schema.Boolean,\n\t\t\t},\n\t\t})\n\n\t\t// 2. Set up a mock agent\n\t\tmockAgent := newMockAgentWithInputCapture(\"complex-agent\", \"agent with complex schema\", []*AgentEvent{\n\t\t\t{\n\t\t\t\tAgentName: \"complex-agent\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\t\tMessage:     schema.AssistantMessage(\"complex processed\", nil),\n\t\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\t// 3. Create agentTool with the complex schema option\n\t\tagentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(complexSchema))\n\n\t\t// 4. Verify the Info() method returns the complex schema\n\t\tinfo, err := agentTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, complexSchema, info.ParamsOneOf)\n\n\t\t// 5. Run the tool with complex arguments\n\t\t_, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{\"name\":\"John\",\"age\":30,\"active\":true}`)\n\t\tassert.NoError(t, err)\n\n\t\t// 6. Assert that the agent received the complex JSON\n\t\tassert.Len(t, mockAgent.capturedInput, 1)\n\t\tassert.Equal(t, `{\"name\":\"John\",\"age\":30,\"active\":true}`, mockAgent.capturedInput[0].Content)\n\t})\n\n\t// Test Case 4: Both options together\n\tt.Run(\"BothOptionsTogether\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// 1. Define a custom schema\n\t\tcustomSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"query\": {\n\t\t\t\tDesc:     \"search query\",\n\t\t\t\tRequired: true,\n\t\t\t\tType:     schema.String,\n\t\t\t},\n\t\t})\n\n\t\t// 2. Set up a mock agent\n\t\tmockAgent := newMockAgentWithInputCapture(\"combined-agent\", \"agent with both options\", []*AgentEvent{\n\t\t\t{\n\t\t\t\tAgentName: \"combined-agent\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\t\tMessage:     schema.AssistantMessage(\"combined processed\", nil),\n\t\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\t// 3. Create agentTool with both options\n\t\tagentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(customSchema), WithFullChatHistoryAsInput())\n\n\t\t// 4. Set up a context with chat history using a graph\n\t\thistory := []Message{\n\t\t\tschema.UserMessage(\"previous conversation\"),\n\t\t\tschema.AssistantMessage(\"previous response\", nil),\n\t\t}\n\n\t\tg := compose.NewGraph[string, string](compose.WithGenLocalState(func(ctx context.Context) (state *State) {\n\t\t\treturn &State{\n\t\t\t\tMessages: append(history, schema.AssistantMessage(\"tool call\", nil)),\n\t\t\t}\n\t\t}))\n\n\t\tassert.NoError(t, g.AddLambdaNode(\"1\", compose.InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\t// Run the tool within the graph context that has the state\n\t\t\t_, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{\"query\":\"current query\"}`)\n\t\t\treturn \"done\", err\n\t\t})))\n\t\tassert.NoError(t, g.AddEdge(compose.START, \"1\"))\n\t\tassert.NoError(t, g.AddEdge(\"1\", compose.END))\n\n\t\tctx, _ = initRunCtx(ctx, \"react-agent\", nil)\n\t\trunner, err := g.Compile(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// 5. Run the graph which will execute the tool with the state\n\t\t_, err = runner.Invoke(ctx, \"\")\n\t\tassert.NoError(t, err)\n\n\t\t// 6. Verify both options work together\n\t\tinfo, err := agentTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, customSchema, info.ParamsOneOf)\n\n\t\t// The agent should receive full history + the custom query\n\t\tassert.Len(t, mockAgent.capturedInput, 4) // 2 history + 2 transfer messages\n\t\tassert.Equal(t, \"previous conversation\", mockAgent.capturedInput[0].Content)\n\t\tassert.Equal(t, \"For context: [react-agent] said: previous response.\", mockAgent.capturedInput[1].Content)\n\t\tassert.Equal(t, \"For context: [react-agent] called tool: `transfer_to_agent` with arguments: combined-agent.\", mockAgent.capturedInput[2].Content)\n\t\tassert.Equal(t, \"For context: [react-agent] `transfer_to_agent` tool returned result: successfully transferred to agent [combined-agent].\", mockAgent.capturedInput[3].Content)\n\t})\n}\n\ntype fakeTCM struct{}\n\nfunc (f *fakeTCM) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\to := model.GetCommonOptions(&model.Options{}, opts...)\n\ttc := schema.ToolCall{ID: \"id-1\", Type: \"function\"}\n\tif len(o.Tools) > 0 {\n\t\ttc.Function.Name = o.Tools[0].Name\n\t}\n\ttc.Function.Arguments = `{\"request\":\"hello\"}`\n\treturn schema.AssistantMessage(\"\", []schema.ToolCall{tc}), nil\n}\nfunc (f *fakeTCM) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tmsg, _ := f.Generate(ctx, input, opts...)\n\treturn schema.StreamReaderFromArray([]*schema.Message{msg}), nil\n}\nfunc (f *fakeTCM) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\treturn f, nil\n}\n\ntype emitOnceModel struct{}\n\nfunc (e *emitOnceModel) Generate(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) {\n\treturn schema.AssistantMessage(\"inner2\", nil), nil\n}\nfunc (e *emitOnceModel) Stream(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tm, _ := e.Generate(ctx, input)\n\treturn schema.StreamReaderFromArray([]*schema.Message{m}), nil\n}\nfunc (e *emitOnceModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\treturn e, nil\n}\n\ntype emitEventsAgent struct{ events []*AgentEvent }\n\nfunc (e *emitEventsAgent) Name(context.Context) string        { return \"emit\" }\nfunc (e *emitEventsAgent) Description(context.Context) string { return \"test\" }\nfunc (e *emitEventsAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tit, gen := NewAsyncIteratorPair[*AgentEvent]()\n\tgo func() {\n\t\tfor _, ev := range e.events {\n\t\t\tgen.Send(ev)\n\t\t}\n\t\tgen.Close()\n\t}()\n\treturn it\n}\n\n// spyAgent captures runSession from ctx in a single nested run\ntype spyAgent struct {\n\ta        Agent\n\tmu       sync.Mutex\n\tcaptured *runSession\n}\n\nfunc (s *spyAgent) Name(ctx context.Context) string        { return s.a.Name(ctx) }\nfunc (s *spyAgent) Description(ctx context.Context) string { return s.a.Description(ctx) }\nfunc (s *spyAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tif rc := getRunCtx(ctx); rc != nil {\n\t\ts.mu.Lock()\n\t\ts.captured = rc.Session\n\t\ts.mu.Unlock()\n\t}\n\treturn s.a.Run(ctx, input, options...)\n}\n\nfunc (s *spyAgent) getCaptured() *runSession {\n\ts.mu.Lock()\n\tdefer s.mu.Unlock()\n\treturn s.captured\n}\n\nfunc TestNestedAgentTool_RunPath(t *testing.T) {\n\tctx := context.Background()\n\n\tinner2, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"inner2\",\n\t\tDescription: \"leaf\",\n\t\tModel:       &emitOnceModel{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: true},\n\t})\n\tinner2Spy := &spyAgent{a: inner2}\n\tinner2Tool := NewAgentTool(ctx, inner2Spy)\n\n\tinner, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"inner\",\n\t\tDescription: \"mid\",\n\t\tModel:       &fakeTCM{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: true, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{inner2Tool}}},\n\t})\n\tinnerSpy := &spyAgent{a: inner}\n\tinnerTool := NewAgentTool(ctx, innerSpy)\n\n\touter, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"outer\",\n\t\tDescription: \"top\",\n\t\tModel:       &fakeTCM{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: true, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{innerTool}}},\n\t})\n\n\tinput := &AgentInput{Messages: []Message{schema.UserMessage(\"q\")}}\n\tctx, outerRunCtx := initRunCtx(ctx, \"outer\", input)\n\tr := NewRunner(ctx, RunnerConfig{Agent: outer, EnableStreaming: false, CheckPointStore: newBridgeStore()})\n\tit := r.Run(ctx, []Message{schema.UserMessage(\"q\")})\n\n\tvar target *AgentEvent\n\tfor {\n\t\tev, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif ev.Output != nil && ev.Output.MessageOutput != nil && !ev.Output.MessageOutput.IsStreaming {\n\t\t\tif ev.Output.MessageOutput.Message != nil && ev.Output.MessageOutput.Message.Content == \"inner2\" {\n\t\t\t\ttarget = ev\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\tif target == nil {\n\t\tt.Fatalf(\"no inner2 event found in ephemerals\")\n\t}\n\n\tgot := make([]string, len(target.RunPath))\n\tfor i := range target.RunPath {\n\t\tgot[i] = target.RunPath[i].agentName\n\t}\n\twant := []string{\"outer\", \"inner\", \"inner2\"}\n\tif len(got) != len(want) {\n\t\tt.Fatalf(\"unexpected runPath len: got %d want %d: %+v\", len(got), len(want), got)\n\t}\n\tfor i := range want {\n\t\tif got[i] != want[i] {\n\t\t\tt.Fatalf(\"runPath mismatch at %d: got %s want %s; full: %+v\", i, got[i], want[i], got)\n\t\t}\n\t}\n\n\tfor _, w := range outerRunCtx.Session.getEvents() {\n\t\tif w.AgentName != \"outer\" {\n\t\t\tt.Fatalf(\"outer session contains non-outer event: %s\", w.AgentName)\n\t\t}\n\t}\n\tif innerSpy.getCaptured() == nil {\n\t\tt.Fatalf(\"inner spy did not capture session\")\n\t}\n\tfor _, w := range innerSpy.getCaptured().getEvents() {\n\t\tif w.AgentName != \"inner\" {\n\t\t\tt.Fatalf(\"inner session contains non-inner event: %s\", w.AgentName)\n\t\t}\n\t}\n\tif inner2Spy.getCaptured() == nil {\n\t\tt.Fatalf(\"inner2 spy did not capture session\")\n\t}\n\tfor _, w := range inner2Spy.getCaptured().getEvents() {\n\t\tif w.AgentName != \"inner2\" {\n\t\t\tt.Fatalf(\"inner2 session contains non-inner2 event: %s\", w.AgentName)\n\t\t}\n\t}\n}\n\nfunc TestNestedAgentTool_NoInternalEventsWhenDisabled(t *testing.T) {\n\tctx := context.Background()\n\n\tinner2, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"inner2\",\n\t\tDescription: \"leaf\",\n\t\tModel:       &emitOnceModel{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: false},\n\t})\n\tinner2Tool := NewAgentTool(ctx, inner2)\n\n\tinner, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"inner\",\n\t\tDescription: \"mid\",\n\t\tModel:       &fakeTCM{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: false, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{inner2Tool}}},\n\t})\n\tinnerTool := NewAgentTool(ctx, inner)\n\n\touter, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"outer\",\n\t\tDescription: \"top\",\n\t\tModel:       &fakeTCM{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: false, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{innerTool}}},\n\t})\n\n\tr := NewRunner(ctx, RunnerConfig{Agent: outer, EnableStreaming: false, CheckPointStore: newBridgeStore()})\n\tit := r.Run(ctx, []Message{schema.UserMessage(\"q\")})\n\n\tfor {\n\t\tev, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif ev.AgentName == \"inner2\" {\n\t\t\tt.Fatalf(\"inner2 internal event should not be emitted when disabled\")\n\t\t}\n\t}\n}\n\nfunc TestNestedAgentTool_InnerToolResultNotEmittedToOuter(t *testing.T) {\n\tctx := context.Background()\n\n\tinnerTool := &simpleTool{name: \"inner_tool\", result: \"inner_tool_result\"}\n\tinner, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"inner\",\n\t\tDescription: \"inner agent with tool\",\n\t\tModel:       &fakeTCM{},\n\t\tToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{innerTool}}},\n\t})\n\tinnerAgentTool := NewAgentTool(ctx, inner)\n\n\touter, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"outer\",\n\t\tDescription: \"outer agent\",\n\t\tModel:       &fakeTCM{},\n\t\tToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{innerAgentTool}}},\n\t})\n\n\tr := NewRunner(ctx, RunnerConfig{Agent: outer, EnableStreaming: false, CheckPointStore: newBridgeStore()})\n\tit := r.Run(ctx, []Message{schema.UserMessage(\"q\")})\n\n\tvar allEvents []*AgentEvent\n\tfor {\n\t\tev, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tallEvents = append(allEvents, ev)\n\t}\n\n\tfor _, ev := range allEvents {\n\t\tif ev.Output != nil && ev.Output.MessageOutput != nil &&\n\t\t\tev.Output.MessageOutput.Message != nil &&\n\t\t\tev.Output.MessageOutput.Message.Role == schema.Tool &&\n\t\t\tev.AgentName == \"outer\" &&\n\t\t\tev.Output.MessageOutput.Message.Content == \"inner_tool_result\" {\n\t\t\tt.Fatalf(\"inner agent's tool result (inner_tool_result) should not be emitted as outer agent's event, but got event with AgentName=%s, Content=%s\",\n\t\t\t\tev.AgentName, ev.Output.MessageOutput.Message.Content)\n\t\t}\n\t}\n}\n\ntype simpleTool struct {\n\tname   string\n\tresult string\n}\n\nfunc (s *simpleTool) Info(context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: s.name, Desc: \"simple tool\"}, nil\n}\n\nfunc (s *simpleTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\treturn s.result, nil\n}\n\nfunc TestAgentTool_InterruptWithoutCheckpoint(t *testing.T) {\n\tctx := context.Background()\n\tctx, _ = initRunCtx(ctx, \"TestAgent\", &AgentInput{Messages: []Message{}})\n\n\tinterrupted := &AgentEvent{AgentName: \"TestAgent\"}\n\tinterrupted.Action = StatefulInterrupt(ctx, \"info\", \"state\").Action\n\n\terr := compositeInterruptFromLast(ctx, &bridgeStore{}, interrupted)\n\tif err == nil {\n\t\tt.Fatalf(\"expected error for interrupt without checkpoint\")\n\t}\n\tif !strings.Contains(err.Error(), \"interrupt occurred but checkpoint data is missing\") {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n}\n\nfunc compositeInterruptFromLast(ctx context.Context, ms *bridgeStore, lastEvent *AgentEvent) error {\n\tif lastEvent == nil || lastEvent.Action == nil || lastEvent.Action.Interrupted == nil {\n\t\treturn nil\n\t}\n\tdata, existed, err := ms.Get(ctx, bridgeCheckpointID)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to get interrupt info: %w\", err)\n\t}\n\tif !existed {\n\t\treturn fmt.Errorf(\"interrupt occurred but checkpoint data is missing\")\n\t}\n\treturn tool.CompositeInterrupt(ctx, \"agent tool interrupt\", data, lastEvent.Action.internalInterrupted)\n}\n\nfunc TestAgentTool_InvokableRun_FinalOnly(t *testing.T) {\n\tctx := context.Background()\n\n\tinner2, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"inner2\",\n\t\tDescription: \"leaf\",\n\t\tModel:       &emitOnceModel{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: true},\n\t})\n\tinvTool := NewAgentTool(ctx, inner2)\n\tout, err := invTool.(tool.InvokableTool).InvokableRun(ctx, `{\"request\":\"q\"}`)\n\tif err != nil {\n\t\tt.Fatalf(\"invokable run error: %v\", err)\n\t}\n\tif out != \"inner2\" {\n\t\tt.Fatalf(\"unexpected output: %s\", out)\n\t}\n}\n\ntype streamingAgent struct{}\n\nfunc (s *streamingAgent) Name(context.Context) string        { return \"stream\" }\nfunc (s *streamingAgent) Description(context.Context) string { return \"test\" }\nfunc (s *streamingAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tit, gen := NewAsyncIteratorPair[*AgentEvent]()\n\tgo func() {\n\t\tmv := &MessageVariant{IsStreaming: true, MessageStream: schema.StreamReaderFromArray([]Message{schema.AssistantMessage(\"1\", nil), schema.AssistantMessage(\"2\", nil)})}\n\t\tgen.Send(&AgentEvent{AgentName: \"stream\", Output: &AgentOutput{MessageOutput: mv}})\n\t\tmv = &MessageVariant{IsStreaming: true, MessageStream: schema.StreamReaderFromArray([]Message{schema.AssistantMessage(\"a\", nil), schema.AssistantMessage(\"b\", nil)})}\n\t\tgen.Send(&AgentEvent{AgentName: \"stream\", Output: &AgentOutput{MessageOutput: mv}})\n\t\tgen.Close()\n\t}()\n\treturn it\n}\n\nfunc TestAgentTool_InvokableRun_StreamingVariant(t *testing.T) {\n\tctx := context.Background()\n\tagent := &streamingAgent{}\n\tit := NewAgentTool(ctx, agent)\n\tout, err := it.(tool.InvokableTool).InvokableRun(ctx, `{\"request\":\"q\"}`)\n\tif err != nil {\n\t\tt.Fatalf(\"invokable run error: %v\", err)\n\t}\n\tif out != \"ab\" {\n\t\tt.Fatalf(\"unexpected output: %s\", out)\n\t}\n}\n\nfunc TestSequentialWorkflow_WithChatModelAgentTool_NestedRunPathAndSessions(t *testing.T) {\n\tctx := context.Background()\n\n\tinner2, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"inner2\",\n\t\tDescription: \"leaf\",\n\t\tModel:       &emitOnceModel{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: true},\n\t})\n\tinner2Spy := &spyAgent{a: inner2}\n\tinner2ToolSpy := NewAgentTool(ctx, inner2Spy)\n\n\tinnerWithSpy, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"inner\",\n\t\tDescription: \"mid\",\n\t\tModel:       &fakeTCM{},\n\t\tToolsConfig: ToolsConfig{EmitInternalEvents: true, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{inner2ToolSpy}}},\n\t})\n\tinnerSpy := &spyAgent{a: innerWithSpy}\n\n\touter, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\tName:        \"outer-seq\",\n\t\tDescription: \"workflow\",\n\t\tSubAgents:   []Agent{innerSpy},\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"new sequential agent err: %v\", err)\n\t}\n\n\tinput := &AgentInput{Messages: []Message{schema.UserMessage(\"q\")}}\n\tctx, outerRunCtx := initRunCtx(ctx, \"outer-seq\", input)\n\tr := NewRunner(ctx, RunnerConfig{Agent: outer, EnableStreaming: false, CheckPointStore: newBridgeStore()})\n\tit := r.Run(ctx, []Message{schema.UserMessage(\"q\")})\n\n\tvar target *AgentEvent\n\tfor {\n\t\tev, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif ev.Output != nil && ev.Output.MessageOutput != nil && !ev.Output.MessageOutput.IsStreaming {\n\t\t\tif ev.Output.MessageOutput.Message != nil && ev.Output.MessageOutput.Message.Content == \"inner2\" {\n\t\t\t\ttarget = ev\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\tif target == nil {\n\t\tt.Fatalf(\"no inner2 event found\")\n\t}\n\n\tgot := make([]string, len(target.RunPath))\n\tfor i := range target.RunPath {\n\t\tgot[i] = target.RunPath[i].agentName\n\t}\n\twant := []string{\"outer-seq\", \"inner\", \"inner2\"}\n\tif len(got) != len(want) {\n\t\tt.Fatalf(\"unexpected runPath len: got %d want %d: %+v\", len(got), len(want), got)\n\t}\n\tfor i := range want {\n\t\tif got[i] != want[i] {\n\t\t\tt.Fatalf(\"runPath mismatch at %d: got %s want %s; full: %+v\", i, got[i], want[i], got)\n\t\t}\n\t}\n\n\tfor _, w := range outerRunCtx.Session.getEvents() {\n\t\tif w.AgentName != \"outer-seq\" {\n\t\t\tt.Fatalf(\"outer session contains non-outer event: %s\", w.AgentName)\n\t\t}\n\t}\n\tif innerSpy.getCaptured() == nil {\n\t\tt.Fatalf(\"inner spy did not capture session\")\n\t}\n\tfor _, w := range innerSpy.getCaptured().getEvents() {\n\t\tif w.AgentName != \"inner\" {\n\t\t\tt.Fatalf(\"inner session contains non-inner event: %s\", w.AgentName)\n\t\t}\n\t}\n\tif inner2Spy.getCaptured() == nil {\n\t\tt.Fatalf(\"inner2 spy did not capture session\")\n\t}\n\tfor _, w := range inner2Spy.getCaptured().getEvents() {\n\t\tif w.AgentName != \"inner2\" {\n\t\t\tt.Fatalf(\"inner2 session contains non-inner2 event: %s\", w.AgentName)\n\t\t}\n\t}\n}\n\nfunc TestRunPathGating_IgnoresInnerExitAndAllowsOutput(t *testing.T) {\n\tctx := context.Background()\n\n\tinnerExit := &AgentEvent{Action: &AgentAction{Exit: true}, RunPath: []RunStep{{agentName: \"inner\"}}}\n\tfinalOut := EventFromMessage(schema.AssistantMessage(\"ok\", nil), nil, schema.Assistant, \"\")\n\n\tsub := &emitEventsAgent{events: []*AgentEvent{innerExit, finalOut}}\n\tfa := toFlowAgent(ctx, sub)\n\n\tit := fa.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"q\")}})\n\n\tvar sawFinal bool\n\tfor {\n\t\tev, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif ev.Output != nil && ev.Output.MessageOutput != nil && !ev.Output.MessageOutput.IsStreaming {\n\t\t\tif ev.Output.MessageOutput.Message != nil && ev.Output.MessageOutput.Message.Content == \"ok\" {\n\t\t\t\tsawFinal = true\n\t\t\t}\n\t\t}\n\t}\n\tif !sawFinal {\n\t\tt.Fatalf(\"final output not observed; parent may have exited on inner Exit action\")\n\t}\n}\n\nfunc TestRunPathGating_IgnoresInnerTransfer(t *testing.T) {\n\tctx := context.Background()\n\n\tinnerTransfer := &AgentEvent{Action: NewTransferToAgentAction(\"ghost\"), RunPath: []RunStep{{agentName: \"inner\"}}}\n\tfinalOut := EventFromMessage(schema.AssistantMessage(\"done\", nil), nil, schema.Assistant, \"\")\n\n\tsub := &emitEventsAgent{events: []*AgentEvent{innerTransfer, finalOut}}\n\tfa := toFlowAgent(ctx, sub)\n\n\tit := fa.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"q\")}})\n\n\tvar outputs int\n\tfor {\n\t\tev, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif ev.Output != nil && ev.Output.MessageOutput != nil && !ev.Output.MessageOutput.IsStreaming {\n\t\t\tif ev.Output.MessageOutput.Message != nil {\n\t\t\t\toutputs++\n\t\t\t}\n\t\t}\n\t}\n\tif outputs == 0 {\n\t\tt.Fatalf(\"no outputs observed; parent may have transferred on inner transfer action\")\n\t}\n}\n\ntype streamAgent struct{}\n\nfunc (s *streamAgent) Name(context.Context) string        { return \"s\" }\nfunc (s *streamAgent) Description(context.Context) string { return \"s\" }\nfunc (s *streamAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tit, gen := NewAsyncIteratorPair[*AgentEvent]()\n\tgo func() {\n\t\tframes := []*schema.Message{\n\t\t\tschema.AssistantMessage(\"hello \", nil),\n\t\t\tschema.AssistantMessage(\"world\", nil),\n\t\t}\n\t\tstream := schema.StreamReaderFromArray(frames)\n\t\tgen.Send(EventFromMessage(nil, stream, schema.Assistant, \"\"))\n\t\tgen.Close()\n\t}()\n\treturn it\n}\n\nfunc TestInvokableAgentTool_InfoAndRun(t *testing.T) {\n\tctx := context.Background()\n\n\tat := NewAgentTool(ctx, &streamAgent{})\n\tinfo, err := at.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"s\", info.Name)\n\tassert.Equal(t, \"s\", info.Desc)\n\tjs, err := info.ParamsOneOf.ToJSONSchema()\n\tassert.NoError(t, err)\n\tfound := false\n\tfor _, r := range js.Required {\n\t\tif r == \"request\" {\n\t\t\tfound = true\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.True(t, found)\n\tprop, ok := js.Properties.Get(\"request\")\n\tassert.True(t, ok)\n\tassert.Equal(t, string(schema.String), prop.Type)\n\n\tcustom := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\"x\": {Desc: \"arg\", Required: true, Type: schema.String},\n\t})\n\tat2 := NewAgentTool(ctx, &streamAgent{}, WithAgentInputSchema(custom))\n\tinfo2, err := at2.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, custom, info2.ParamsOneOf)\n\tout, err := at.(tool.InvokableTool).InvokableRun(ctx, `{\"request\":\"x\"}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"hello world\", out)\n}\n\ntype emptyAgent struct{}\n\nfunc (e *emptyAgent) Name(context.Context) string        { return \"empty\" }\nfunc (e *emptyAgent) Description(context.Context) string { return \"empty\" }\nfunc (e *emptyAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tit, gen := NewAsyncIteratorPair[*AgentEvent]()\n\tgo func() { gen.Close() }()\n\treturn it\n}\n\ntype noOutputAgent struct{}\n\nfunc (n *noOutputAgent) Name(context.Context) string        { return \"no\" }\nfunc (n *noOutputAgent) Description(context.Context) string { return \"no\" }\nfunc (n *noOutputAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tit, gen := NewAsyncIteratorPair[*AgentEvent]()\n\tgo func() { gen.Send(&AgentEvent{}); gen.Close() }()\n\treturn it\n}\n\nfunc TestInvokableAgentTool_ErrorCases(t *testing.T) {\n\tctx := context.Background()\n\n\tatEmpty := NewAgentTool(ctx, &emptyAgent{})\n\tout, err := atEmpty.(tool.InvokableTool).InvokableRun(ctx, `{\"request\":\"x\"}`)\n\tassert.Equal(t, \"\", out)\n\tassert.Error(t, err)\n\n\tatNo := NewAgentTool(ctx, &noOutputAgent{})\n\tout2, err := atNo.(tool.InvokableTool).InvokableRun(ctx, `{\"request\":\"x\"}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"\", out2)\n}\n"
  },
  {
    "path": "adk/call_option.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport \"github.com/cloudwego/eino/callbacks\"\n\ntype options struct {\n\tsharedParentSession  bool\n\tsessionValues        map[string]any\n\tcheckPointID         *string\n\tskipTransferMessages bool\n\thandlers             []callbacks.Handler\n}\n\n// AgentRunOption is the call option for adk Agent.\ntype AgentRunOption struct {\n\timplSpecificOptFn any\n\n\t// specify which Agent can see this AgentRunOption, if empty, all Agents can see this AgentRunOption\n\tagentNames []string\n}\n\nfunc (o AgentRunOption) DesignateAgent(name ...string) AgentRunOption {\n\to.agentNames = append(o.agentNames, name...)\n\treturn o\n}\n\nfunc getCommonOptions(base *options, opts ...AgentRunOption) *options {\n\tif base == nil {\n\t\tbase = &options{}\n\t}\n\n\treturn GetImplSpecificOptions(base, opts...)\n}\n\n// WithSessionValues sets session-scoped values for the agent run.\nfunc WithSessionValues(v map[string]any) AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(o *options) {\n\t\to.sessionValues = v\n\t})\n}\n\n// WithSkipTransferMessages disables forwarding transfer messages during execution.\nfunc WithSkipTransferMessages() AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(t *options) {\n\t\tt.skipTransferMessages = true\n\t})\n}\n\nfunc withSharedParentSession() AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(o *options) {\n\t\to.sharedParentSession = true\n\t})\n}\n\n// WithCallbacks adds callback handlers to receive agent lifecycle events.\n// Handlers receive OnStart with AgentCallbackInput and OnEnd with AgentCallbackOutput.\n// Multiple handlers can be added; each receives an independent copy of the event stream.\nfunc WithCallbacks(handlers ...callbacks.Handler) AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(o *options) {\n\t\to.handlers = append(o.handlers, handlers...)\n\t})\n}\n\n// WrapImplSpecificOptFn is the option to wrap the implementation specific option function.\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) AgentRunOption {\n\treturn AgentRunOption{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetImplSpecificOptions extract the implementation specific options from AgentRunOption list, optionally providing a base options with default values.\n// e.g.\n//\n//\tmyOption := &MyOption{\n//\t\tField1: \"default_value\",\n//\t}\n//\n//\tmyOption := model.GetImplSpecificOptions(myOption, opts...)\nfunc GetImplSpecificOptions[T any](base *T, opts ...AgentRunOption) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\toptFn, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\toptFn(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n\n// filterCallbackHandlersForNestedAgents removes callback handlers that have already been applied\n// to the current agent before passing opts to nested inner agents.\n//\n// This is necessary for workflow agents (LoopAgent, SequentialAgent, ParallelAgent) because:\n//  1. Callback handlers designated for the current agent are applied via initAgentCallbacks(),\n//     which stores them in the context.\n//  2. Nested inner agents inherit this context, so they automatically receive these callbacks.\n//  3. If we also pass these handlers in opts to inner agents, they would be applied twice,\n//     causing duplicate callback invocations.\n//\n// Note: This only applies to workflow agents where inner agents inherit context from the parent.\n// For flowAgent's sub-agents (which are peer agents that transfer to each other), the full opts\n// are passed since they don't inherit the parent's callback context.\nfunc filterCallbackHandlersForNestedAgents(currentAgentName string, opts []AgentRunOption) []AgentRunOption {\n\tif len(opts) == 0 {\n\t\treturn nil\n\t}\n\tvar filteredOpts []AgentRunOption\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn == nil {\n\t\t\tfilteredOpts = append(filteredOpts, opt)\n\t\t\tcontinue\n\t\t}\n\t\tif _, isCallbackOpt := opt.implSpecificOptFn.(func(*options)); isCallbackOpt {\n\t\t\ttestOpt := &options{}\n\t\t\topt.implSpecificOptFn.(func(*options))(testOpt)\n\t\t\tif len(testOpt.handlers) > 0 {\n\t\t\t\tif len(opt.agentNames) == 0 {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\t\t\t\tmatched := false\n\t\t\t\tfor _, name := range opt.agentNames {\n\t\t\t\t\tif name == currentAgentName {\n\t\t\t\t\t\tmatched = true\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\tif matched {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t\tfilteredOpts = append(filteredOpts, opt)\n\t}\n\treturn filteredOpts\n}\n\nfunc filterOptions(agentName string, opts []AgentRunOption) []AgentRunOption {\n\tif len(opts) == 0 {\n\t\treturn nil\n\t}\n\tvar filteredOpts []AgentRunOption\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif len(opt.agentNames) == 0 {\n\t\t\tfilteredOpts = append(filteredOpts, opt)\n\t\t\tcontinue\n\t\t}\n\t\tfor j := range opt.agentNames {\n\t\t\tif opt.agentNames[j] == agentName {\n\t\t\t\tfilteredOpts = append(filteredOpts, opt)\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\treturn filteredOpts\n}\n"
  },
  {
    "path": "adk/call_option_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n)\n\ntype mockAgentForOption struct {\n\topts []AgentRunOption\n\n\toptions *options\n}\n\nfunc (m *mockAgentForOption) Name(ctx context.Context) string {\n\treturn \"agent_1\"\n}\n\nfunc (m *mockAgentForOption) Description(ctx context.Context) string {\n\treturn \"\"\n}\n\nfunc (m *mockAgentForOption) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tm.opts = opts\n\tm.options = getCommonOptions(&options{}, opts...)\n\n\treturn nil\n}\n"
  },
  {
    "path": "adk/callback.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components\"\n\ticb \"github.com/cloudwego/eino/internal/callbacks\"\n)\n\n// AgentCallbackInput represents the input passed to agent callbacks during OnStart.\n// Use ConvAgentCallbackInput to safely convert from callbacks.CallbackInput.\ntype AgentCallbackInput struct {\n\t// Input contains the agent input for a new run. Nil when resuming.\n\tInput *AgentInput\n\t// ResumeInfo contains resume information when resuming from an interrupt. Nil for new runs.\n\tResumeInfo *ResumeInfo\n}\n\n// AgentCallbackOutput represents the output passed to agent callbacks during OnEnd.\n// Use ConvAgentCallbackOutput to safely convert from callbacks.CallbackOutput.\n//\n// Important: The Events iterator should be consumed asynchronously to avoid blocking\n// the agent execution. Each callback handler receives an independent copy of the iterator.\ntype AgentCallbackOutput struct {\n\t// Events provides a stream of agent events. Each handler receives its own copy.\n\tEvents *AsyncIterator[*AgentEvent]\n}\n\nfunc copyEventIterator(iter *AsyncIterator[*AgentEvent], n int) []*AsyncIterator[*AgentEvent] {\n\tif n <= 0 {\n\t\treturn nil\n\t}\n\tif n == 1 {\n\t\treturn []*AsyncIterator[*AgentEvent]{iter}\n\t}\n\n\titerators := make([]*AsyncIterator[*AgentEvent], n)\n\tgenerators := make([]*AsyncGenerator[*AgentEvent], n)\n\tfor i := 0; i < n; i++ {\n\t\titerators[i], generators[i] = NewAsyncIteratorPair[*AgentEvent]()\n\t}\n\n\tgo func() {\n\t\tdefer func() {\n\t\t\tfor _, g := range generators {\n\t\t\t\tg.Close()\n\t\t\t}\n\t\t}()\n\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tfor i := 0; i < n-1; i++ {\n\t\t\t\tgenerators[i].Send(copyAgentEvent(event))\n\t\t\t}\n\t\t\tgenerators[n-1].Send(event)\n\t\t}\n\t}()\n\n\treturn iterators\n}\n\nfunc copyAgentCallbackOutput(out *AgentCallbackOutput, n int) []*AgentCallbackOutput {\n\tif out == nil || out.Events == nil {\n\t\tresult := make([]*AgentCallbackOutput, n)\n\t\tfor i := 0; i < n; i++ {\n\t\t\tresult[i] = out\n\t\t}\n\t\treturn result\n\t}\n\titers := copyEventIterator(out.Events, n)\n\tresult := make([]*AgentCallbackOutput, n)\n\tfor i, iter := range iters {\n\t\tresult[i] = &AgentCallbackOutput{Events: iter}\n\t}\n\treturn result\n}\n\n// ConvAgentCallbackInput converts a generic CallbackInput to AgentCallbackInput.\n// Returns nil if the input is not an AgentCallbackInput.\nfunc ConvAgentCallbackInput(input callbacks.CallbackInput) *AgentCallbackInput {\n\tif v, ok := input.(*AgentCallbackInput); ok {\n\t\treturn v\n\t}\n\treturn nil\n}\n\n// ConvAgentCallbackOutput converts a generic CallbackOutput to AgentCallbackOutput.\n// Returns nil if the output is not an AgentCallbackOutput.\nfunc ConvAgentCallbackOutput(output callbacks.CallbackOutput) *AgentCallbackOutput {\n\tif v, ok := output.(*AgentCallbackOutput); ok {\n\t\treturn v\n\t}\n\treturn nil\n}\n\nfunc initAgentCallbacks(ctx context.Context, agentName, agentType string, opts ...AgentRunOption) context.Context {\n\tri := &callbacks.RunInfo{\n\t\tName:      agentName,\n\t\tType:      agentType,\n\t\tComponent: ComponentOfAgent,\n\t}\n\n\to := getCommonOptions(nil, opts...)\n\tif len(o.handlers) == 0 {\n\t\treturn icb.ReuseHandlers(ctx, ri)\n\t}\n\treturn icb.AppendHandlers(ctx, ri, o.handlers...)\n}\n\nfunc getAgentType(agent Agent) string {\n\tif typer, ok := agent.(components.Typer); ok {\n\t\treturn typer.GetType()\n\t}\n\treturn \"\"\n}\n"
  },
  {
    "path": "adk/callback_integration_test.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype callbackRecorder struct {\n\tmu             sync.Mutex\n\tonStartCalled  bool\n\tonEndCalled    bool\n\trunInfo        *callbacks.RunInfo\n\tinputReceived  *AgentCallbackInput\n\teventsReceived []*AgentEvent\n\teventsDone     chan struct{}\n\tcloseOnce      sync.Once\n}\n\nfunc (r *callbackRecorder) getOnStartCalled() bool {\n\tr.mu.Lock()\n\tdefer r.mu.Unlock()\n\treturn r.onStartCalled\n}\n\nfunc (r *callbackRecorder) getOnEndCalled() bool {\n\tr.mu.Lock()\n\tdefer r.mu.Unlock()\n\treturn r.onEndCalled\n}\n\nfunc (r *callbackRecorder) getEventsReceived() []*AgentEvent {\n\tr.mu.Lock()\n\tdefer r.mu.Unlock()\n\tresult := make([]*AgentEvent, len(r.eventsReceived))\n\tcopy(result, r.eventsReceived)\n\treturn result\n}\n\nfunc newRecordingHandler(recorder *callbackRecorder) callbacks.Handler {\n\trecorder.eventsDone = make(chan struct{})\n\treturn callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\trecorder.mu.Lock()\n\t\t\tdefer recorder.mu.Unlock()\n\t\t\trecorder.onStartCalled = true\n\t\t\trecorder.runInfo = info\n\t\t\tif agentInput := ConvAgentCallbackInput(input); agentInput != nil {\n\t\t\t\trecorder.inputReceived = agentInput\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\trecorder.mu.Lock()\n\t\t\trecorder.onEndCalled = true\n\t\t\trecorder.runInfo = info\n\t\t\trecorder.mu.Unlock()\n\n\t\t\tif agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil {\n\t\t\t\tif agentOutput.Events != nil {\n\t\t\t\t\tgo func() {\n\t\t\t\t\t\tdefer recorder.closeOnce.Do(func() { close(recorder.eventsDone) })\n\t\t\t\t\t\tfor {\n\t\t\t\t\t\t\tevent, ok := agentOutput.Events.Next()\n\t\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\trecorder.mu.Lock()\n\t\t\t\t\t\t\trecorder.eventsReceived = append(recorder.eventsReceived, event)\n\t\t\t\t\t\t\trecorder.mu.Unlock()\n\t\t\t\t\t\t}\n\t\t\t\t\t}()\n\t\t\t\t\treturn ctx\n\t\t\t\t}\n\t\t\t}\n\t\t\trecorder.closeOnce.Do(func() { close(recorder.eventsDone) })\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n}\n\nfunc TestCallbackOnStartInvocation(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"test response\", nil), nil).\n\t\tTimes(1)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Test agent for callback\",\n\t\tInstruction: \"You are a test agent\",\n\t\tModel:       cm,\n\t})\n\tassert.NoError(t, err)\n\n\trecorder := &callbackRecorder{}\n\thandler := newRecordingHandler(recorder)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agent})\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\t<-recorder.eventsDone\n\n\tassert.True(t, recorder.onStartCalled, \"OnStart should be called\")\n\tassert.NotNil(t, recorder.inputReceived, \"Input should be received\")\n\tassert.NotNil(t, recorder.inputReceived.Input, \"AgentInput should be set\")\n\tassert.Len(t, recorder.inputReceived.Input.Messages, 1)\n}\n\nfunc TestCallbackOnEndInvocation(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"test response\", nil), nil).\n\t\tTimes(1)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Test agent for callback\",\n\t\tInstruction: \"You are a test agent\",\n\t\tModel:       cm,\n\t})\n\tassert.NoError(t, err)\n\n\trecorder := &callbackRecorder{}\n\thandler := newRecordingHandler(recorder)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agent})\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\t<-recorder.eventsDone\n\n\tassert.True(t, recorder.onEndCalled, \"OnEnd should be called\")\n\tassert.NotEmpty(t, recorder.eventsReceived, \"Events should be received\")\n}\n\nfunc TestCallbackRunInfoForChatModelAgent(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"test response\", nil), nil).\n\t\tTimes(1)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestChatAgent\",\n\t\tDescription: \"Test chat agent\",\n\t\tInstruction: \"You are a test agent\",\n\t\tModel:       cm,\n\t})\n\tassert.NoError(t, err)\n\n\trecorder := &callbackRecorder{}\n\thandler := newRecordingHandler(recorder)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agent})\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\t<-recorder.eventsDone\n\n\tassert.NotNil(t, recorder.runInfo)\n\tassert.Equal(t, \"TestChatAgent\", recorder.runInfo.Name)\n\tassert.Equal(t, \"ChatModel\", recorder.runInfo.Type)\n\tassert.Equal(t, ComponentOfAgent, recorder.runInfo.Component)\n}\n\nfunc TestMultipleCallbackHandlers(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"test response\", nil), nil).\n\t\tTimes(1)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Test agent\",\n\t\tInstruction: \"You are a test agent\",\n\t\tModel:       cm,\n\t})\n\tassert.NoError(t, err)\n\n\trecorder1 := &callbackRecorder{}\n\trecorder2 := &callbackRecorder{}\n\thandler1 := newRecordingHandler(recorder1)\n\thandler2 := newRecordingHandler(recorder2)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agent})\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler1, handler2))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\t<-recorder1.eventsDone\n\t<-recorder2.eventsDone\n\n\tassert.True(t, recorder1.onStartCalled, \"Handler1 OnStart should be called\")\n\tassert.True(t, recorder2.onStartCalled, \"Handler2 OnStart should be called\")\n\tassert.True(t, recorder1.onEndCalled, \"Handler1 OnEnd should be called\")\n\tassert.True(t, recorder2.onEndCalled, \"Handler2 OnEnd should be called\")\n\n\tassert.NotEmpty(t, recorder1.eventsReceived, \"Handler1 should receive events\")\n\tassert.NotEmpty(t, recorder2.eventsReceived, \"Handler2 should receive events\")\n}\n\nfunc TestCallbackWithWorkflowAgent(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm1 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"response 1\", nil), nil).\n\t\tTimes(1)\n\tcm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes()\n\n\tcm2 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"response 2\", nil), nil).\n\t\tTimes(1)\n\tcm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes()\n\n\tagent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent1\",\n\t\tDescription: \"First agent\",\n\t\tInstruction: \"You are agent 1\",\n\t\tModel:       cm1,\n\t})\n\tassert.NoError(t, err)\n\n\tagent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent2\",\n\t\tDescription: \"Second agent\",\n\t\tInstruction: \"You are agent 2\",\n\t\tModel:       cm2,\n\t})\n\tassert.NoError(t, err)\n\n\tseqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\tName:        \"SequentialAgent\",\n\t\tDescription: \"Sequential workflow\",\n\t\tSubAgents:   []Agent{agent1, agent2},\n\t})\n\tassert.NoError(t, err)\n\n\tvar callbackInfos []*callbacks.RunInfo\n\thandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component == ComponentOfAgent {\n\t\t\t\tcallbackInfos = append(callbackInfos, info)\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: seqAgent})\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassert.NotEmpty(t, callbackInfos, \"OnStart should be called for agents\")\n\tfoundAgent1 := false\n\tfoundAgent2 := false\n\tfor _, info := range callbackInfos {\n\t\tif info.Name == \"Agent1\" && info.Type == \"ChatModel\" {\n\t\t\tfoundAgent1 = true\n\t\t}\n\t\tif info.Name == \"Agent2\" && info.Type == \"ChatModel\" {\n\t\t\tfoundAgent2 = true\n\t\t}\n\t}\n\tassert.True(t, foundAgent1, \"Agent1 callback should be invoked\")\n\tassert.True(t, foundAgent2, \"Agent2 callback should be invoked\")\n}\n\nfunc TestCallbackEventsMatchAgentOutput(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\texpectedContent := \"This is the test response content\"\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(expectedContent, nil), nil).\n\t\tTimes(1)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Test agent\",\n\t\tInstruction: \"You are a test agent\",\n\t\tModel:       cm,\n\t})\n\tassert.NoError(t, err)\n\n\trecorder := &callbackRecorder{}\n\thandler := newRecordingHandler(recorder)\n\n\tvar agentEvents []*AgentEvent\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agent})\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler))\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tagentEvents = append(agentEvents, event)\n\t}\n\n\t<-recorder.eventsDone\n\n\tassert.NotEmpty(t, agentEvents, \"Agent should emit events\")\n\tassert.NotEmpty(t, recorder.eventsReceived, \"Callback should receive events\")\n\n\tfoundExpectedContent := false\n\tfor _, event := range recorder.eventsReceived {\n\t\tif event.Output != nil && event.Output.MessageOutput != nil {\n\t\t\tmsg := event.Output.MessageOutput.Message\n\t\t\tif msg != nil && msg.Content == expectedContent {\n\t\t\t\tfoundExpectedContent = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\tassert.True(t, foundExpectedContent, \"Callback events should contain the expected content\")\n}\n\nfunc TestCallbackOnEndForWorkflowAgent(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm1 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"response 1\", nil), nil).\n\t\tTimes(1)\n\tcm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes()\n\n\tcm2 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"response 2\", nil), nil).\n\t\tTimes(1)\n\tcm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes()\n\n\tagent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent1\",\n\t\tDescription: \"First agent\",\n\t\tInstruction: \"You are agent 1\",\n\t\tModel:       cm1,\n\t})\n\tassert.NoError(t, err)\n\n\tagent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent2\",\n\t\tDescription: \"Second agent\",\n\t\tInstruction: \"You are agent 2\",\n\t\tModel:       cm2,\n\t})\n\tassert.NoError(t, err)\n\n\tseqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\tName:        \"SequentialAgent\",\n\t\tDescription: \"Sequential workflow\",\n\t\tSubAgents:   []Agent{agent1, agent2},\n\t})\n\tassert.NoError(t, err)\n\n\trecorder := &callbackRecorder{}\n\thandler := newRecordingHandler(recorder)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: seqAgent})\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\t<-recorder.eventsDone\n\n\tassert.True(t, recorder.getOnStartCalled(), \"OnStart should be called for workflow agent\")\n\tassert.True(t, recorder.getOnEndCalled(), \"OnEnd should be called for workflow agent\")\n\tassert.NotEmpty(t, recorder.getEventsReceived(), \"Events should be received for workflow agent\")\n}\n\ntype ctxKeyForTest string\n\nconst testOnStartMarkerKey ctxKeyForTest = \"onStartMarker\"\n\nfunc TestSubAgentContextIsolation(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm1 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"transferring to Agent2\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"transfer_1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      TransferToAgentToolName,\n\t\t\t\t\t\tArguments: `{\"agent_name\": \"Agent2\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\tcm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes()\n\n\tcm2 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"final response from Agent2\", nil), nil).\n\t\tTimes(1)\n\tcm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes()\n\n\tagent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent1\",\n\t\tDescription: \"First agent that transfers to Agent2\",\n\t\tInstruction: \"You are agent 1\",\n\t\tModel:       cm1,\n\t})\n\tassert.NoError(t, err)\n\n\tagent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent2\",\n\t\tDescription: \"Second agent\",\n\t\tInstruction: \"You are agent 2\",\n\t\tModel:       cm2,\n\t})\n\tassert.NoError(t, err)\n\n\tagentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents})\n\n\tvar mu sync.Mutex\n\tonStartContextMarkers := make(map[string][]string)\n\n\thandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tmarker, _ := ctx.Value(testOnStartMarkerKey).(string)\n\t\t\tonStartContextMarkers[info.Name] = append(onStartContextMarkers[info.Name], marker)\n\t\t\tmu.Unlock()\n\n\t\t\treturn context.WithValue(ctx, testOnStartMarkerKey, info.Name+\"_marker\")\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tmu.Lock()\n\tdefer mu.Unlock()\n\n\tassert.NotEmpty(t, onStartContextMarkers[\"Agent1\"], \"Agent1's OnStart should be called\")\n\tassert.NotEmpty(t, onStartContextMarkers[\"Agent2\"], \"Agent2's OnStart should be called\")\n\n\tif len(onStartContextMarkers[\"Agent1\"]) > 0 {\n\t\tassert.Equal(t, \"\", onStartContextMarkers[\"Agent1\"][0],\n\t\t\t\"Agent1's OnStart should receive context without marker (initial context)\")\n\t}\n\tif len(onStartContextMarkers[\"Agent2\"]) > 0 {\n\t\tassert.Equal(t, \"\", onStartContextMarkers[\"Agent2\"][0],\n\t\t\t\"Agent2's first OnStart should NOT inherit Agent1's marker - context should be isolated\")\n\t}\n}\n\nfunc TestCallbackDesignatedToSpecificAgent(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm1 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"transferring to Agent2\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"transfer_1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      TransferToAgentToolName,\n\t\t\t\t\t\tArguments: `{\"agent_name\": \"Agent2\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\tcm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes()\n\n\tcm2 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"final response from Agent2\", nil), nil).\n\t\tTimes(1)\n\tcm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes()\n\n\tagent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent1\",\n\t\tDescription: \"First agent that transfers to Agent2\",\n\t\tInstruction: \"You are agent 1\",\n\t\tModel:       cm1,\n\t})\n\tassert.NoError(t, err)\n\n\tagent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent2\",\n\t\tDescription: \"Second agent\",\n\t\tInstruction: \"You are agent 2\",\n\t\tModel:       cm2,\n\t})\n\tassert.NoError(t, err)\n\n\tagentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents})\n\n\tvar mu sync.Mutex\n\tonStartCalls := make(map[string]int)\n\n\tagent2OnlyHandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tonStartCalls[info.Name]++\n\t\t\tmu.Unlock()\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(agent2OnlyHandler).DesignateAgent(\"Agent2\"))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tmu.Lock()\n\tdefer mu.Unlock()\n\n\tassert.Equal(t, 0, onStartCalls[\"Agent1\"], \"Agent1's OnStart should NOT be called when handler is designated to Agent2\")\n\tassert.Equal(t, 1, onStartCalls[\"Agent2\"], \"Agent2's OnStart should be called exactly once\")\n}\n\nfunc TestCallbackDesignatedToMultipleAgents(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm1 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"transferring to Agent2\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"transfer_1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      TransferToAgentToolName,\n\t\t\t\t\t\tArguments: `{\"agent_name\": \"Agent2\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\tcm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes()\n\n\tcm2 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"final response from Agent2\", nil), nil).\n\t\tTimes(1)\n\tcm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes()\n\n\tagent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent1\",\n\t\tDescription: \"First agent\",\n\t\tInstruction: \"You are agent 1\",\n\t\tModel:       cm1,\n\t})\n\tassert.NoError(t, err)\n\n\tagent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent2\",\n\t\tDescription: \"Second agent\",\n\t\tInstruction: \"You are agent 2\",\n\t\tModel:       cm2,\n\t})\n\tassert.NoError(t, err)\n\n\tagentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents})\n\n\tvar mu sync.Mutex\n\tonStartCalls := make(map[string]int)\n\n\tagent1And2Handler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tonStartCalls[info.Name]++\n\t\t\tmu.Unlock()\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(agent1And2Handler).DesignateAgent(\"Agent1\", \"Agent2\"))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tmu.Lock()\n\tdefer mu.Unlock()\n\n\tassert.Equal(t, 1, onStartCalls[\"Agent1\"], \"Agent1's OnStart should be called exactly once\")\n\tassert.Equal(t, 1, onStartCalls[\"Agent2\"], \"Agent2's OnStart should be called exactly once\")\n}\n\nfunc TestCallbackDesignatedExcludesNonMatchingAgents(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm1 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"transferring to Agent2\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"transfer_1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      TransferToAgentToolName,\n\t\t\t\t\t\tArguments: `{\"agent_name\": \"Agent2\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\tcm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes()\n\n\tcm2 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"final response from Agent2\", nil), nil).\n\t\tTimes(1)\n\tcm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes()\n\n\tagent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent1\",\n\t\tDescription: \"First agent\",\n\t\tInstruction: \"You are agent 1\",\n\t\tModel:       cm1,\n\t})\n\tassert.NoError(t, err)\n\n\tagent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent2\",\n\t\tDescription: \"Second agent\",\n\t\tInstruction: \"You are agent 2\",\n\t\tModel:       cm2,\n\t})\n\tassert.NoError(t, err)\n\n\tagentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents})\n\n\tvar mu sync.Mutex\n\tonStartCalls := make(map[string]int)\n\n\tagent1OnlyHandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tonStartCalls[info.Name]++\n\t\t\tmu.Unlock()\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(agent1OnlyHandler).DesignateAgent(\"Agent1\"))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tmu.Lock()\n\tdefer mu.Unlock()\n\n\tassert.Equal(t, 1, onStartCalls[\"Agent1\"], \"Agent1's OnStart should be called exactly once\")\n\tassert.Equal(t, 0, onStartCalls[\"Agent2\"], \"Agent2's OnStart should NOT be called when handler is designated only to Agent1\")\n}\n\nfunc TestMixedDesignatedAndGlobalCallbacks(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm1 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"transferring to Agent2\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"transfer_1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      TransferToAgentToolName,\n\t\t\t\t\t\tArguments: `{\"agent_name\": \"Agent2\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\tcm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes()\n\n\tcm2 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"final response from Agent2\", nil), nil).\n\t\tTimes(1)\n\tcm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes()\n\n\tagent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent1\",\n\t\tDescription: \"First agent that transfers to Agent2\",\n\t\tInstruction: \"You are agent 1\",\n\t\tModel:       cm1,\n\t})\n\tassert.NoError(t, err)\n\n\tagent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent2\",\n\t\tDescription: \"Second agent\",\n\t\tInstruction: \"You are agent 2\",\n\t\tModel:       cm2,\n\t})\n\tassert.NoError(t, err)\n\n\tagentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents})\n\n\tvar mu sync.Mutex\n\tglobalHandlerCalls := make(map[string]int)\n\tagent2OnlyHandlerCalls := make(map[string]int)\n\n\tglobalHandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tglobalHandlerCalls[info.Name]++\n\t\t\tmu.Unlock()\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\tagent2OnlyHandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tagent2OnlyHandlerCalls[info.Name]++\n\t\t\tmu.Unlock()\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\titer := runner.Query(ctx, \"hello\",\n\t\tWithCallbacks(globalHandler),\n\t\tWithCallbacks(agent2OnlyHandler).DesignateAgent(\"Agent2\"),\n\t)\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tmu.Lock()\n\tdefer mu.Unlock()\n\n\tassert.Equal(t, 1, globalHandlerCalls[\"Agent1\"], \"Global handler should fire for Agent1\")\n\tassert.Equal(t, 1, globalHandlerCalls[\"Agent2\"], \"Global handler should fire for Agent2\")\n\n\tassert.Equal(t, 0, agent2OnlyHandlerCalls[\"Agent1\"], \"Agent2-only handler should NOT fire for Agent1\")\n\tassert.Equal(t, 1, agent2OnlyHandlerCalls[\"Agent2\"], \"Agent2-only handler should fire for Agent2\")\n}\n\nfunc TestOnStartCalledOncePerAgentWithDesignation(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm1 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"transferring to Agent2\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"transfer_1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      TransferToAgentToolName,\n\t\t\t\t\t\tArguments: `{\"agent_name\": \"Agent2\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\tcm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes()\n\n\tcm2 := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"final response from Agent2\", nil), nil).\n\t\tTimes(1)\n\tcm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes()\n\n\tagent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent1\",\n\t\tDescription: \"First agent that transfers to Agent2\",\n\t\tInstruction: \"You are agent 1\",\n\t\tModel:       cm1,\n\t})\n\tassert.NoError(t, err)\n\n\tagent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"Agent2\",\n\t\tDescription: \"Second agent\",\n\t\tInstruction: \"You are agent 2\",\n\t\tModel:       cm2,\n\t})\n\tassert.NoError(t, err)\n\n\tagentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents})\n\n\tvar mu sync.Mutex\n\tonStartCalls := make(map[string]int)\n\n\thandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tonStartCalls[info.Name]++\n\t\t\tmu.Unlock()\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\titer := runner.Query(ctx, \"hello\", WithCallbacks(handler))\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tmu.Lock()\n\tdefer mu.Unlock()\n\n\tassert.Equal(t, 1, onStartCalls[\"Agent1\"], \"Agent1's OnStart should be called exactly once\")\n\tassert.Equal(t, 1, onStartCalls[\"Agent2\"], \"Agent2's OnStart should be called exactly once\")\n}\n"
  },
  {
    "path": "adk/callback_test.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestCopyEventIterator(t *testing.T) {\n\tt.Run(\"n=0 returns nil\", func(t *testing.T) {\n\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\tgo func() {\n\t\t\tgen.Send(&AgentEvent{AgentName: \"test\"})\n\t\t\tgen.Close()\n\t\t}()\n\n\t\tresult := copyEventIterator(iter, 0)\n\t\tassert.Nil(t, result)\n\t})\n\n\tt.Run(\"n=1 returns original iterator\", func(t *testing.T) {\n\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\tgo func() {\n\t\t\tgen.Send(&AgentEvent{AgentName: \"test\"})\n\t\t\tgen.Close()\n\t\t}()\n\n\t\tresult := copyEventIterator(iter, 1)\n\t\tassert.Len(t, result, 1)\n\t\tassert.Equal(t, iter, result[0])\n\t})\n\n\tt.Run(\"n>1 creates n independent copies\", func(t *testing.T) {\n\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\tevents := []*AgentEvent{\n\t\t\t{AgentName: \"agent1\", Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.AssistantMessage(\"msg1\", nil)}}},\n\t\t\t{AgentName: \"agent2\", Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.AssistantMessage(\"msg2\", nil)}}},\n\t\t}\n\n\t\tgo func() {\n\t\t\tfor _, e := range events {\n\t\t\t\tgen.Send(e)\n\t\t\t}\n\t\t\tgen.Close()\n\t\t}()\n\n\t\tn := 3\n\t\tcopies := copyEventIterator(iter, n)\n\t\tassert.Len(t, copies, n)\n\n\t\tvar wg sync.WaitGroup\n\t\treceivedEvents := make([][]*AgentEvent, n)\n\n\t\tfor i := 0; i < n; i++ {\n\t\t\twg.Add(1)\n\t\t\tgo func(idx int) {\n\t\t\t\tdefer wg.Done()\n\t\t\t\tfor {\n\t\t\t\t\tevent, ok := copies[idx].Next()\n\t\t\t\t\tif !ok {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\treceivedEvents[idx] = append(receivedEvents[idx], event)\n\t\t\t\t}\n\t\t\t}(i)\n\t\t}\n\n\t\twg.Wait()\n\n\t\tfor i := 0; i < n; i++ {\n\t\t\tassert.Len(t, receivedEvents[i], len(events), \"iterator %d should receive all events\", i)\n\t\t\tfor j, e := range receivedEvents[i] {\n\t\t\t\tassert.Equal(t, events[j].AgentName, e.AgentName)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestCopyAgentCallbackOutput(t *testing.T) {\n\tt.Run(\"nil output\", func(t *testing.T) {\n\t\tresult := copyAgentCallbackOutput(nil, 3)\n\t\tassert.Len(t, result, 3)\n\t\tfor _, r := range result {\n\t\t\tassert.Nil(t, r)\n\t\t}\n\t})\n\n\tt.Run(\"output with nil Events\", func(t *testing.T) {\n\t\tout := &AgentCallbackOutput{Events: nil}\n\t\tresult := copyAgentCallbackOutput(out, 3)\n\t\tassert.Len(t, result, 3)\n\t\tfor _, r := range result {\n\t\t\tassert.Equal(t, out, r)\n\t\t}\n\t})\n\n\tt.Run(\"valid output with events\", func(t *testing.T) {\n\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\tgo func() {\n\t\t\tgen.Send(&AgentEvent{AgentName: \"test\"})\n\t\t\tgen.Close()\n\t\t}()\n\n\t\tout := &AgentCallbackOutput{Events: iter}\n\t\tresult := copyAgentCallbackOutput(out, 2)\n\t\tassert.Len(t, result, 2)\n\n\t\tfor i, r := range result {\n\t\t\tassert.NotNil(t, r, \"result[%d] should not be nil\", i)\n\t\t\tassert.NotNil(t, r.Events, \"result[%d].Events should not be nil\", i)\n\t\t}\n\t})\n}\n\nfunc TestConvAgentCallbackInput(t *testing.T) {\n\tt.Run(\"valid AgentCallbackInput\", func(t *testing.T) {\n\t\tinput := &AgentCallbackInput{\n\t\t\tInput: &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}},\n\t\t}\n\t\tresult := ConvAgentCallbackInput(input)\n\t\tassert.Equal(t, input, result)\n\t})\n\n\tt.Run(\"invalid type returns nil\", func(t *testing.T) {\n\t\tresult := ConvAgentCallbackInput(\"invalid\")\n\t\tassert.Nil(t, result)\n\t})\n\n\tt.Run(\"nil returns nil\", func(t *testing.T) {\n\t\tresult := ConvAgentCallbackInput(nil)\n\t\tassert.Nil(t, result)\n\t})\n}\n\nfunc TestConvAgentCallbackOutput(t *testing.T) {\n\tt.Run(\"valid AgentCallbackOutput\", func(t *testing.T) {\n\t\titer, _ := NewAsyncIteratorPair[*AgentEvent]()\n\t\toutput := &AgentCallbackOutput{Events: iter}\n\t\tresult := ConvAgentCallbackOutput(output)\n\t\tassert.Equal(t, output, result)\n\t})\n\n\tt.Run(\"invalid type returns nil\", func(t *testing.T) {\n\t\tresult := ConvAgentCallbackOutput(\"invalid\")\n\t\tassert.Nil(t, result)\n\t})\n\n\tt.Run(\"nil returns nil\", func(t *testing.T) {\n\t\tresult := ConvAgentCallbackOutput(nil)\n\t\tassert.Nil(t, result)\n\t})\n}\n\ntype mockTyperAgent struct {\n\tname      string\n\tagentType string\n}\n\nfunc (a *mockTyperAgent) Name(_ context.Context) string        { return a.name }\nfunc (a *mockTyperAgent) Description(_ context.Context) string { return \"mock agent\" }\nfunc (a *mockTyperAgent) GetType() string                      { return a.agentType }\nfunc (a *mockTyperAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\tgen.Close()\n\treturn iter\n}\n\ntype mockNonTyperAgent struct {\n\tname string\n}\n\nfunc (a *mockNonTyperAgent) Name(_ context.Context) string        { return a.name }\nfunc (a *mockNonTyperAgent) Description(_ context.Context) string { return \"mock agent\" }\nfunc (a *mockNonTyperAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\tgen.Close()\n\treturn iter\n}\n\nfunc TestGetAgentType(t *testing.T) {\n\tt.Run(\"agent implementing Typer\", func(t *testing.T) {\n\t\tagent := &mockTyperAgent{name: \"test\", agentType: \"CustomType\"}\n\t\tresult := getAgentType(agent)\n\t\tassert.Equal(t, \"CustomType\", result)\n\t})\n\n\tt.Run(\"agent not implementing Typer\", func(t *testing.T) {\n\t\tagent := &mockNonTyperAgent{name: \"test\"}\n\t\tresult := getAgentType(agent)\n\t\tassert.Equal(t, \"\", result)\n\t})\n}\n\nfunc TestWithCallbacksOption(t *testing.T) {\n\thandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\topt := WithCallbacks(handler)\n\topts := getCommonOptions(nil, opt)\n\n\tassert.Len(t, opts.handlers, 1)\n}\n\nfunc TestWithMultipleCallbacksOption(t *testing.T) {\n\thandler1 := callbacks.NewHandlerBuilder().Build()\n\thandler2 := callbacks.NewHandlerBuilder().Build()\n\topt := WithCallbacks(handler1, handler2)\n\n\topts := getCommonOptions(nil, opt)\n\n\tassert.Len(t, opts.handlers, 2)\n}\n"
  },
  {
    "path": "adk/chatmodel.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/gob\"\n\t\"errors\"\n\t\"fmt\"\n\t\"math\"\n\t\"runtime/debug\"\n\t\"sync\"\n\t\"sync/atomic\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype chatModelAgentExecCtx struct {\n\truntimeReturnDirectly map[string]bool\n\tgenerator             *AsyncGenerator[*AgentEvent]\n}\n\nfunc (e *chatModelAgentExecCtx) send(event *AgentEvent) {\n\tif e != nil && e.generator != nil {\n\t\te.generator.Send(event)\n\t}\n}\n\ntype chatModelAgentExecCtxKey struct{}\n\nfunc withChatModelAgentExecCtx(ctx context.Context, execCtx *chatModelAgentExecCtx) context.Context {\n\treturn context.WithValue(ctx, chatModelAgentExecCtxKey{}, execCtx)\n}\n\nfunc getChatModelAgentExecCtx(ctx context.Context) *chatModelAgentExecCtx {\n\tif v := ctx.Value(chatModelAgentExecCtxKey{}); v != nil {\n\t\treturn v.(*chatModelAgentExecCtx)\n\t}\n\treturn nil\n}\n\ntype chatModelAgentRunOptions struct {\n\tchatModelOptions []model.Option\n\ttoolOptions      []tool.Option\n\tagentToolOptions map[string][]AgentRunOption\n\n\thistoryModifier func(context.Context, []Message) []Message\n}\n\n// WithChatModelOptions sets options for the underlying chat model.\nfunc WithChatModelOptions(opts []model.Option) AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) {\n\t\tt.chatModelOptions = opts\n\t})\n}\n\n// WithToolOptions sets options for tools used by the chat model agent.\nfunc WithToolOptions(opts []tool.Option) AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) {\n\t\tt.toolOptions = opts\n\t})\n}\n\n// WithAgentToolRunOptions specifies per-tool run options for the agent.\nfunc WithAgentToolRunOptions(opts map[string][]AgentRunOption) AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) {\n\t\tt.agentToolOptions = opts\n\t})\n}\n\n// WithHistoryModifier sets a function to modify history during resume.\n// Deprecated: use ResumeWithData and ChatModelAgentResumeData instead.\nfunc WithHistoryModifier(f func(context.Context, []Message) []Message) AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) {\n\t\tt.historyModifier = f\n\t})\n}\n\ntype ToolsConfig struct {\n\tcompose.ToolsNodeConfig\n\n\t// ReturnDirectly specifies tools that cause the agent to return immediately when called.\n\t// If multiple listed tools are called simultaneously, only the first one triggers the return.\n\t// The map keys are tool names indicate whether the tool should trigger immediate return.\n\tReturnDirectly map[string]bool\n\n\t// EmitInternalEvents indicates whether internal events from agentTool should be emitted\n\t// to the parent agent's AsyncGenerator, allowing real-time streaming of nested agent output\n\t// to the end-user via Runner.\n\t//\n\t// Note that these forwarded events are NOT recorded in the parent agent's runSession.\n\t// They are only emitted to the end-user and have no effect on the parent agent's state\n\t// or checkpoint.\n\t//\n\t// Action Scoping:\n\t// Actions emitted by the inner agent are scoped to the agent tool boundary:\n\t//   - Interrupted: Propagated via CompositeInterrupt to allow proper interrupt/resume\n\t//   - Exit, TransferToAgent, BreakLoop: Ignored outside the agent tool\n\tEmitInternalEvents bool\n}\n\n// GenModelInput transforms agent instructions and input into a format suitable for the model.\ntype GenModelInput func(ctx context.Context, instruction string, input *AgentInput) ([]Message, error)\n\nfunc defaultGenModelInput(ctx context.Context, instruction string, input *AgentInput) ([]Message, error) {\n\tmsgs := make([]Message, 0, len(input.Messages)+1)\n\n\tif instruction != \"\" {\n\t\tsp := schema.SystemMessage(instruction)\n\n\t\tvs := GetSessionValues(ctx)\n\t\tif len(vs) > 0 {\n\t\t\tct := prompt.FromMessages(schema.FString, sp)\n\t\t\tms, err := ct.Format(ctx, vs)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"defaultGenModelInput: failed to format instruction using FString template. \"+\n\t\t\t\t\t\"This formatting is triggered automatically when SessionValues are present. \"+\n\t\t\t\t\t\"If your instruction contains literal curly braces (e.g., JSON), provide a custom GenModelInput that uses another format. If you are using \"+\n\t\t\t\t\t\"SessionValues for purposes other than instruction formatting, provide a custom GenModelInput that does no formatting at all: %w\", err)\n\t\t\t}\n\n\t\t\tsp = ms[0]\n\t\t}\n\n\t\tmsgs = append(msgs, sp)\n\t}\n\n\tmsgs = append(msgs, input.Messages...)\n\n\treturn msgs, nil\n}\n\n// ChatModelAgentState represents the state of a chat model agent during conversation.\n// This is the primary state type for both ChatModelAgentMiddleware and AgentMiddleware callbacks.\ntype ChatModelAgentState struct {\n\t// Messages contains all messages in the current conversation session.\n\tMessages []Message\n}\n\n// AgentMiddleware provides hooks to customize agent behavior at various stages of execution.\n//\n// Limitations of AgentMiddleware (struct-based):\n//   - Struct types are closed: users cannot add new methods\n//   - Callbacks only return error, cannot return modified context\n//   - Configuration is scattered across closures when using factory functions\n//\n// For new code requiring extensibility, consider using ChatModelAgentMiddleware (interface-based) instead.\n// AgentMiddleware is kept for backward compatibility and remains suitable for simple,\n// static additions like extra instruction or tools.\n//\n// See ChatModelAgentMiddleware documentation for detailed comparison.\ntype AgentMiddleware struct {\n\t// AdditionalInstruction adds supplementary text to the agent's system instruction.\n\t// This instruction is concatenated with the base instruction before each chat model call.\n\tAdditionalInstruction string\n\n\t// AdditionalTools adds supplementary tools to the agent's available toolset.\n\t// These tools are combined with the tools configured for the agent.\n\tAdditionalTools []tool.BaseTool\n\n\t// BeforeChatModel is called before each ChatModel invocation, allowing modification of the agent state.\n\tBeforeChatModel func(context.Context, *ChatModelAgentState) error\n\n\t// AfterChatModel is called after each ChatModel invocation, allowing modification of the agent state.\n\tAfterChatModel func(context.Context, *ChatModelAgentState) error\n\n\t// WrapToolCall wraps tool calls with custom middleware logic.\n\t// Each middleware contains Invokable and/or Streamable functions for tool calls.\n\tWrapToolCall compose.ToolMiddleware\n}\n\ntype ChatModelAgentConfig struct {\n\t// Name of the agent. Better be unique across all agents.\n\tName string\n\t// Description of the agent's capabilities.\n\t// Helps other agents determine whether to transfer tasks to this agent.\n\tDescription string\n\t// Instruction used as the system prompt for this agent.\n\t// Optional. If empty, no system prompt will be used.\n\t// Supports f-string placeholders for session values in default GenModelInput, for example:\n\t// \"You are a helpful assistant. The current time is {Time}. The current user is {User}.\"\n\t// These placeholders will be replaced with session values for \"Time\" and \"User\".\n\tInstruction string\n\n\t// Model is the chat model used by the agent.\n\t// If your ChatModelAgent uses any tools, this model must support the model.WithTools\n\t// call option, as that's how ChatModelAgent configures the model with tool information.\n\tModel model.BaseChatModel\n\n\tToolsConfig ToolsConfig\n\n\t// GenModelInput transforms instructions and input messages into the model's input format.\n\t// Optional. Defaults to defaultGenModelInput which combines instruction and messages.\n\tGenModelInput GenModelInput\n\n\t// Exit defines the tool used to terminate the agent process.\n\t// Optional. If nil, no Exit Action will be generated.\n\t// You can use the provided 'ExitTool' implementation directly.\n\tExit tool.BaseTool\n\n\t// OutputKey stores the agent's response in the session.\n\t// Optional. When set, stores output via AddSessionValue(ctx, outputKey, msg.Content).\n\tOutputKey string\n\n\t// MaxIterations defines the upper limit of ChatModel generation cycles.\n\t// The agent will terminate with an error if this limit is exceeded.\n\t// Optional. Defaults to 20.\n\tMaxIterations int\n\n\t// Middlewares configures agent middleware for extending functionality.\n\t// Use for simple, static additions like extra instruction or tools.\n\t// Kept for backward compatibility; for new code, consider using Handlers instead.\n\tMiddlewares []AgentMiddleware\n\n\t// Handlers configures interface-based handlers for extending agent behavior.\n\t// Unlike Middlewares (struct-based), Handlers allow users to:\n\t//   - Add custom methods to their handler implementations\n\t//   - Return modified context from handler methods\n\t//   - Centralize configuration in struct fields instead of closures\n\t//\n\t// Handlers are processed after Middlewares, in registration order.\n\t// See ChatModelAgentMiddleware documentation for when to use Handlers vs Middlewares.\n\t//\n\t// Execution Order (relative to AgentMiddleware and ToolsConfig):\n\t//\n\t// Model call lifecycle (outermost to innermost wrapper chain):\n\t//  1. AgentMiddleware.BeforeChatModel (hook, runs before model call)\n\t//  2. ChatModelAgentMiddleware.BeforeModelRewriteState (hook, can modify state before model call)\n\t//  3. retryModelWrapper (internal - retries on failure, if configured)\n\t//  4. eventSenderModelWrapper (internal - sends model response events)\n\t//  5. ChatModelAgentMiddleware.WrapModel (wrapper, first registered is outermost)\n\t//  6. callbackInjectionModelWrapper (internal - injects callbacks if not enabled)\n\t//  7. Model.Generate/Stream\n\t//  8. ChatModelAgentMiddleware.AfterModelRewriteState (hook, can modify state after model call)\n\t//  9. AgentMiddleware.AfterChatModel (hook, runs after model call)\n\t//\n\t// Custom Event Sender Position:\n\t// By default, events are sent after all user middlewares (WrapModel) have processed the output,\n\t// containing the modified messages. To send events with original (unmodified) output, pass\n\t// NewEventSenderModelWrapper as a Handler after the modifying middleware:\n\t//\n\t//   agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t//       Handlers: []adk.ChatModelAgentMiddleware{\n\t//           myCustomHandler,                   // First registered = outermost wrapper\n\t//           adk.NewEventSenderModelWrapper(),  // Last registered = innermost, events sent with original output\n\t//       },\n\t//   })\n\t//\n\t// Handler order: first registered is outermost. So [A, B, C] becomes A(B(C(model))).\n\t// EventSenderModelWrapper sends events in post-processing, so placing it innermost\n\t// means it receives the original model output before outer handlers modify it.\n\t//\n\t// When EventSenderModelWrapper is detected in Handlers, the framework skips\n\t// the default event sender to avoid duplicate events.\n\t//\n\t// Tool call lifecycle (outermost to innermost):\n\t//  1. eventSenderToolHandler (internal ToolMiddleware - sends tool result events after all processing)\n\t//  2. ToolsConfig.ToolCallMiddlewares (ToolMiddleware)\n\t//  3. AgentMiddleware.WrapToolCall (ToolMiddleware)\n\t//  4. ChatModelAgentMiddleware.WrapToolCall (wrapper, first registered is outermost)\n\t//  5. callbackInjectedToolCall (internal - injects callbacks if tool doesn't handle them)\n\t//  6. Tool.InvokableRun/StreamableRun\n\t//\n\t// Tool List Modification:\n\t//\n\t// There are two ways to modify the tool list:\n\t//\n\t//  1. In BeforeAgent: Modify ChatModelAgentContext.Tools ([]tool.BaseTool) directly. This affects\n\t//     both the tool info list passed to ChatModel AND the actual tools available for\n\t//     execution. Changes persist for the entire agent run.\n\t//\n\t//  2. In WrapModel: Create a model wrapper that modifies the tool info list per model\n\t//     request using model.WithTools(toolInfos). This ONLY affects the tool info list\n\t//     passed to ChatModel, NOT the actual tools available for execution. Use this for\n\t//     dynamic tool filtering/selection based on conversation context. The modification\n\t//     is scoped to this model request only.\n\tHandlers []ChatModelAgentMiddleware\n\n\t// ModelRetryConfig configures retry behavior for the ChatModel.\n\t// When set, the agent will automatically retry failed ChatModel calls\n\t// based on the configured policy.\n\t// Optional. If nil, no retry will be performed.\n\tModelRetryConfig *ModelRetryConfig\n}\n\ntype ChatModelAgent struct {\n\tname        string\n\tdescription string\n\tinstruction string\n\n\tmodel       model.BaseChatModel\n\ttoolsConfig ToolsConfig\n\n\tgenModelInput GenModelInput\n\n\toutputKey     string\n\tmaxIterations int\n\n\tsubAgents   []Agent\n\tparentAgent Agent\n\n\tdisallowTransferToParent bool\n\n\texit tool.BaseTool\n\n\thandlers    []ChatModelAgentMiddleware\n\tmiddlewares []AgentMiddleware\n\n\tmodelRetryConfig *ModelRetryConfig\n\n\tonce   sync.Once\n\trun    runFunc\n\tfrozen uint32\n\texeCtx *execContext\n}\n\ntype runFunc func(ctx context.Context, input *AgentInput, generator *AsyncGenerator[*AgentEvent], store *bridgeStore, instruction string, returnDirectly map[string]bool, opts ...compose.Option)\n\n// NewChatModelAgent constructs a chat model-backed agent with the provided config.\nfunc NewChatModelAgent(ctx context.Context, config *ChatModelAgentConfig) (*ChatModelAgent, error) {\n\tif config.Name == \"\" {\n\t\treturn nil, errors.New(\"agent 'Name' is required\")\n\t}\n\tif config.Description == \"\" {\n\t\treturn nil, errors.New(\"agent 'Description' is required\")\n\t}\n\tif config.Model == nil {\n\t\treturn nil, errors.New(\"agent 'Model' is required\")\n\t}\n\n\tgenInput := defaultGenModelInput\n\tif config.GenModelInput != nil {\n\t\tgenInput = config.GenModelInput\n\t}\n\n\ttc := config.ToolsConfig\n\n\t// Tool call middleware execution order (outermost to innermost):\n\t// 1. eventSenderToolHandler (internal - sends tool result events after all modifications)\n\t// 2. User-provided ToolsConfig.ToolCallMiddlewares (original order preserved)\n\t// 3. Middlewares' WrapToolCall (in registration order)\n\t// 4. ChatModelAgentMiddleware.WrapToolCall (in registration order)\n\t// 5. callbackInjectedToolCall (internal - injects callbacks if tool doesn't handle them)\n\teventSender := &eventSenderToolHandler{}\n\ttc.ToolCallMiddlewares = append(\n\t\t[]compose.ToolMiddleware{{Invokable: eventSender.WrapInvokableToolCall,\n\t\t\tStreamable:         eventSender.WrapStreamableToolCall,\n\t\t\tEnhancedInvokable:  eventSender.WrapEnhancedInvokableToolCall,\n\t\t\tEnhancedStreamable: eventSender.WrapEnhancedStreamableToolCall,\n\t\t}},\n\t\ttc.ToolCallMiddlewares...,\n\t)\n\ttc.ToolCallMiddlewares = append(tc.ToolCallMiddlewares, collectToolMiddlewaresFromMiddlewares(config.Middlewares)...)\n\n\treturn &ChatModelAgent{\n\t\tname:             config.Name,\n\t\tdescription:      config.Description,\n\t\tinstruction:      config.Instruction,\n\t\tmodel:            config.Model,\n\t\ttoolsConfig:      tc,\n\t\tgenModelInput:    genInput,\n\t\texit:             config.Exit,\n\t\toutputKey:        config.OutputKey,\n\t\tmaxIterations:    config.MaxIterations,\n\t\thandlers:         config.Handlers,\n\t\tmiddlewares:      config.Middlewares,\n\t\tmodelRetryConfig: config.ModelRetryConfig,\n\t}, nil\n}\n\nfunc collectToolMiddlewaresFromMiddlewares(mws []AgentMiddleware) []compose.ToolMiddleware {\n\tvar middlewares []compose.ToolMiddleware\n\tfor _, m := range mws {\n\t\tif m.WrapToolCall.Invokable == nil && m.WrapToolCall.Streamable == nil {\n\t\t\tcontinue\n\t\t}\n\t\tmiddlewares = append(middlewares, m.WrapToolCall)\n\t}\n\treturn middlewares\n}\n\nconst (\n\tTransferToAgentToolName        = \"transfer_to_agent\"\n\tTransferToAgentToolDesc        = \"Transfer the question to another agent.\"\n\tTransferToAgentToolDescChinese = \"将问题移交给其他 Agent。\"\n)\n\nvar (\n\ttoolInfoTransferToAgent = &schema.ToolInfo{\n\t\tName: TransferToAgentToolName,\n\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"agent_name\": {\n\t\t\t\tDesc:     \"the name of the agent to transfer to\",\n\t\t\t\tRequired: true,\n\t\t\t\tType:     schema.String,\n\t\t\t},\n\t\t}),\n\t}\n\n\tToolInfoExit = &schema.ToolInfo{\n\t\tName: \"exit\",\n\t\tDesc: \"Exit the agent process and return the final result.\",\n\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"final_result\": {\n\t\t\t\tDesc:     \"the final result to return\",\n\t\t\t\tRequired: true,\n\t\t\t\tType:     schema.String,\n\t\t\t},\n\t\t}),\n\t}\n)\n\ntype ExitTool struct{}\n\nfunc (et ExitTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn ToolInfoExit, nil\n}\n\nfunc (et ExitTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\ttype exitParams struct {\n\t\tFinalResult string `json:\"final_result\"`\n\t}\n\n\tparams := &exitParams{}\n\terr := sonic.UnmarshalString(argumentsInJSON, params)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\terr = SendToolGenAction(ctx, \"exit\", NewExitAction())\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn params.FinalResult, nil\n}\n\ntype transferToAgent struct{}\n\nfunc (tta transferToAgent) Info(_ context.Context) (*schema.ToolInfo, error) {\n\tdesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: TransferToAgentToolDesc,\n\t\tChinese: TransferToAgentToolDescChinese,\n\t})\n\tinfo := *toolInfoTransferToAgent\n\tinfo.Desc = desc\n\treturn &info, nil\n}\n\nfunc transferToAgentToolOutput(destName string) string {\n\ttpl := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: \"successfully transferred to agent [%s]\",\n\t\tChinese: \"成功移交任务至 agent [%s]\",\n\t})\n\treturn fmt.Sprintf(tpl, destName)\n}\n\nfunc (tta transferToAgent) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\ttype transferParams struct {\n\t\tAgentName string `json:\"agent_name\"`\n\t}\n\n\tparams := &transferParams{}\n\terr := sonic.UnmarshalString(argumentsInJSON, params)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\terr = SendToolGenAction(ctx, TransferToAgentToolName, NewTransferToAgentAction(params.AgentName))\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn transferToAgentToolOutput(params.AgentName), nil\n}\n\nfunc (a *ChatModelAgent) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *ChatModelAgent) Description(_ context.Context) string {\n\treturn a.description\n}\n\nfunc (a *ChatModelAgent) GetType() string {\n\treturn \"ChatModel\"\n}\n\nfunc (a *ChatModelAgent) OnSetSubAgents(_ context.Context, subAgents []Agent) error {\n\tif atomic.LoadUint32(&a.frozen) == 1 {\n\t\treturn errors.New(\"agent has been frozen after run\")\n\t}\n\n\tif len(a.subAgents) > 0 {\n\t\treturn errors.New(\"agent's sub-agents has already been set\")\n\t}\n\n\ta.subAgents = subAgents\n\treturn nil\n}\n\nfunc (a *ChatModelAgent) OnSetAsSubAgent(_ context.Context, parent Agent) error {\n\tif atomic.LoadUint32(&a.frozen) == 1 {\n\t\treturn errors.New(\"agent has been frozen after run\")\n\t}\n\n\tif a.parentAgent != nil {\n\t\treturn errors.New(\"agent has already been set as a sub-agent of another agent\")\n\t}\n\n\ta.parentAgent = parent\n\treturn nil\n}\n\nfunc (a *ChatModelAgent) OnDisallowTransferToParent(_ context.Context) error {\n\tif atomic.LoadUint32(&a.frozen) == 1 {\n\t\treturn errors.New(\"agent has been frozen after run\")\n\t}\n\n\ta.disallowTransferToParent = true\n\n\treturn nil\n}\n\ntype ChatModelAgentInterruptInfo struct {\n\tInfo *compose.InterruptInfo\n\tData []byte\n}\n\nfunc init() {\n\tschema.RegisterName[*ChatModelAgentInterruptInfo](\"_eino_adk_chat_model_agent_interrupt_info\")\n}\n\nfunc setOutputToSession(ctx context.Context, msg Message, msgStream MessageStream, outputKey string) error {\n\tif msg != nil {\n\t\tAddSessionValue(ctx, outputKey, msg.Content)\n\t\treturn nil\n\t}\n\n\tconcatenated, err := schema.ConcatMessageStream(msgStream)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tAddSessionValue(ctx, outputKey, concatenated.Content)\n\treturn nil\n}\n\nfunc errFunc(err error) runFunc {\n\treturn func(ctx context.Context, input *AgentInput, generator *AsyncGenerator[*AgentEvent], store *bridgeStore, _ string, _ map[string]bool, _ ...compose.Option) {\n\t\tgenerator.Send(&AgentEvent{Err: err})\n\t}\n}\n\n// ChatModelAgentResumeData holds data that can be provided to a ChatModelAgent during a resume operation\n// to modify its behavior. It is provided via the adk.ResumeWithData function.\ntype ChatModelAgentResumeData struct {\n\t// HistoryModifier is a function that can transform the agent's message history before it is sent to the model.\n\t// This allows for adding new information or context upon resumption.\n\tHistoryModifier func(ctx context.Context, history []Message) []Message\n}\n\ntype execContext struct {\n\tinstruction    string\n\ttoolsNodeConf  compose.ToolsNodeConfig\n\treturnDirectly map[string]bool\n\n\ttoolInfos      []*schema.ToolInfo\n\tunwrappedTools []tool.BaseTool\n\n\trebuildGraph bool // whether needs to instantiate a new graph because of topology changes due to tool modifications\n\ttoolUpdated  bool // whether needs to pass a compose.WithToolList option to ToolsNode due to tool list change\n}\n\nfunc (a *ChatModelAgent) applyBeforeAgent(ctx context.Context, ec *execContext) (context.Context, *execContext, error) {\n\trunCtx := &ChatModelAgentContext{\n\t\tInstruction:    ec.instruction,\n\t\tTools:          cloneSlice(ec.unwrappedTools),\n\t\tReturnDirectly: copyMap(ec.returnDirectly),\n\t}\n\n\tvar err error\n\tfor i, handler := range a.handlers {\n\t\tctx, runCtx, err = handler.BeforeAgent(ctx, runCtx)\n\t\tif err != nil {\n\t\t\treturn ctx, nil, fmt.Errorf(\"handler[%d] (%T) BeforeAgent failed: %w\", i, handler, err)\n\t\t}\n\t}\n\n\truntimeEC := &execContext{\n\t\tinstruction: runCtx.Instruction,\n\t\ttoolsNodeConf: compose.ToolsNodeConfig{\n\t\t\tTools:               runCtx.Tools,\n\t\t\tToolCallMiddlewares: cloneSlice(ec.toolsNodeConf.ToolCallMiddlewares),\n\t\t},\n\t\treturnDirectly: runCtx.ReturnDirectly,\n\t\ttoolUpdated:    true,\n\t\trebuildGraph: (len(ec.toolsNodeConf.Tools) == 0 && len(runCtx.Tools) > 0) ||\n\t\t\t(len(ec.returnDirectly) == 0 && len(runCtx.ReturnDirectly) > 0),\n\t}\n\n\ttoolInfos, err := genToolInfos(ctx, &runtimeEC.toolsNodeConf)\n\tif err != nil {\n\t\treturn ctx, nil, err\n\t}\n\n\truntimeEC.toolInfos = toolInfos\n\n\treturn ctx, runtimeEC, nil\n}\n\nfunc (a *ChatModelAgent) prepareExecContext(ctx context.Context) (*execContext, error) {\n\tinstruction := a.instruction\n\ttoolsNodeConf := compose.ToolsNodeConfig{\n\t\tTools:                cloneSlice(a.toolsConfig.Tools),\n\t\tToolCallMiddlewares:  cloneSlice(a.toolsConfig.ToolCallMiddlewares),\n\t\tUnknownToolsHandler:  a.toolsConfig.UnknownToolsHandler,\n\t\tExecuteSequentially:  a.toolsConfig.ExecuteSequentially,\n\t\tToolArgumentsHandler: a.toolsConfig.ToolArgumentsHandler,\n\t}\n\treturnDirectly := copyMap(a.toolsConfig.ReturnDirectly)\n\n\ttransferToAgents := a.subAgents\n\tif a.parentAgent != nil && !a.disallowTransferToParent {\n\t\ttransferToAgents = append(transferToAgents, a.parentAgent)\n\t}\n\n\tif len(transferToAgents) > 0 {\n\t\ttransferInstruction := genTransferToAgentInstruction(ctx, transferToAgents)\n\t\tinstruction = concatInstructions(instruction, transferInstruction)\n\n\t\ttoolsNodeConf.Tools = append(toolsNodeConf.Tools, &transferToAgent{})\n\t\treturnDirectly[TransferToAgentToolName] = true\n\t}\n\n\tif a.exit != nil {\n\t\ttoolsNodeConf.Tools = append(toolsNodeConf.Tools, a.exit)\n\t\texitInfo, err := a.exit.Info(ctx)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturnDirectly[exitInfo.Name] = true\n\t}\n\n\tfor _, m := range a.middlewares {\n\t\tif m.AdditionalInstruction != \"\" {\n\t\t\tinstruction = concatInstructions(instruction, m.AdditionalInstruction)\n\t\t}\n\t\ttoolsNodeConf.Tools = append(toolsNodeConf.Tools, m.AdditionalTools...)\n\t}\n\n\tunwrappedTools := cloneSlice(toolsNodeConf.Tools)\n\n\thandlerMiddlewares := handlersToToolMiddlewares(a.handlers)\n\ttoolsNodeConf.ToolCallMiddlewares = append(toolsNodeConf.ToolCallMiddlewares, handlerMiddlewares...)\n\n\ttoolInfos, err := genToolInfos(ctx, &toolsNodeConf)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &execContext{\n\t\tinstruction:    instruction,\n\t\ttoolsNodeConf:  toolsNodeConf,\n\t\treturnDirectly: returnDirectly,\n\t\ttoolInfos:      toolInfos,\n\t\tunwrappedTools: unwrappedTools,\n\t}, nil\n}\n\nfunc (a *ChatModelAgent) buildNoToolsRunFunc(_ context.Context) runFunc {\n\twrappedModel := buildModelWrappers(a.model, &modelWrapperConfig{\n\t\thandlers:    a.handlers,\n\t\tmiddlewares: a.middlewares,\n\t\tretryConfig: a.modelRetryConfig,\n\t})\n\n\ttype noToolsInput struct {\n\t\tinput       *AgentInput\n\t\tinstruction string\n\t}\n\n\treturn func(ctx context.Context, input *AgentInput, generator *AsyncGenerator[*AgentEvent],\n\t\tstore *bridgeStore, instruction string, _ map[string]bool, opts ...compose.Option) {\n\n\t\tchain := compose.NewChain[noToolsInput, Message](\n\t\t\tcompose.WithGenLocalState(func(ctx context.Context) (state *State) {\n\t\t\t\treturn &State{}\n\t\t\t})).\n\t\t\tAppendLambda(compose.InvokableLambda(func(ctx context.Context, in noToolsInput) ([]Message, error) {\n\t\t\t\tmessages, err := a.genModelInput(ctx, in.instruction, in.input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\treturn messages, nil\n\t\t\t})).\n\t\t\tAppendChatModel(wrappedModel)\n\n\t\tr, err := chain.Compile(ctx, compose.WithGraphName(a.name),\n\t\t\tcompose.WithCheckPointStore(store),\n\t\t\tcompose.WithSerializer(&gobSerializer{}))\n\t\tif err != nil {\n\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t\treturn\n\t\t}\n\n\t\tctx = withChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{\n\t\t\tgenerator: generator,\n\t\t})\n\n\t\tin := noToolsInput{input: input, instruction: instruction}\n\n\t\tvar msg Message\n\t\tvar msgStream MessageStream\n\t\tif input.EnableStreaming {\n\t\t\tmsgStream, err = r.Stream(ctx, in, opts...)\n\t\t} else {\n\t\t\tmsg, err = r.Invoke(ctx, in, opts...)\n\t\t}\n\n\t\tif err == nil {\n\t\t\tif a.outputKey != \"\" {\n\t\t\t\terr = setOutputToSession(ctx, msg, msgStream, a.outputKey)\n\t\t\t\tif err != nil {\n\t\t\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t\t\t}\n\t\t\t} else if msgStream != nil {\n\t\t\t\tmsgStream.Close()\n\t\t\t}\n\t\t} else {\n\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t}\n\t}\n}\n\nfunc (a *ChatModelAgent) buildReactRunFunc(ctx context.Context, bc *execContext) (runFunc, error) {\n\tconf := &reactConfig{\n\t\tmodel:       a.model,\n\t\ttoolsConfig: &bc.toolsNodeConf,\n\t\tmodelWrapperConf: &modelWrapperConfig{\n\t\t\thandlers:    a.handlers,\n\t\t\tmiddlewares: a.middlewares,\n\t\t\tretryConfig: a.modelRetryConfig,\n\t\t\ttoolInfos:   bc.toolInfos,\n\t\t},\n\t\ttoolsReturnDirectly: bc.returnDirectly,\n\t\tagentName:           a.name,\n\t\tmaxIterations:       a.maxIterations,\n\t}\n\n\ttype reactRunInput struct {\n\t\tinput       *AgentInput\n\t\tinstruction string\n\t}\n\n\treturn func(ctx context.Context, input *AgentInput, generator *AsyncGenerator[*AgentEvent], store *bridgeStore,\n\t\tinstruction string, returnDirectly map[string]bool, opts ...compose.Option) {\n\t\tg, err := newReact(ctx, conf)\n\t\tif err != nil {\n\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t\treturn\n\t\t}\n\n\t\tchain := compose.NewChain[reactRunInput, Message]().\n\t\t\tAppendLambda(\n\t\t\t\tcompose.InvokableLambda(func(ctx context.Context, in reactRunInput) (*reactInput, error) {\n\t\t\t\t\tmessages, genErr := a.genModelInput(ctx, in.instruction, in.input)\n\t\t\t\t\tif genErr != nil {\n\t\t\t\t\t\treturn nil, genErr\n\t\t\t\t\t}\n\t\t\t\t\treturn &reactInput{\n\t\t\t\t\t\tmessages: messages,\n\t\t\t\t\t}, nil\n\t\t\t\t}),\n\t\t\t).\n\t\t\tAppendGraph(g, compose.WithNodeName(\"ReAct\"), compose.WithGraphCompileOptions(compose.WithMaxRunSteps(math.MaxInt)))\n\n\t\tvar compileOptions []compose.GraphCompileOption\n\t\tcompileOptions = append(compileOptions,\n\t\t\tcompose.WithGraphName(a.name),\n\t\t\tcompose.WithCheckPointStore(store),\n\t\t\tcompose.WithSerializer(&gobSerializer{}),\n\t\t\tcompose.WithMaxRunSteps(math.MaxInt))\n\n\t\trunnable, err_ := chain.Compile(ctx, compileOptions...)\n\t\tif err_ != nil {\n\t\t\tgenerator.Send(&AgentEvent{Err: err_})\n\t\t\treturn\n\t\t}\n\n\t\tctx = withChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{\n\t\t\truntimeReturnDirectly: returnDirectly,\n\t\t\tgenerator:             generator,\n\t\t})\n\n\t\tin := reactRunInput{\n\t\t\tinput:       input,\n\t\t\tinstruction: instruction,\n\t\t}\n\n\t\tvar runOpts []compose.Option\n\t\trunOpts = append(runOpts, opts...)\n\t\tif a.toolsConfig.EmitInternalEvents {\n\t\t\trunOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEventGenerator(generator))))\n\t\t}\n\t\tif input.EnableStreaming {\n\t\t\trunOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEnableStreaming(true))))\n\t\t}\n\n\t\tvar msg Message\n\t\tvar msgStream MessageStream\n\t\tif input.EnableStreaming {\n\t\t\tmsgStream, err_ = runnable.Stream(ctx, in, runOpts...)\n\t\t} else {\n\t\t\tmsg, err_ = runnable.Invoke(ctx, in, runOpts...)\n\t\t}\n\n\t\tif err_ == nil {\n\t\t\tif a.outputKey != \"\" {\n\t\t\t\terr_ = setOutputToSession(ctx, msg, msgStream, a.outputKey)\n\t\t\t\tif err_ != nil {\n\t\t\t\t\tgenerator.Send(&AgentEvent{Err: err_})\n\t\t\t\t}\n\t\t\t} else if msgStream != nil {\n\t\t\t\tmsgStream.Close()\n\t\t\t}\n\n\t\t\treturn\n\t\t}\n\n\t\tinfo, ok := compose.ExtractInterruptInfo(err_)\n\t\tif !ok {\n\t\t\tgenerator.Send(&AgentEvent{Err: err_})\n\t\t\treturn\n\t\t}\n\n\t\tdata, existed, err := store.Get(ctx, bridgeCheckpointID)\n\t\tif err != nil {\n\t\t\tgenerator.Send(&AgentEvent{AgentName: a.name, Err: fmt.Errorf(\"failed to get interrupt info: %w\", err)})\n\t\t\treturn\n\t\t}\n\t\tif !existed {\n\t\t\tgenerator.Send(&AgentEvent{AgentName: a.name, Err: fmt.Errorf(\"interrupt occurred but checkpoint data is missing\")})\n\t\t\treturn\n\t\t}\n\n\t\tis := FromInterruptContexts(info.InterruptContexts)\n\n\t\tevent := CompositeInterrupt(ctx, info, data, is)\n\t\tevent.Action.Interrupted.Data = &ChatModelAgentInterruptInfo{\n\t\t\tInfo: info,\n\t\t\tData: data,\n\t\t}\n\t\tevent.AgentName = a.name\n\t\tgenerator.Send(event)\n\t}, nil\n}\n\nfunc (a *ChatModelAgent) buildRunFunc(ctx context.Context) runFunc {\n\ta.once.Do(func() {\n\t\tec, err := a.prepareExecContext(ctx)\n\t\tif err != nil {\n\t\t\ta.run = errFunc(err)\n\t\t\treturn\n\t\t}\n\n\t\ta.exeCtx = ec\n\n\t\tif len(ec.toolsNodeConf.Tools) == 0 {\n\t\t\ta.run = a.buildNoToolsRunFunc(ctx)\n\t\t\treturn\n\t\t}\n\n\t\trun, err := a.buildReactRunFunc(ctx, ec)\n\t\tif err != nil {\n\t\t\ta.run = errFunc(err)\n\t\t\treturn\n\t\t}\n\t\ta.run = run\n\t})\n\n\tatomic.StoreUint32(&a.frozen, 1)\n\n\treturn a.run\n}\n\nfunc (a *ChatModelAgent) getRunFunc(ctx context.Context) (context.Context, runFunc, *execContext, error) {\n\tdefaultRun := a.buildRunFunc(ctx)\n\tbc := a.exeCtx\n\n\tif bc == nil {\n\t\treturn ctx, defaultRun, bc, nil\n\t}\n\n\tif len(a.handlers) == 0 {\n\t\truntimeBC := &execContext{\n\t\t\tinstruction:    bc.instruction,\n\t\t\ttoolsNodeConf:  bc.toolsNodeConf,\n\t\t\treturnDirectly: bc.returnDirectly,\n\t\t\ttoolInfos:      bc.toolInfos,\n\t\t}\n\t\treturn ctx, defaultRun, runtimeBC, nil\n\t}\n\n\tctx, runtimeBC, err := a.applyBeforeAgent(ctx, bc)\n\tif err != nil {\n\t\treturn ctx, nil, nil, err\n\t}\n\n\tif !runtimeBC.rebuildGraph {\n\t\treturn ctx, defaultRun, runtimeBC, nil\n\t}\n\n\tvar tempRun runFunc\n\tif len(runtimeBC.toolsNodeConf.Tools) == 0 {\n\t\ttempRun = a.buildNoToolsRunFunc(ctx)\n\t} else {\n\t\ttempRun, err = a.buildReactRunFunc(ctx, runtimeBC)\n\t\tif err != nil {\n\t\t\treturn ctx, nil, nil, err\n\t\t}\n\t}\n\n\treturn ctx, tempRun, runtimeBC, nil\n}\n\nfunc (a *ChatModelAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tctx, run, bc, err := a.getRunFunc(ctx)\n\tif err != nil {\n\t\tgo func() {\n\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t\tgenerator.Close()\n\t\t}()\n\t\treturn iterator\n\t}\n\n\tco := getComposeOptions(opts)\n\tco = append(co, compose.WithCheckPointID(bridgeCheckpointID))\n\n\tif bc != nil {\n\t\tco = append(co, compose.WithChatModelOption(model.WithTools(bc.toolInfos)))\n\t\tif bc.toolUpdated {\n\t\t\tco = append(co, compose.WithToolsNodeOption(compose.WithToolList(bc.toolsNodeConf.Tools...)))\n\t\t}\n\t}\n\n\tgo func() {\n\t\tdefer func() {\n\t\t\tpanicErr := recover()\n\t\t\tif panicErr != nil {\n\t\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\t\tgenerator.Send(&AgentEvent{Err: e})\n\t\t\t}\n\n\t\t\tgenerator.Close()\n\t\t}()\n\n\t\tvar (\n\t\t\tinstruction    string\n\t\t\treturnDirectly map[string]bool\n\t\t)\n\n\t\tif bc != nil {\n\t\t\tinstruction = bc.instruction\n\t\t\treturnDirectly = bc.returnDirectly\n\t\t}\n\n\t\trun(ctx, input, generator, newBridgeStore(), instruction, returnDirectly, co...)\n\t}()\n\n\treturn iterator\n}\n\nfunc (a *ChatModelAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tctx, run, bc, err := a.getRunFunc(ctx)\n\tif err != nil {\n\t\tgo func() {\n\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t\tgenerator.Close()\n\t\t}()\n\t\treturn iterator\n\t}\n\n\tco := getComposeOptions(opts)\n\tco = append(co, compose.WithCheckPointID(bridgeCheckpointID))\n\n\tif bc != nil {\n\t\tco = append(co, compose.WithChatModelOption(model.WithTools(bc.toolInfos)))\n\t\tif bc.toolUpdated {\n\t\t\tco = append(co, compose.WithToolsNodeOption(compose.WithToolList(bc.toolsNodeConf.Tools...)))\n\t\t}\n\t}\n\n\tif info.InterruptState == nil {\n\t\tpanic(fmt.Sprintf(\"ChatModelAgent.Resume: agent '%s' was asked to resume but has no state\", a.Name(ctx)))\n\t}\n\n\tstateByte, ok := info.InterruptState.([]byte)\n\tif !ok {\n\t\tpanic(fmt.Sprintf(\"ChatModelAgent.Resume: agent '%s' was asked to resume but has invalid interrupt state type: %T\",\n\t\t\ta.Name(ctx), info.InterruptState))\n\t}\n\n\t// Migrate legacy checkpoints before resume.\n\t// This covers both:\n\t// - v0.7.*: state is stored as a struct wire type (stateV07) under the legacy name.\n\t// - v0.8.0-v0.8.3: state is stored as a GobEncoder payload under the same legacy name and must\n\t//   be routed to a GobDecode-compatible compat type via byte-patching.\n\t// The result is re-encoded so the resume path always operates on the current *State.\n\tstateByte, err = preprocessComposeCheckpoint(stateByte)\n\tif err != nil {\n\t\tgo func() {\n\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t\tgenerator.Close()\n\t\t}()\n\t\treturn iterator\n\t}\n\n\tvar historyModifier func(ctx context.Context, history []Message) []Message\n\tif info.ResumeData != nil {\n\t\tresumeData, ok := info.ResumeData.(*ChatModelAgentResumeData)\n\t\tif !ok {\n\t\t\tpanic(fmt.Sprintf(\"ChatModelAgent.Resume: agent '%s' was asked to resume but has invalid resume data type: %T\",\n\t\t\t\ta.Name(ctx), info.ResumeData))\n\t\t}\n\t\thistoryModifier = resumeData.HistoryModifier\n\t}\n\n\tif historyModifier != nil {\n\t\tco = append(co, compose.WithStateModifier(func(ctx context.Context, path compose.NodePath, state any) error {\n\t\t\ts, ok := state.(*State)\n\t\t\tif !ok {\n\t\t\t\treturn nil\n\t\t\t}\n\t\t\ts.Messages = historyModifier(ctx, s.Messages)\n\t\t\treturn nil\n\t\t}))\n\t}\n\n\tgo func() {\n\t\tdefer func() {\n\t\t\tpanicErr := recover()\n\t\t\tif panicErr != nil {\n\t\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\t\tgenerator.Send(&AgentEvent{Err: e})\n\t\t\t}\n\n\t\t\tgenerator.Close()\n\t\t}()\n\n\t\tvar (\n\t\t\tinstruction    string\n\t\t\treturnDirectly map[string]bool\n\t\t)\n\n\t\tif bc != nil {\n\t\t\tinstruction = bc.instruction\n\t\t\treturnDirectly = bc.returnDirectly\n\t\t}\n\n\t\trun(ctx, &AgentInput{EnableStreaming: info.EnableStreaming}, generator,\n\t\t\tnewResumeBridgeStore(stateByte), instruction, returnDirectly, co...)\n\t}()\n\n\treturn iterator\n}\n\nfunc getComposeOptions(opts []AgentRunOption) []compose.Option {\n\to := GetImplSpecificOptions[chatModelAgentRunOptions](nil, opts...)\n\tvar co []compose.Option\n\tif len(o.chatModelOptions) > 0 {\n\t\tco = append(co, compose.WithChatModelOption(o.chatModelOptions...))\n\t}\n\tvar to []tool.Option\n\tif len(o.toolOptions) > 0 {\n\t\tto = append(to, o.toolOptions...)\n\t}\n\tfor toolName, atos := range o.agentToolOptions {\n\t\tto = append(to, withAgentToolOptions(toolName, atos))\n\t}\n\tif len(to) > 0 {\n\t\tco = append(co, compose.WithToolsNodeOption(compose.WithToolOption(to...)))\n\t}\n\tif o.historyModifier != nil {\n\t\tco = append(co, compose.WithStateModifier(func(ctx context.Context, path compose.NodePath, state any) error {\n\t\t\ts, ok := state.(*State)\n\t\t\tif !ok {\n\t\t\t\treturn fmt.Errorf(\"unexpected state type: %T, expected: %T\", state, &State{})\n\t\t\t}\n\t\t\ts.Messages = o.historyModifier(ctx, s.Messages)\n\t\t\treturn nil\n\t\t}))\n\t}\n\treturn co\n}\n\ntype gobSerializer struct{}\n\nfunc (g *gobSerializer) Marshal(v any) ([]byte, error) {\n\tbuf := new(bytes.Buffer)\n\terr := gob.NewEncoder(buf).Encode(v)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn buf.Bytes(), nil\n}\n\nfunc (g *gobSerializer) Unmarshal(data []byte, v any) error {\n\tbuf := bytes.NewBuffer(data)\n\treturn gob.NewDecoder(buf).Decode(v)\n}\n\n// preprocessComposeCheckpoint migrates legacy compose checkpoints to the current format.\n// It handles the v0.8.0-v0.8.3 format:\n//   - gob name \"_eino_adk_state_v080_\" (already byte-patched by preprocessADKCheckpoint\n//     from \"_eino_adk_react_state\"), opaque-bytes wire format → decoded as *stateV080\n//\n// v0.7 checkpoints need no migration — State is now a plain struct registered under the\n// same gob name, and gob handles missing fields gracefully.\n//\n// Fast path: if the legacy name is not present, skip entirely.\nfunc preprocessComposeCheckpoint(data []byte) ([]byte, error) {\n\tconst lenPrefixedCompatName = \"\\x15\" + stateGobNameV080\n\tif bytes.Contains(data, []byte(lenPrefixedCompatName)) {\n\t\t// v0.8.0-v0.8.3: already byte-patched by preprocessADKCheckpoint; decode as *stateV080.\n\t\tmigrated, err := compose.MigrateCheckpointState(data, &gobSerializer{}, func(state any) (any, bool, error) {\n\t\t\tsc, ok := state.(*stateV080)\n\t\t\tif !ok {\n\t\t\t\treturn state, false, nil\n\t\t\t}\n\t\t\treturn stateV080ToState(sc), true, nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to migrate v0.8.0-v0.8.3 compose checkpoint: %w\", err)\n\t\t}\n\t\treturn migrated, nil\n\t}\n\n\treturn data, nil\n}\n"
  },
  {
    "path": "adk/chatmodel_retry_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"io\"\n\t\"strings\"\n\t\"sync/atomic\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nvar errRetryAble = errors.New(\"retry-able error\")\nvar errNonRetryAble = errors.New(\"non-retry-able error\")\n\nfunc TestChatModelAgentRetry_NoTools_DirectError_Generate(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tvar callCount int32\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\tcount := atomic.AddInt32(&callCount, 1)\n\t\t\tif count < 3 {\n\t\t\t\treturn nil, errRetryAble\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"Success after retry\", nil), nil\n\t\t}).Times(3)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent for retry functionality\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.Nil(t, event.Err)\n\tassert.NotNil(t, event.Output)\n\tassert.Equal(t, \"Success after retry\", event.Output.MessageOutput.Message.Content)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n\tassert.Equal(t, int32(3), atomic.LoadInt32(&callCount))\n}\n\nfunc TestChatModelAgentRetry_NoTools_DirectError_Stream(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tvar callCount int32\n\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\t\t\tcount := atomic.AddInt32(&callCount, 1)\n\t\t\tif count < 2 {\n\t\t\t\treturn nil, errRetryAble\n\t\t\t}\n\t\t\treturn schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"Success\", nil),\n\t\t\t}), nil\n\t\t}).Times(2)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent for retry functionality\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages:        []Message{schema.UserMessage(\"Hello\")},\n\t\tEnableStreaming: true,\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.Nil(t, event.Err)\n\tassert.NotNil(t, event.Output)\n\tassert.True(t, event.Output.MessageOutput.IsStreaming)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n\tassert.Equal(t, int32(2), atomic.LoadInt32(&callCount))\n}\n\ntype streamErrorModel struct {\n\tcallCount   int32\n\tfailAtChunk int\n\tmaxFailures int\n\ttools       []*schema.ToolInfo\n\treturnTool  bool\n}\n\nfunc (m *streamErrorModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {\n\treturn schema.AssistantMessage(\"Generated\", nil), nil\n}\n\nfunc (m *streamErrorModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tcount := atomic.AddInt32(&m.callCount, 1)\n\n\tsr, sw := schema.Pipe[*schema.Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tfor i := 0; i < 5; i++ {\n\t\t\tif i == m.failAtChunk && int(count) <= m.maxFailures {\n\t\t\t\tsw.Send(nil, errRetryAble)\n\t\t\t\treturn\n\t\t\t}\n\t\t\tif m.returnTool && i == 0 {\n\t\t\t\tsw.Send(schema.AssistantMessage(\"\", []schema.ToolCall{{\n\t\t\t\t\tID:       \"call-1\",\n\t\t\t\t\tFunction: schema.FunctionCall{Name: \"test_tool\", Arguments: \"{}\"},\n\t\t\t\t}}), nil)\n\t\t\t} else {\n\t\t\t\tsw.Send(schema.AssistantMessage(\"chunk\", nil), nil)\n\t\t\t}\n\t\t}\n\t}()\n\treturn sr, nil\n}\n\nfunc (m *streamErrorModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\tm.tools = tools\n\treturn m, nil\n}\n\nfunc TestChatModelAgentRetry_StreamError(t *testing.T) {\n\tt.Run(\"WithTools\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tm := &streamErrorModel{\n\t\t\tfailAtChunk: 2,\n\t\t\tmaxFailures: 2,\n\t\t\treturnTool:  false,\n\t\t}\n\n\t\tconfig := &ChatModelAgentConfig{\n\t\t\tName:        \"RetryTestAgent\",\n\t\t\tDescription: \"Test agent for retry functionality\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       m,\n\t\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\t\tMaxRetries:  3,\n\t\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t\t},\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{&fakeToolForTest{tarCount: 0}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tagent, err := NewChatModelAgent(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tinput := &AgentInput{\n\t\t\tMessages:        []Message{schema.UserMessage(\"Hello\")},\n\t\t\tEnableStreaming: true,\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\n\t\tvar events []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iterator.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tassert.Equal(t, 3, len(events))\n\n\t\tvar streamErrEventCount int\n\t\tvar errs []error\n\t\tfor i, event := range events {\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming {\n\t\t\t\tsr := event.Output.MessageOutput.MessageStream\n\t\t\t\tfor {\n\t\t\t\t\tmsg, err := sr.Recv()\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\tstreamErrEventCount++\n\t\t\t\t\t\terrs = append(errs, err)\n\t\t\t\t\t\tt.Logf(\"event %d: err: %v\", i, err)\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\tt.Logf(\"event %d: %v\", i, msg.Content)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 2, streamErrEventCount)\n\t\tassert.Equal(t, 2, len(errs))\n\t\tvar willRetryErr *WillRetryError\n\t\tassert.True(t, errors.As(errs[0], &willRetryErr))\n\t\tassert.True(t, errors.As(errs[1], &willRetryErr))\n\t\tassert.Equal(t, int32(3), atomic.LoadInt32(&m.callCount))\n\t})\n\n\tt.Run(\"NoTools\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tm := &streamErrorModel{\n\t\t\tfailAtChunk: 2,\n\t\t\tmaxFailures: 2,\n\t\t\treturnTool:  false,\n\t\t}\n\n\t\tconfig := &ChatModelAgentConfig{\n\t\t\tName:        \"RetryTestAgent\",\n\t\t\tDescription: \"Test agent for retry functionality\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       m,\n\t\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\t\tMaxRetries:  3,\n\t\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t\t},\n\t\t}\n\n\t\tagent, err := NewChatModelAgent(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tinput := &AgentInput{\n\t\t\tMessages:        []Message{schema.UserMessage(\"Hello\")},\n\t\t\tEnableStreaming: true,\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\n\t\tvar events []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iterator.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tassert.Equal(t, 3, len(events))\n\n\t\tvar streamErrEventCount int\n\t\tvar errs []error\n\t\tfor i, event := range events {\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming {\n\t\t\t\tsr := event.Output.MessageOutput.MessageStream\n\t\t\t\tfor {\n\t\t\t\t\tmsg, err := sr.Recv()\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\tstreamErrEventCount++\n\t\t\t\t\t\terrs = append(errs, err)\n\t\t\t\t\t\tt.Logf(\"event %d: err: %v\", i, err)\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\tt.Logf(\"event %d: %v\", i, msg.Content)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 2, streamErrEventCount)\n\t\tassert.Equal(t, 2, len(errs))\n\t\tvar willRetryErr *WillRetryError\n\t\tassert.True(t, errors.As(errs[0], &willRetryErr))\n\t\tassert.True(t, errors.As(errs[1], &willRetryErr))\n\t\tassert.Equal(t, int32(3), atomic.LoadInt32(&m.callCount))\n\t})\n}\n\nfunc TestChatModelAgentRetry_WithTools_DirectError_Generate(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tvar callCount int32\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\tcount := atomic.AddInt32(&callCount, 1)\n\t\t\tif count < 2 {\n\t\t\t\treturn nil, errRetryAble\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"Success after retry\", nil), nil\n\t\t}).Times(2)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tfakeTool := &fakeToolForTest{tarCount: 0}\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent for retry functionality\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t},\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.Nil(t, event.Err)\n\tassert.NotNil(t, event.Output)\n\tassert.Equal(t, \"Success after retry\", event.Output.MessageOutput.Message.Content)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n\tassert.Equal(t, int32(2), atomic.LoadInt32(&callCount))\n}\n\nfunc TestChatModelAgentRetry_NonRetryableError(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(nil, errNonRetryAble).Times(1)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent for retry functionality\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.NotNil(t, event.Err)\n\tassert.True(t, errors.Is(event.Err, errNonRetryAble))\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\ntype inputCapturingModel struct {\n\tcapturedInputs [][]Message\n}\n\nfunc (m *inputCapturingModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) {\n\tm.capturedInputs = append(m.capturedInputs, input)\n\treturn schema.AssistantMessage(\"Response from capturing model\", nil), nil\n}\n\nfunc (m *inputCapturingModel) Stream(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tm.capturedInputs = append(m.capturedInputs, input)\n\treturn schema.StreamReaderFromArray([]*schema.Message{\n\t\tschema.AssistantMessage(\"Response from capturing model\", nil),\n\t}), nil\n}\n\nfunc (m *inputCapturingModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\treturn m, nil\n}\n\nfunc TestChatModelAgentRetry_MaxRetriesExhausted(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(nil, errRetryAble).Times(4)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent for retry functionality\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.NotNil(t, event.Err)\n\tassert.True(t, errors.Is(event.Err, ErrExceedMaxRetries))\n\tvar retryErr *RetryExhaustedError\n\tassert.True(t, errors.As(event.Err, &retryErr))\n\tassert.True(t, errors.Is(retryErr.LastErr, errRetryAble))\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestChatModelAgentRetry_BackoffFunction(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tvar backoffCalls []int\n\tvar callCount int32\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\tcount := atomic.AddInt32(&callCount, 1)\n\t\t\tif count < 3 {\n\t\t\t\treturn nil, errRetryAble\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"Success\", nil), nil\n\t\t}).Times(3)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent for retry functionality\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t\tBackoffFunc: func(ctx context.Context, attempt int) time.Duration {\n\t\t\t\tbackoffCalls = append(backoffCalls, attempt)\n\t\t\t\treturn time.Millisecond\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Nil(t, event.Err)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n\n\tassert.Equal(t, []int{1, 2}, backoffCalls)\n}\n\nfunc TestChatModelAgentRetry_NoRetryConfig(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(nil, errRetryAble).Times(1)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Test agent without retry config\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.NotNil(t, event.Err)\n\tassert.True(t, errors.Is(event.Err, errRetryAble))\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestChatModelAgentRetry_WithTools_NonRetryAbleStreamError(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(nil, errNonRetryAble).Times(1)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tfakeTool := &fakeToolForTest{tarCount: 0}\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent for retry functionality\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t},\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages:        []Message{schema.UserMessage(\"Hello\")},\n\t\tEnableStreaming: true,\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.NotNil(t, event.Err)\n\tassert.True(t, errors.Is(event.Err, errNonRetryAble))\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\ntype nonRetryAbleStreamErrorModel struct {\n\ttools []*schema.ToolInfo\n}\n\nfunc (m *nonRetryAbleStreamErrorModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {\n\treturn schema.AssistantMessage(\"Generated\", nil), nil\n}\n\nfunc (m *nonRetryAbleStreamErrorModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tsr, sw := schema.Pipe[*schema.Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(schema.AssistantMessage(\"chunk1\", nil), nil)\n\t\tsw.Send(nil, errNonRetryAble)\n\t}()\n\treturn sr, nil\n}\n\nfunc (m *nonRetryAbleStreamErrorModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\tm.tools = tools\n\treturn m, nil\n}\n\nfunc TestChatModelAgentRetry_NoTools_NonRetryAbleStreamError(t *testing.T) {\n\tctx := context.Background()\n\n\tm := &nonRetryAbleStreamErrorModel{}\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent for retry functionality\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       m,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages:        []Message{schema.UserMessage(\"Hello\")},\n\t\tEnableStreaming: true,\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tvar events []*AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\tassert.Equal(t, 2, len(events))\n\n\tevent0 := events[0]\n\tassert.NotNil(t, event0.Output)\n\tassert.NotNil(t, event0.Output.MessageOutput)\n\tassert.True(t, event0.Output.MessageOutput.IsStreaming)\n\tsr := event0.Output.MessageOutput.MessageStream\n\tvar streamErr error\n\tfor {\n\t\t_, err := sr.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tif err != nil {\n\t\t\tstreamErr = err\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.NotNil(t, streamErr)\n\tassert.True(t, errors.Is(streamErr, errNonRetryAble), \"Stream error should be the original error\")\n\n\tevent1 := events[1]\n\tassert.NotNil(t, event1.Err)\n\tassert.True(t, errors.Is(event1.Err, errNonRetryAble))\n}\n\nfunc TestDefaultBackoff(t *testing.T) {\n\tctx := context.Background()\n\n\td1 := defaultBackoff(ctx, 1)\n\td2 := defaultBackoff(ctx, 2)\n\td3 := defaultBackoff(ctx, 3)\n\n\tt.Logf(\"Backoff delays: d1=%v, d2=%v, d3=%v\", d1, d2, d3)\n\n\tassert.True(t, d1 >= 100*time.Millisecond && d1 < 150*time.Millisecond,\n\t\t\"First retry should be ~100ms + jitter (0-50ms), got %v\", d1)\n\tassert.True(t, d2 >= 200*time.Millisecond && d2 < 300*time.Millisecond,\n\t\t\"Second retry should be ~200ms + jitter (0-100ms), got %v\", d2)\n\tassert.True(t, d3 >= 400*time.Millisecond && d3 < 600*time.Millisecond,\n\t\t\"Third retry should be ~400ms + jitter (0-200ms), got %v\", d3)\n\n\td10 := defaultBackoff(ctx, 10)\n\tt.Logf(\"Backoff delay for attempt 10: %v\", d10)\n\tassert.True(t, d10 >= 10*time.Second && d10 <= 15*time.Second,\n\t\t\"Delay should be capped at 10s + jitter (0-5s), got %v\", d10)\n\n\td100 := defaultBackoff(ctx, 100)\n\tt.Logf(\"Backoff delay for attempt 100: %v\", d100)\n\tassert.True(t, d100 >= 10*time.Second && d100 <= 15*time.Second,\n\t\t\"Delay should still be capped at 10s + jitter for very high attempts, got %v\", d100)\n}\n\nfunc TestRetryExhaustedError_ErrorString(t *testing.T) {\n\terrWithLast := &RetryExhaustedError{\n\t\tLastErr:      errors.New(\"connection timeout\"),\n\t\tTotalRetries: 3,\n\t}\n\tassert.Contains(t, errWithLast.Error(), \"exceeds max retries\")\n\tassert.Contains(t, errWithLast.Error(), \"connection timeout\")\n\n\terrWithoutLast := &RetryExhaustedError{\n\t\tLastErr:      nil,\n\t\tTotalRetries: 3,\n\t}\n\tassert.Equal(t, \"exceeds max retries\", errWithoutLast.Error())\n}\n\nfunc TestWillRetryError_ErrorString(t *testing.T) {\n\twillRetry := &WillRetryError{ErrStr: \"transient error\", RetryAttempt: 1}\n\tassert.Equal(t, \"transient error\", willRetry.Error())\n}\n\ntype customError struct {\n\tcode int\n\tmsg  string\n}\n\nfunc (e *customError) Error() string {\n\treturn e.msg\n}\n\nfunc TestWillRetryError_Unwrap(t *testing.T) {\n\toriginalErr := &customError{code: 500, msg: \"internal error\"}\n\twillRetry := &WillRetryError{ErrStr: originalErr.Error(), RetryAttempt: 1, err: originalErr}\n\n\tassert.True(t, errors.Is(willRetry, originalErr))\n\n\tvar targetErr *customError\n\tassert.True(t, errors.As(willRetry, &targetErr))\n\tassert.Equal(t, 500, targetErr.code)\n\tassert.Equal(t, \"internal error\", targetErr.msg)\n}\n\nfunc TestChatModelAgentRetry_DefaultIsRetryAble(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tvar callCount int32\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\tcount := atomic.AddInt32(&callCount, 1)\n\t\t\tif count < 2 {\n\t\t\t\treturn nil, errors.New(\"any error should be retried\")\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"Success\", nil), nil\n\t\t}).Times(2)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"RetryTestAgent\",\n\t\tDescription: \"Test agent with default IsRetryAble\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries: 3,\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t}\n\titerator := agent.Run(ctx, input)\n\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.Nil(t, event.Err)\n\tassert.Equal(t, \"Success\", event.Output.MessageOutput.Message.Content)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n\tassert.Equal(t, int32(2), atomic.LoadInt32(&callCount))\n}\n\nfunc TestSequentialWorkflow_RetryAbleStreamError_SuccessfulRetry(t *testing.T) {\n\tctx := context.Background()\n\n\tretryModel := &streamErrorModel{\n\t\tfailAtChunk: 2,\n\t\tmaxFailures: 2,\n\t}\n\n\tagentA, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"AgentA\",\n\t\tDescription: \"Agent A with retry that emits stream errors then succeeds\",\n\t\tInstruction: \"You are agent A.\",\n\t\tModel:       retryModel,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tcapturingModel := &inputCapturingModel{}\n\tagentB, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"AgentB\",\n\t\tDescription: \"Agent B that captures input\",\n\t\tInstruction: \"You are agent B.\",\n\t\tModel:       capturingModel,\n\t})\n\tassert.NoError(t, err)\n\n\tsequentialAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\tName:        \"SequentialAgent\",\n\t\tDescription: \"Sequential agent A->B\",\n\t\tSubAgents:   []Agent{agentA, agentB},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages:        []Message{schema.UserMessage(\"Hello\")},\n\t\tEnableStreaming: true,\n\t}\n\tctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input)\n\titerator := sequentialAgent.Run(ctx, input)\n\n\tvar events []*AgentEvent\n\tvar willRetryErrCount int\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t\tif event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming {\n\t\t\tsr := event.Output.MessageOutput.MessageStream\n\t\t\tfor {\n\t\t\t\t_, err := sr.Recv()\n\t\t\t\tif err == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tif err != nil {\n\t\t\t\t\tvar retryErr *WillRetryError\n\t\t\t\t\tif errors.As(err, &retryErr) {\n\t\t\t\t\t\twillRetryErrCount++\n\t\t\t\t\t}\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tassert.Equal(t, 2, willRetryErrCount, \"End-user should receive 2 WillRetryError events\")\n\tassert.Equal(t, 1, len(capturingModel.capturedInputs), \"Agent B should be called exactly once\")\n\n\tsuccessorInput := capturingModel.capturedInputs[0]\n\tvar hasSuccessfulMessage bool\n\tfor _, msg := range successorInput {\n\t\tif strings.Contains(msg.Content, \"chunkchunkchunkchunkchunk\") {\n\t\t\thasSuccessfulMessage = true\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.True(t, hasSuccessfulMessage, \"Agent B should receive the successful message from Agent A\")\n\n\tfor _, msg := range successorInput {\n\t\tassert.NotContains(t, msg.Content, \"retry-able error\", \"Agent B should not receive failed stream messages\")\n\t}\n}\n\ntype streamErrorModelNoRetry struct {\n\tcallCount int32\n}\n\nfunc (m *streamErrorModelNoRetry) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {\n\treturn schema.AssistantMessage(\"Generated\", nil), nil\n}\n\nfunc (m *streamErrorModelNoRetry) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tatomic.AddInt32(&m.callCount, 1)\n\tsr, sw := schema.Pipe[*schema.Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(schema.AssistantMessage(\"chunk1\", nil), nil)\n\t\tsw.Send(schema.AssistantMessage(\"chunk2\", nil), nil)\n\t\tsw.Send(nil, errRetryAble)\n\t}()\n\treturn sr, nil\n}\n\nfunc (m *streamErrorModelNoRetry) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\treturn m, nil\n}\n\nfunc TestSequentialWorkflow_NonRetryAbleStreamError_StopsFlow(t *testing.T) {\n\tctx := context.Background()\n\n\tnonRetryModel := &nonRetryAbleStreamErrorModel{}\n\n\tagentA, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"AgentA\",\n\t\tDescription: \"Agent A that emits non-retryable stream error\",\n\t\tInstruction: \"You are agent A.\",\n\t\tModel:       nonRetryModel,\n\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\tMaxRetries:  3,\n\t\t\tIsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tcapturingModel := &inputCapturingModel{}\n\tagentB, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"AgentB\",\n\t\tDescription: \"Agent B that captures input\",\n\t\tInstruction: \"You are agent B.\",\n\t\tModel:       capturingModel,\n\t})\n\tassert.NoError(t, err)\n\n\tsequentialAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\tName:        \"SequentialAgent\",\n\t\tDescription: \"Sequential agent A->B\",\n\t\tSubAgents:   []Agent{agentA, agentB},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages:        []Message{schema.UserMessage(\"Hello\")},\n\t\tEnableStreaming: true,\n\t}\n\tctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input)\n\titerator := sequentialAgent.Run(ctx, input)\n\n\tvar events []*AgentEvent\n\tvar streamErrFound bool\n\tvar finalErrEvent *AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t\tif event.Err != nil {\n\t\t\tfinalErrEvent = event\n\t\t}\n\t\tif event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming {\n\t\t\tsr := event.Output.MessageOutput.MessageStream\n\t\t\tfor {\n\t\t\t\t_, err := sr.Recv()\n\t\t\t\tif err == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tif err != nil {\n\t\t\t\t\tstreamErrFound = true\n\t\t\t\t\tassert.True(t, errors.Is(err, errNonRetryAble), \"Stream error should be the original error\")\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tassert.True(t, streamErrFound, \"End-user should receive stream error\")\n\tassert.NotNil(t, finalErrEvent, \"Should receive a final error event\")\n\tassert.True(t, errors.Is(finalErrEvent.Err, errNonRetryAble), \"Final error should be the non-retryable error\")\n\tassert.Equal(t, 0, len(capturingModel.capturedInputs), \"Agent B should NOT be called due to error\")\n}\n\nfunc TestSequentialWorkflow_NoRetryConfig_StreamError_StopsFlow(t *testing.T) {\n\tctx := context.Background()\n\n\tnoRetryModel := &streamErrorModelNoRetry{}\n\n\tagentA, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"AgentA\",\n\t\tDescription: \"Agent A without retry config that emits stream error\",\n\t\tInstruction: \"You are agent A.\",\n\t\tModel:       noRetryModel,\n\t})\n\tassert.NoError(t, err)\n\n\tcapturingModel := &inputCapturingModel{}\n\tagentB, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"AgentB\",\n\t\tDescription: \"Agent B that captures input\",\n\t\tInstruction: \"You are agent B.\",\n\t\tModel:       capturingModel,\n\t})\n\tassert.NoError(t, err)\n\n\tsequentialAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\tName:        \"SequentialAgent\",\n\t\tDescription: \"Sequential agent A->B\",\n\t\tSubAgents:   []Agent{agentA, agentB},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{\n\t\tMessages:        []Message{schema.UserMessage(\"Hello\")},\n\t\tEnableStreaming: true,\n\t}\n\tctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input)\n\titerator := sequentialAgent.Run(ctx, input)\n\n\tvar events []*AgentEvent\n\tvar streamErrFound bool\n\tvar finalErrEvent *AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t\tif event.Err != nil {\n\t\t\tfinalErrEvent = event\n\t\t}\n\t\tif event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming {\n\t\t\tsr := event.Output.MessageOutput.MessageStream\n\t\t\tfor {\n\t\t\t\t_, err := sr.Recv()\n\t\t\t\tif err == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tif err != nil {\n\t\t\t\t\tstreamErrFound = true\n\t\t\t\t\tassert.True(t, errors.Is(err, errRetryAble), \"Stream error should be the original error\")\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tassert.True(t, streamErrFound, \"End-user should receive stream error\")\n\tassert.NotNil(t, finalErrEvent, \"Should receive a final error event\")\n\tassert.True(t, errors.Is(finalErrEvent.Err, errRetryAble), \"Final error should be the original error\")\n\tassert.Equal(t, 0, len(capturingModel.capturedInputs), \"Agent B should NOT be called due to error\")\n\tassert.Equal(t, int32(1), atomic.LoadInt32(&noRetryModel.callCount), \"Model should only be called once (no retry)\")\n}\n"
  },
  {
    "path": "adk/chatmodel_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// TestChatModelAgentRun tests the Run method of ChatModelAgent\nfunc TestChatModelAgentRun(t *testing.T) {\n\t// Basic test for Run method\n\tt.Run(\"BasicFunctionality\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Hello, I am an AI assistant.\", nil), nil).\n\t\t\tTimes(1)\n\n\t\t// Create a ChatModelAgent\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for unit testing\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, agent)\n\n\t\t// Run the agent\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{\n\t\t\t\tschema.UserMessage(\"Hello, who are you?\"),\n\t\t\t},\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\t\tassert.NotNil(t, iterator)\n\n\t\t// Get the event from the iterator\n\t\tevent, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event)\n\t\tassert.Nil(t, event.Err)\n\t\tassert.NotNil(t, event.Output)\n\t\tassert.NotNil(t, event.Output.MessageOutput)\n\n\t\t// Verify the message content\n\t\tmsg := event.Output.MessageOutput.Message\n\t\tassert.NotNil(t, msg)\n\t\tassert.Equal(t, \"Hello, I am an AI assistant.\", msg.Content)\n\n\t\t// No more events\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\t})\n\n\tt.Run(\"BasicChatModelWithAgentMiddleware\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Hello, I am an AI assistant.\", nil), nil).\n\t\t\tTimes(1)\n\n\t\tafterChatModelExecuted := false\n\n\t\t// Create a ChatModelAgent\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for unit testing\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t\tMiddlewares: []AgentMiddleware{\n\t\t\t\t{\n\t\t\t\t\tBeforeChatModel: func(ctx context.Context, state *ChatModelAgentState) error {\n\t\t\t\t\t\tstate.Messages = append(state.Messages, schema.UserMessage(\"m\"))\n\t\t\t\t\t\treturn nil\n\t\t\t\t\t},\n\t\t\t\t\tAfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error {\n\t\t\t\t\t\tassert.Len(t, state.Messages, 4)\n\t\t\t\t\t\tafterChatModelExecuted = true\n\t\t\t\t\t\treturn nil\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, agent)\n\n\t\t// Run the agent\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{\n\t\t\t\tschema.UserMessage(\"Hello, who are you?\"),\n\t\t\t},\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\t\tassert.NotNil(t, iterator)\n\n\t\t// Get the event from the iterator\n\t\tevent, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event)\n\t\tassert.Nil(t, event.Err)\n\t\tassert.NotNil(t, event.Output)\n\t\tassert.NotNil(t, event.Output.MessageOutput)\n\n\t\t// Verify the message content\n\t\tmsg := event.Output.MessageOutput.Message\n\t\tassert.NotNil(t, msg)\n\t\tassert.Equal(t, \"Hello, I am an AI assistant.\", msg.Content)\n\n\t\t// No more events\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\n\t\tassert.True(t, afterChatModelExecuted)\n\t})\n\n\tt.Run(\"AfterChatModel_NoTools_ModifyDoesNotAffectEvent\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"original content\", nil), nil).\n\t\t\tTimes(1)\n\n\t\tvar capturedMessages []*schema.Message\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for AfterChatModel NoTools scenario\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t\tMiddlewares: []AgentMiddleware{\n\t\t\t\t{\n\t\t\t\t\tAfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error {\n\t\t\t\t\t\tcapturedMessages = make([]*schema.Message, len(state.Messages))\n\t\t\t\t\t\tcopy(capturedMessages, state.Messages)\n\t\t\t\t\t\tstate.Messages = append(state.Messages, schema.AssistantMessage(\"appended content\", nil))\n\t\t\t\t\t\treturn nil\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\n\t\tevent, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event)\n\t\tassert.Nil(t, event.Err)\n\t\tassert.NotNil(t, event.Output)\n\t\tassert.NotNil(t, event.Output.MessageOutput)\n\n\t\tmsg := event.Output.MessageOutput.Message\n\t\tassert.NotNil(t, msg)\n\t\tassert.Equal(t, \"original content\", msg.Content)\n\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\n\t\tassert.Len(t, capturedMessages, 3)\n\t})\n\n\tt.Run(\"AfterChatModel_ReAct_ModifyAffectsFlow\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tgenerateCount := 0\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\tgenerateCount++\n\t\t\t\tif generateCount == 1 {\n\t\t\t\t\treturn schema.AssistantMessage(\"first response with tool call\", []schema.ToolCall{\n\t\t\t\t\t\t{ID: \"tc1\", Function: schema.FunctionCall{Name: \"test_tool\", Arguments: \"{}\"}},\n\t\t\t\t\t}), nil\n\t\t\t\t}\n\t\t\t\treturn schema.AssistantMessage(\"final response\", nil), nil\n\t\t\t}).AnyTimes()\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\ttoolCalled := false\n\t\ttestTool := &fakeToolForTest{tarCount: 0}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for AfterChatModel ReAct scenario\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tMiddlewares: []AgentMiddleware{\n\t\t\t\t{\n\t\t\t\t\tAfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error {\n\t\t\t\t\t\tlastMsg := state.Messages[len(state.Messages)-1]\n\t\t\t\t\t\tif len(lastMsg.ToolCalls) > 0 {\n\t\t\t\t\t\t\ttoolCalled = true\n\t\t\t\t\t\t\tstate.Messages[len(state.Messages)-1] = schema.AssistantMessage(\"modified to remove tool call\", nil)\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn nil\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\n\t\tvar events []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iterator.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tassert.True(t, toolCalled)\n\t\tassert.Equal(t, 1, generateCount)\n\n\t\tassert.Equal(t, 1, len(events))\n\t\tevent := events[0]\n\t\tassert.NotNil(t, event.Output)\n\t\tassert.NotNil(t, event.Output.MessageOutput)\n\t\tassert.Equal(t, \"first response with tool call\", event.Output.MessageOutput.Message.Content)\n\t\tassert.Len(t, event.Output.MessageOutput.Message.ToolCalls, 1)\n\t})\n\n\tt.Run(\"AfterChatModel_ReAct_AppendToolCall_AffectsFlow\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tgenerateCount := 0\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\tgenerateCount++\n\t\t\t\tif generateCount == 1 {\n\t\t\t\t\treturn schema.AssistantMessage(\"first response no tool\", nil), nil\n\t\t\t\t}\n\t\t\t\treturn schema.AssistantMessage(\"final response\", nil), nil\n\t\t\t}).AnyTimes()\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\ttestTool := &fakeToolForTest{tarCount: 0}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for AfterChatModel ReAct append tool call\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tMiddlewares: []AgentMiddleware{\n\t\t\t\t{\n\t\t\t\t\tAfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error {\n\t\t\t\t\t\tif generateCount == 1 {\n\t\t\t\t\t\t\tstate.Messages[len(state.Messages)-1] = schema.AssistantMessage(\"modified with tool call\", []schema.ToolCall{\n\t\t\t\t\t\t\t\t{ID: \"tc1\", Function: schema.FunctionCall{Name: \"test_tool\", Arguments: \"{}\"}},\n\t\t\t\t\t\t\t})\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn nil\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"Hello\")},\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\n\t\tvar events []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iterator.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tassert.Equal(t, 2, generateCount)\n\n\t\tassert.Equal(t, 3, len(events))\n\n\t\tevent0 := events[0]\n\t\tassert.NotNil(t, event0.Output)\n\t\tassert.NotNil(t, event0.Output.MessageOutput)\n\t\tassert.Equal(t, \"first response no tool\", event0.Output.MessageOutput.Message.Content)\n\t\tassert.Empty(t, event0.Output.MessageOutput.Message.ToolCalls)\n\n\t\tevent2 := events[2]\n\t\tassert.NotNil(t, event2.Output)\n\t\tassert.NotNil(t, event2.Output.MessageOutput)\n\t\tassert.Equal(t, \"final response\", event2.Output.MessageOutput.Message.Content)\n\t})\n\n\t// Test with streaming output\n\tt.Run(\"StreamOutput\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Create a stream reader for the mock response\n\t\tsr := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\tschema.AssistantMessage(\"Hello\", nil),\n\t\t\tschema.AssistantMessage(\", I am\", nil),\n\t\t\tschema.AssistantMessage(\" an AI assistant.\", nil),\n\t\t})\n\n\t\t// Set up expectations for the mock model\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(sr, nil).\n\t\t\tTimes(1)\n\n\t\t// Create a ChatModelAgent\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for unit testing\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, agent)\n\n\t\t// Run the agent with streaming enabled\n\t\tinput := &AgentInput{\n\t\t\tMessages:        []Message{schema.UserMessage(\"Hello, who are you?\")},\n\t\t\tEnableStreaming: true,\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\t\tassert.NotNil(t, iterator)\n\n\t\t// Get the event from the iterator\n\t\tevent, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event)\n\t\tassert.Nil(t, event.Err)\n\t\tassert.NotNil(t, event.Output)\n\t\tassert.NotNil(t, event.Output.MessageOutput)\n\t\tassert.True(t, event.Output.MessageOutput.IsStreaming)\n\n\t\t// No more events\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\t})\n\n\t// Test error handling\n\tt.Run(\"ErrorHandling\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model to return an error\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(nil, errors.New(\"model error\")).\n\t\t\tTimes(1)\n\n\t\t// Create a ChatModelAgent\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for unit testing\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, agent)\n\n\t\t// Run the agent\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"Hello, who are you?\")},\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\t\tassert.NotNil(t, iterator)\n\n\t\t// Get the event from the iterator, should contain an error\n\t\tevent, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event)\n\t\tassert.NotNil(t, event.Err)\n\t\tassert.Contains(t, event.Err.Error(), \"model error\")\n\n\t\t// No more events\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\t})\n\n\t// Test with tools\n\tt.Run(\"WithTools\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a fake tool for testing\n\t\tfakeTool := &fakeToolForTest{\n\t\t\ttarCount: 1,\n\t\t}\n\n\t\tinfo, err := fakeTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Using tool\",\n\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"name\": \"test user\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t}}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Task completed\", nil), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create a ChatModelAgent with tools\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for unit testing\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, agent)\n\n\t\t// Run the agent\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"Use the test tool\")},\n\t\t}\n\t\titerator := agent.Run(ctx, input)\n\t\tassert.NotNil(t, iterator)\n\n\t\t// Get events from the iterator\n\t\t// First event should be the model output with tool call\n\t\tevent1, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event1)\n\t\tassert.Nil(t, event1.Err)\n\t\tassert.NotNil(t, event1.Output)\n\t\tassert.NotNil(t, event1.Output.MessageOutput)\n\t\tassert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role)\n\n\t\t// Second event should be the tool output\n\t\tevent2, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event2)\n\t\tassert.Nil(t, event2.Err)\n\t\tassert.NotNil(t, event2.Output)\n\t\tassert.NotNil(t, event2.Output.MessageOutput)\n\t\tassert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role)\n\n\t\t// Third event should be the final model output\n\t\tevent3, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event3)\n\t\tassert.Nil(t, event3.Err)\n\t\tassert.NotNil(t, event3.Output)\n\t\tassert.NotNil(t, event3.Output.MessageOutput)\n\t\tassert.Equal(t, schema.Assistant, event3.Output.MessageOutput.Role)\n\n\t\t// No more events\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\t})\n}\n\n// TestExitTool tests the Exit tool functionality\nfunc TestExitTool(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock chat model\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// Set up expectations for the mock model\n\t// First call: model generates a message with Exit tool call\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"I'll exit with a final result\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      \"exit\",\n\t\t\t\t\t\tArguments: `{\"final_result\": \"This is the final result\"}`},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\n\t// Model should implement WithTools\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t// Create an agent with the Exit tool\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Test agent with Exit tool\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tExit:        &ExitTool{},\n\t})\n\tassert.NoError(t, err)\n\tassert.NotNil(t, agent)\n\n\t// Run the agent\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Please exit with a final result\"),\n\t\t},\n\t}\n\titerator := agent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// First event: model output with tool call\n\tevent1, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event1)\n\tassert.Nil(t, event1.Err)\n\tassert.NotNil(t, event1.Output)\n\tassert.NotNil(t, event1.Output.MessageOutput)\n\tassert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role)\n\n\t// Second event: tool output (Exit)\n\tevent2, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event2)\n\tassert.Nil(t, event2.Err)\n\tassert.NotNil(t, event2.Output)\n\tassert.NotNil(t, event2.Output.MessageOutput)\n\tassert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role)\n\n\t// Verify the action is Exit\n\tassert.NotNil(t, event2.Action)\n\tassert.True(t, event2.Action.Exit)\n\n\t// Verify the final result\n\tassert.Equal(t, \"This is the final result\", event2.Output.MessageOutput.Message.Content)\n\n\t// No more events\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestParallelReturnDirectlyToolCall(t *testing.T) {\n\tctx := context.Background()\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock chat model\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// Set up expectations for the mock model\n\t// First call: model generates a message with Exit tool call\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"I'll exit with a final result\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID:       \"tool-call-1\",\n\t\t\t\t\tFunction: schema.FunctionCall{Name: \"tool1\"},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tID:       \"tool-call-2\",\n\t\t\t\t\tFunction: schema.FunctionCall{Name: \"tool2\"},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tID:       \"tool-call-3\",\n\t\t\t\t\tFunction: schema.FunctionCall{Name: \"tool3\"},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\n\t// Model should implement WithTools\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t// Create an agent with the Exit tool\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Test agent with Exit tool\",\n\t\tInstruction: \"You are a helpful assistant.\",\n\t\tModel:       cm,\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{\n\t\t\t\t\t&myTool{name: \"tool1\", desc: \"tool1\", waitTime: time.Millisecond},\n\t\t\t\t\t&myTool{name: \"tool2\", desc: \"tool2\", waitTime: 10 * time.Millisecond},\n\t\t\t\t\t&myTool{name: \"tool3\", desc: \"tool3\", waitTime: 100 * time.Millisecond},\n\t\t\t\t},\n\t\t\t},\n\t\t\tReturnDirectly: map[string]bool{\n\t\t\t\t\"tool1\": true,\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\tassert.NotNil(t, agent)\n\n\tr := NewRunner(ctx, RunnerConfig{\n\t\tAgent: agent,\n\t})\n\titer := r.Query(ctx, \"\")\n\ttimes := 0\n\tfor {\n\t\te, ok := iter.Next()\n\t\tif !ok {\n\t\t\tassert.Equal(t, 4, times)\n\t\t\tbreak\n\t\t}\n\t\tif times == 3 {\n\t\t\tassert.Equal(t, \"tool1\", e.Output.MessageOutput.Message.ToolName)\n\t\t}\n\t\ttimes++\n\t}\n}\n\nfunc TestConcurrentSameToolSendToolGenActionUsesToolCallID(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"tools\", []schema.ToolCall{\n\t\t\t{ID: \"id1\", Function: schema.FunctionCall{Name: \"action_tool\", Arguments: \"A\"}},\n\t\t\t{ID: \"id2\", Function: schema.FunctionCall{Name: \"action_tool\", Arguments: \"B\"}},\n\t\t}), nil).\n\t\tTimes(1)\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).\n\t\tTimes(1)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Agent with action tool\",\n\t\tInstruction: \"\",\n\t\tModel:       cm,\n\t\tToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{actionTool{}}}},\n\t})\n\tassert.NoError(t, err)\n\n\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"go\")}})\n\tseen := map[string]bool{}\n\tfor {\n\t\te, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.Message != nil && e.Output.MessageOutput.Message.Role == schema.Tool {\n\t\t\tif e.Action != nil && e.Action.CustomizedAction != nil {\n\t\t\t\tif s, ok := e.Action.CustomizedAction.(string); ok {\n\t\t\t\t\tseen[s] = true\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\tassert.True(t, seen[\"A\"])\n\tassert.True(t, seen[\"B\"])\n}\n\ntype myTool struct {\n\tname     string\n\tdesc     string\n\twaitTime time.Duration\n}\n\nfunc (m *myTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: m.name,\n\t\tDesc: m.desc,\n\t}, nil\n}\n\nfunc (m *myTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\ttime.Sleep(m.waitTime)\n\treturn \"success\", nil\n}\n\ntype actionTool struct{}\n\nfunc (a actionTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"action_tool\", Desc: \"action tool\"}, nil\n}\n\nfunc (a actionTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\t_ = SendToolGenAction(ctx, \"action_tool\", &AgentAction{CustomizedAction: argumentsInJSON})\n\treturn \"ok\", nil\n}\n\ntype streamActionTool struct{}\n\nfunc (s streamActionTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"action_tool_stream\", Desc: \"action stream tool\"}, nil\n}\n\nfunc (s streamActionTool) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) {\n\t_ = SendToolGenAction(ctx, \"action_tool_stream\", &AgentAction{CustomizedAction: argumentsInJSON})\n\tsr, sw := schema.Pipe[string](1)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\t_ = sw.Send(\"o\", nil)\n\t\t_ = sw.Send(\"k\", nil)\n\t}()\n\treturn sr, nil\n}\n\ntype legacyStreamActionTool struct{}\n\nfunc (s legacyStreamActionTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"legacy_action_tool_stream\", Desc: \"legacy action stream tool\"}, nil\n}\n\nfunc (s legacyStreamActionTool) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) {\n\t_ = compose.ProcessState(ctx, func(ctx context.Context, st *State) error {\n\t\tst.setToolGenAction(\"legacy_action_tool_stream\", &AgentAction{CustomizedAction: argumentsInJSON})\n\t\treturn nil\n\t})\n\tsr, sw := schema.Pipe[string](1)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\t_ = sw.Send(\"o\", nil)\n\t\t_ = sw.Send(\"k\", nil)\n\t}()\n\treturn sr, nil\n}\n\n// TestChatModelAgentOutputKey tests the outputKey configuration and setOutputToSession function\nfunc TestChatModelAgentOutputKey(t *testing.T) {\n\t// Test outputKey configuration - stores output in session\n\tt.Run(\"OutputKeyStoresInSession\", func(t *testing.T) {\n\t\tfor i := 0; i < 1000; i++ {\n\n\t\t}\n\t\tctx := context.Background()\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Hello, I am an AI assistant.\", nil), nil).\n\t\t\tTimes(1)\n\n\t\t// Create a ChatModelAgent with outputKey configured\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for unit testing\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t\tOutputKey:   \"agent_output\", // This should store output in session\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, agent)\n\n\t\t// Initialize a run context to enable session storage\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{\n\t\t\t\tschema.UserMessage(\"Hello, who are you?\"),\n\t\t\t},\n\t\t}\n\t\tctx, runCtx := initRunCtx(ctx, \"TestAgent\", input)\n\t\tassert.NotNil(t, runCtx)\n\t\tassert.NotNil(t, runCtx.Session)\n\n\t\t// Run the agent\n\t\titerator := agent.Run(ctx, input)\n\t\tassert.NotNil(t, iterator)\n\n\t\t// Get the event from the iterator\n\t\tevent, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event)\n\t\tassert.Nil(t, event.Err)\n\n\t\t// Verify the message content\n\t\tmsg := event.Output.MessageOutput.Message\n\t\tassert.Equal(t, \"Hello, I am an AI assistant.\", msg.Content)\n\n\t\t// Verify that the output was stored in the session\n\t\ttime.AfterFunc(100*time.Millisecond, func() {\n\t\t\tsessionValues := GetSessionValues(ctx)\n\t\t\tassert.Contains(t, sessionValues, \"agent_output\")\n\t\t\tassert.Equal(t, \"Hello, I am an AI assistant.\", sessionValues[\"agent_output\"])\n\t\t})\n\n\t\t// No more events\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\t})\n\n\t// Test outputKey configuration with streaming - stores concatenated output in session\n\tt.Run(\"OutputKeyWithStreamingStoresInSession\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Create a stream reader for the mock response\n\t\tsr := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\tschema.AssistantMessage(\"Hello\", nil),\n\t\t\tschema.AssistantMessage(\", I am\", nil),\n\t\t\tschema.AssistantMessage(\" an AI assistant.\", nil),\n\t\t})\n\n\t\t// Set up expectations for the mock model\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(sr, nil).\n\t\t\tTimes(1)\n\n\t\t// Create a ChatModelAgent with outputKey configured\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent for unit testing\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t\tOutputKey:   \"agent_output\", // This should store concatenated stream in session\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, agent)\n\n\t\t// Initialize a run context to enable session storage\n\t\tinput := &AgentInput{\n\t\t\tMessages:        []Message{schema.UserMessage(\"Hello, who are you?\")},\n\t\t\tEnableStreaming: true,\n\t\t}\n\t\tctx, runCtx := initRunCtx(ctx, \"TestAgent\", input)\n\t\tassert.NotNil(t, runCtx)\n\t\tassert.NotNil(t, runCtx.Session)\n\n\t\t// Run the agent\n\t\titerator := agent.Run(ctx, input)\n\t\tassert.NotNil(t, iterator)\n\n\t\t// Get the event from the iterator\n\t\tevent, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event)\n\t\tassert.Nil(t, event.Err)\n\t\tassert.True(t, event.Output.MessageOutput.IsStreaming)\n\n\t\ttime.AfterFunc(100*time.Millisecond, func() {\n\t\t\t// Verify that the concatenated output was stored in the session\n\t\t\tsessionValues := GetSessionValues(ctx)\n\t\t\tassert.Contains(t, sessionValues, \"agent_output\")\n\t\t\tassert.Equal(t, \"Hello, I am an AI assistant.\", sessionValues[\"agent_output\"])\n\t\t})\n\n\t\t// No more events\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\t})\n\n\t// Test setOutputToSession function directly - regular message\n\tt.Run(\"SetOutputToSessionRegularMessage\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Initialize a run context to enable session storage\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"test\")},\n\t\t}\n\t\tctx, runCtx := initRunCtx(ctx, \"TestAgent\", input)\n\t\tassert.NotNil(t, runCtx)\n\t\tassert.NotNil(t, runCtx.Session)\n\n\t\t// Test with a regular message\n\t\tmsg := schema.AssistantMessage(\"Test response\", nil)\n\t\terr := setOutputToSession(ctx, msg, nil, \"test_output\")\n\t\tassert.NoError(t, err)\n\n\t\t// Verify the message content was stored\n\t\tsessionValues := GetSessionValues(ctx)\n\t\tassert.Contains(t, sessionValues, \"test_output\")\n\t\tassert.Equal(t, \"Test response\", sessionValues[\"test_output\"])\n\t})\n\n\t// Test setOutputToSession function directly - streaming message\n\tt.Run(\"SetOutputToSessionStreamingMessage\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Initialize a run context to enable session storage\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"test\")},\n\t\t}\n\t\tctx, runCtx := initRunCtx(ctx, \"TestAgent\", input)\n\t\tassert.NotNil(t, runCtx)\n\t\tassert.NotNil(t, runCtx.Session)\n\n\t\t// Test with a streaming message\n\t\tsr := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\tschema.AssistantMessage(\"Stream\", nil),\n\t\t\tschema.AssistantMessage(\" response\", nil),\n\t\t\tschema.AssistantMessage(\" content\", nil),\n\t\t})\n\t\terr := setOutputToSession(ctx, nil, sr, \"test_output\")\n\t\tassert.NoError(t, err)\n\n\t\t// Verify the concatenated stream content was stored\n\t\tsessionValues := GetSessionValues(ctx)\n\t\tassert.Contains(t, sessionValues, \"test_output\")\n\t\tassert.Equal(t, \"Stream response content\", sessionValues[\"test_output\"])\n\t})\n\n\t// Test setOutputToSession function directly - error case\n\tt.Run(\"SetOutputToSessionErrorCase\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Initialize a run context to enable session storage\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"test\")},\n\t\t}\n\t\tctx, runCtx := initRunCtx(ctx, \"TestAgent\", input)\n\t\tassert.NotNil(t, runCtx)\n\t\tassert.NotNil(t, runCtx.Session)\n\n\t\t// Test with an invalid stream (simulate error)\n\t\t// Create a stream that will fail when concatenated\n\t\tsr := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\tschema.AssistantMessage(\"test\", nil),\n\t\t})\n\t\t// Close the stream to simulate an error condition\n\t\tsr.Close()\n\n\t\t// This should return an error because the stream is closed\n\t\terr := setOutputToSession(ctx, nil, sr, \"test_output\")\n\t\t// Note: The actual behavior may vary depending on the stream implementation\n\t\t// Some streams may not error when closed, so we'll accept either outcome\n\t\tif err != nil {\n\t\t\t// If there's an error, verify nothing was stored\n\t\t\tsessionValues := GetSessionValues(ctx)\n\t\t\tassert.NotContains(t, sessionValues, \"test_output\")\n\t\t} else {\n\t\t\t// If no error, verify the content was stored\n\t\t\tsessionValues := GetSessionValues(ctx)\n\t\t\tassert.Contains(t, sessionValues, \"test_output\")\n\t\t\tassert.Equal(t, \"test\", sessionValues[\"test_output\"])\n\t\t}\n\t})\n\n\t// Test outputKey with React workflow (tools enabled)\n\tt.Run(\"OutputKeyWithReactWorkflow\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Create a simple tool for testing\n\t\tfakeTool := &fakeToolForTest{\n\t\t\ttarCount: 1,\n\t\t}\n\n\t\t// Set up expectations for the mock model - it will generate a final response\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Final response from React workflow\", nil), nil).\n\t\t\tTimes(1)\n\t\t// Model should implement WithTools\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create a ChatModelAgent with outputKey and tools configured\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent with tools\",\n\t\t\tInstruction: \"You are a helpful assistant.\",\n\t\t\tModel:       cm,\n\t\t\tOutputKey:   \"agent_output\",\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, agent)\n\n\t\t// Initialize a run context to enable session storage\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{schema.UserMessage(\"Use the tool\")},\n\t\t}\n\t\tctx, runCtx := initRunCtx(ctx, \"TestAgent\", input)\n\t\tassert.NotNil(t, runCtx)\n\t\tassert.NotNil(t, runCtx.Session)\n\n\t\t// Run the agent\n\t\titerator := agent.Run(ctx, input)\n\t\tassert.NotNil(t, iterator)\n\n\t\t// Get the event from the iterator\n\t\tevent, ok := iterator.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event)\n\t\tassert.Nil(t, event.Err)\n\n\t\t// Verify the message content\n\t\tmsg := event.Output.MessageOutput.Message\n\t\tassert.Equal(t, \"Final response from React workflow\", msg.Content)\n\n\t\t// Verify that the output was stored in the session\n\t\ttime.AfterFunc(time.Millisecond*10, func() {\n\t\t\tsessionValues := GetSessionValues(ctx)\n\t\t\tassert.Contains(t, sessionValues, \"agent_output\")\n\t\t\tassert.Equal(t, \"Final response from React workflow\", sessionValues[\"agent_output\"])\n\t\t})\n\n\t\t// No more events\n\t\t_, ok = iterator.Next()\n\t\tassert.False(t, ok)\n\t})\n}\n\nfunc TestConcurrentSameStreamToolSendToolGenActionUsesToolCallID(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"tools\", []schema.ToolCall{\n\t\t\t{ID: \"sid1\", Function: schema.FunctionCall{Name: \"action_tool_stream\", Arguments: \"SA\"}},\n\t\t\t{ID: \"sid2\", Function: schema.FunctionCall{Name: \"action_tool_stream\", Arguments: \"SB\"}},\n\t\t}), nil).\n\t\tTimes(1)\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).\n\t\tTimes(1)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Agent with stream action tool\",\n\t\tInstruction: \"\",\n\t\tModel:       cm,\n\t\tToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{streamActionTool{}}}},\n\t})\n\tassert.NoError(t, err)\n\n\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"go\")}})\n\tseen := map[string]bool{}\n\tfor {\n\t\te, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif e.Output != nil && e.Output.MessageOutput != nil {\n\t\t\tif e.Output.MessageOutput.IsStreaming {\n\t\t\t\tif e.Action != nil && e.Action.CustomizedAction != nil {\n\t\t\t\t\tif s, ok := e.Action.CustomizedAction.(string); ok {\n\t\t\t\t\t\tseen[s] = true\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\tassert.True(t, seen[\"SA\"])\n\tassert.True(t, seen[\"SB\"])\n}\n\nfunc TestStreamToolLegacyNameKeyFallback(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"tools\", []schema.ToolCall{\n\t\t\t{ID: \"lsid1\", Function: schema.FunctionCall{Name: \"legacy_action_tool_stream\", Arguments: \"LA\"}},\n\t\t}), nil).\n\t\tTimes(1)\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).\n\t\tTimes(1)\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Agent with legacy stream action tool\",\n\t\tInstruction: \"\",\n\t\tModel:       cm,\n\t\tToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{legacyStreamActionTool{}}}},\n\t})\n\tassert.NoError(t, err)\n\n\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"go\")}})\n\tfound := false\n\tfor {\n\t\te, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.IsStreaming {\n\t\t\tif e.Action != nil && e.Action.CustomizedAction != nil {\n\t\t\t\tif s, ok := e.Action.CustomizedAction.(string); ok {\n\t\t\t\t\tfound = (s == \"LA\")\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\tassert.True(t, found)\n}\n\nfunc TestChatModelAgent_ToolResultMiddleware_EmitsFinalResult(t *testing.T) {\n\toriginalResult := \"original_result\"\n\tmodifiedResult := \"modified_by_middleware\"\n\n\tresultModifyingMiddleware := compose.ToolMiddleware{\n\t\tInvokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\t\t\toutput, err := next(ctx, input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\toutput.Result = modifiedResult\n\t\t\t\treturn output, nil\n\t\t\t}\n\t\t},\n\t\tStreamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\t\t\toutput, err := next(ctx, input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\toutput.Result = schema.StreamReaderFromArray([]string{modifiedResult})\n\t\t\t\treturn output, nil\n\t\t\t}\n\t\t},\n\t}\n\n\tt.Run(\"Invoke\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\ttestTool := &simpleToolForMiddlewareTest{name: \"test_tool\", result: originalResult}\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tinfo, err := testTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"\",\n\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"input\": \"test\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"final response\", nil), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"test_agent\",\n\t\t\tDescription: \"test agent with middleware\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools:               []tool.BaseTool{testTool},\n\t\t\t\t\tToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false, CheckPointStore: newBridgeStore()})\n\t\tit := r.Run(ctx, []Message{schema.UserMessage(\"call the tool\")})\n\n\t\tvar toolResultEvents []*AgentEvent\n\t\tfor {\n\t\t\tev, ok := it.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif ev.Output != nil && ev.Output.MessageOutput != nil &&\n\t\t\t\tev.Output.MessageOutput.Message != nil &&\n\t\t\t\tev.Output.MessageOutput.Message.Role == schema.Tool {\n\t\t\t\ttoolResultEvents = append(toolResultEvents, ev)\n\t\t\t}\n\t\t}\n\n\t\tassert.NotEmpty(t, toolResultEvents, \"should have at least one tool result event\")\n\t\tfor _, ev := range toolResultEvents {\n\t\t\tassert.Equal(t, modifiedResult, ev.Output.MessageOutput.Message.Content,\n\t\t\t\t\"tool result event should contain the middleware-modified result, not the original\")\n\t\t\tassert.NotEqual(t, originalResult, ev.Output.MessageOutput.Message.Content,\n\t\t\t\t\"tool result event should NOT contain the original result\")\n\t\t}\n\t})\n\n\tt.Run(\"Stream\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\ttestTool := &simpleToolForMiddlewareTest{name: \"test_tool_stream\", result: originalResult}\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tinfo, err := testTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"input\": \"test\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"final response\", nil),\n\t\t\t}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"test_agent\",\n\t\t\tDescription: \"test agent with middleware\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools:               []tool.BaseTool{testTool},\n\t\t\t\t\tToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true, CheckPointStore: newBridgeStore()})\n\t\tit := r.Run(ctx, []Message{schema.UserMessage(\"call the tool\")})\n\n\t\tvar toolResultContents []string\n\t\tfor {\n\t\t\tev, ok := it.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif ev.Output != nil && ev.Output.MessageOutput != nil {\n\t\t\t\tif ev.Output.MessageOutput.Message != nil &&\n\t\t\t\t\tev.Output.MessageOutput.Message.Role == schema.Tool {\n\t\t\t\t\ttoolResultContents = append(toolResultContents, ev.Output.MessageOutput.Message.Content)\n\t\t\t\t}\n\t\t\t\tif ev.Output.MessageOutput.IsStreaming &&\n\t\t\t\t\tev.Output.MessageOutput.MessageStream != nil &&\n\t\t\t\t\tev.Output.MessageOutput.Role == schema.Tool {\n\t\t\t\t\tvar msgs []*schema.Message\n\t\t\t\t\tfor {\n\t\t\t\t\t\tmsg, err := ev.Output.MessageOutput.MessageStream.Recv()\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t\tmsgs = append(msgs, msg)\n\t\t\t\t\t}\n\t\t\t\t\tif len(msgs) > 0 {\n\t\t\t\t\t\tconcated, err := schema.ConcatMessages(msgs)\n\t\t\t\t\t\tif err == nil {\n\t\t\t\t\t\t\ttoolResultContents = append(toolResultContents, concated.Content)\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tassert.NotEmpty(t, toolResultContents, \"should have at least one tool result event\")\n\t\tfor _, content := range toolResultContents {\n\t\t\tassert.Equal(t, modifiedResult, content,\n\t\t\t\t\"tool result event should contain the middleware-modified result, not the original\")\n\t\t\tassert.NotEqual(t, originalResult, content,\n\t\t\t\t\"tool result event should NOT contain the original result\")\n\t\t}\n\t})\n}\n\ntype simpleToolForMiddlewareTest struct {\n\tname   string\n\tresult string\n}\n\nfunc (s *simpleToolForMiddlewareTest) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: s.name,\n\t\tDesc: \"simple tool\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"input\": {\n\t\t\t\t\tDesc:     \"input\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t\tType:     schema.String,\n\t\t\t\t},\n\t\t\t}),\n\t}, nil\n}\n\nfunc (s *simpleToolForMiddlewareTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn s.result, nil\n}\n\nfunc (s *simpleToolForMiddlewareTest) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) {\n\treturn schema.StreamReaderFromArray([]string{s.result}), nil\n}\n\nfunc TestGetComposeOptions(t *testing.T) {\n\tt.Run(\"WithChatModelOptions\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar capturedTemperature float32\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\toptions := model.GetCommonOptions(&model.Options{}, opts...)\n\t\t\t\tif options.Temperature != nil {\n\t\t\t\t\tcapturedTemperature = *options.Temperature\n\t\t\t\t}\n\t\t\t\treturn schema.AssistantMessage(\"response\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\ttemp := float32(0.7)\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}},\n\t\t\tWithChatModelOptions([]model.Option{model.WithTemperature(temp)}))\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, temp, capturedTemperature, \"Temperature should be passed through WithChatModelOptions\")\n\t})\n\n\tt.Run(\"WithToolOptions\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar toolOptionsCaptured bool\n\t\ttestTool := &toolOptionCapturingTool{\n\t\t\tname: \"test_tool\",\n\t\t\tonRun: func(opts []tool.Option) {\n\t\t\t\tif len(opts) > 0 {\n\t\t\t\t\ttoolOptionsCaptured = true\n\t\t\t\t}\n\t\t\t},\n\t\t}\n\t\tinfo, _ := testTool.Info(ctx)\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Using tool\", []schema.ToolCall{\n\t\t\t\t{ID: \"call1\", Function: schema.FunctionCall{Name: info.Name, Arguments: \"{}\"}},\n\t\t\t}), nil).Times(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}},\n\t\t\tWithToolOptions([]tool.Option{testToolOption(\"test_value\")}))\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.True(t, toolOptionsCaptured, \"Tool options should be passed through WithToolOptions\")\n\t})\n\n}\n\ntype toolOptionCapturingTool struct {\n\tname  string\n\tonRun func(opts []tool.Option)\n}\n\nfunc (t *toolOptionCapturingTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: t.name, Desc: t.name + \" description\"}, nil\n}\n\nfunc (t *toolOptionCapturingTool) InvokableRun(_ context.Context, _ string, opts ...tool.Option) (string, error) {\n\tif t.onRun != nil {\n\t\tt.onRun(opts)\n\t}\n\treturn t.name + \" result\", nil\n}\n\ntype testToolOptions struct {\n\tvalue string\n}\n\nfunc testToolOption(value string) tool.Option {\n\treturn tool.WrapImplSpecificOptFn(func(o *testToolOptions) {\n\t\to.value = value\n\t})\n}\n\ntype errorTool struct {\n\tinfoErr error\n}\n\nfunc (e *errorTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn nil, e.infoErr\n}\n\nfunc (e *errorTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn \"\", nil\n}\n\nfunc TestChatModelAgent_PrepareExecContextError(t *testing.T) {\n\tt.Run(\"Run_WithToolInfoError_ReturnsError\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\texpectedErr := errors.New(\"tool info error\")\n\t\terrTool := &errorTool{infoErr: expectedErr}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{errTool},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\n\t\tevent, ok := iter.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event.Err)\n\t\tassert.Contains(t, event.Err.Error(), \"tool info error\")\n\n\t\t_, ok = iter.Next()\n\t\tassert.False(t, ok)\n\t})\n\n\tt.Run(\"Resume_WithToolInfoError_ReturnsError\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\texpectedErr := errors.New(\"tool info error for resume\")\n\t\terrTool := &errorTool{infoErr: expectedErr}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{errTool},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Resume(ctx, &ResumeInfo{\n\t\t\tInterruptState:  []byte(\"dummy\"),\n\t\t\tEnableStreaming: false,\n\t\t})\n\n\t\tevent, ok := iter.Next()\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, event.Err)\n\t\tassert.Contains(t, event.Err.Error(), \"tool info error for resume\")\n\n\t\t_, ok = iter.Next()\n\t\tassert.False(t, ok)\n\t})\n}\n\nfunc TestPreprocessComposeCheckpoint_MigrateErrorIsReturned(t *testing.T) {\n\tin := []byte(\"prefix\\u0015\" + stateGobNameV080 + \"suffix\")\n\t_, err := preprocessComposeCheckpoint(in)\n\tassert.Error(t, err)\n}\n"
  },
  {
    "path": "adk/config.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport \"github.com/cloudwego/eino/adk/internal\"\n\n// Language represents the language setting for the ADK built-in prompts.\ntype Language = internal.Language\n\nconst (\n\t// LanguageEnglish represents English language.\n\tLanguageEnglish Language = internal.LanguageEnglish\n\t// LanguageChinese represents Chinese language.\n\tLanguageChinese Language = internal.LanguageChinese\n)\n\n// SetLanguage sets the language for the ADK built-in prompts.\n// The default language is English if not explicitly set.\nfunc SetLanguage(lang Language) error {\n\treturn internal.SetLanguage(lang)\n}\n"
  },
  {
    "path": "adk/deterministic_transfer.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"runtime/debug\"\n\t\"sync\"\n\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc init() {\n\tschema.RegisterName[*deterministicTransferState](\"_eino_adk_deterministic_transfer_state\")\n}\n\ntype deterministicTransferState struct {\n\tEventList []*agentEventWrapper\n}\n\n// AgentWithDeterministicTransferTo wraps an agent to transfer to given agents deterministically.\nfunc AgentWithDeterministicTransferTo(_ context.Context, config *DeterministicTransferConfig) Agent {\n\tif ra, ok := config.Agent.(ResumableAgent); ok {\n\t\treturn &resumableAgentWithDeterministicTransferTo{\n\t\t\tagent:        ra,\n\t\t\ttoAgentNames: config.ToAgentNames,\n\t\t}\n\t}\n\treturn &agentWithDeterministicTransferTo{\n\t\tagent:        config.Agent,\n\t\ttoAgentNames: config.ToAgentNames,\n\t}\n}\n\ntype agentWithDeterministicTransferTo struct {\n\tagent        Agent\n\ttoAgentNames []string\n}\n\nfunc (a *agentWithDeterministicTransferTo) Description(ctx context.Context) string {\n\treturn a.agent.Description(ctx)\n}\n\nfunc (a *agentWithDeterministicTransferTo) Name(ctx context.Context) string {\n\treturn a.agent.Name(ctx)\n}\n\nfunc (a *agentWithDeterministicTransferTo) GetType() string {\n\tif typer, ok := a.agent.(components.Typer); ok {\n\t\treturn typer.GetType()\n\t}\n\treturn \"DeterministicTransfer\"\n}\n\nfunc (a *agentWithDeterministicTransferTo) Run(ctx context.Context,\n\tinput *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\n\tif fa, ok := a.agent.(*flowAgent); ok {\n\t\treturn runFlowAgentWithIsolatedSession(ctx, fa, input, a.toAgentNames, options...)\n\t}\n\n\taIter := a.agent.Run(ctx, input, options...)\n\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\tgo forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames)\n\n\treturn iterator\n}\n\ntype resumableAgentWithDeterministicTransferTo struct {\n\tagent        ResumableAgent\n\ttoAgentNames []string\n}\n\nfunc (a *resumableAgentWithDeterministicTransferTo) Description(ctx context.Context) string {\n\treturn a.agent.Description(ctx)\n}\n\nfunc (a *resumableAgentWithDeterministicTransferTo) Name(ctx context.Context) string {\n\treturn a.agent.Name(ctx)\n}\n\nfunc (a *resumableAgentWithDeterministicTransferTo) GetType() string {\n\tif typer, ok := a.agent.(components.Typer); ok {\n\t\treturn typer.GetType()\n\t}\n\treturn \"DeterministicTransfer\"\n}\n\nfunc (a *resumableAgentWithDeterministicTransferTo) Run(ctx context.Context,\n\tinput *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\n\tif fa, ok := a.agent.(*flowAgent); ok {\n\t\treturn runFlowAgentWithIsolatedSession(ctx, fa, input, a.toAgentNames, options...)\n\t}\n\n\taIter := a.agent.Run(ctx, input, options...)\n\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\tgo forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames)\n\n\treturn iterator\n}\n\nfunc (a *resumableAgentWithDeterministicTransferTo) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tif fa, ok := a.agent.(*flowAgent); ok {\n\t\treturn resumeFlowAgentWithIsolatedSession(ctx, fa, info, a.toAgentNames, opts...)\n\t}\n\n\taIter := a.agent.Resume(ctx, info, opts...)\n\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\tgo forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames)\n\n\treturn iterator\n}\n\nfunc forwardEventsAndAppendTransfer(iter *AsyncIterator[*AgentEvent],\n\tgenerator *AsyncGenerator[*AgentEvent], toAgentNames []string) {\n\n\tdefer func() {\n\t\tif panicErr := recover(); panicErr != nil {\n\t\t\tgenerator.Send(&AgentEvent{Err: safe.NewPanicErr(panicErr, debug.Stack())})\n\t\t}\n\t\tgenerator.Close()\n\t}()\n\n\tvar lastEvent *AgentEvent\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tgenerator.Send(event)\n\t\tlastEvent = event\n\t}\n\n\tif lastEvent != nil && lastEvent.Action != nil && (lastEvent.Action.Interrupted != nil || lastEvent.Action.Exit) {\n\t\treturn\n\t}\n\n\tsendTransferEvents(generator, toAgentNames)\n}\n\nfunc runFlowAgentWithIsolatedSession(ctx context.Context, fa *flowAgent, input *AgentInput,\n\ttoAgentNames []string, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\n\tparentSession := getSession(ctx)\n\tparentRunCtx := getRunCtx(ctx)\n\n\tisolatedSession := &runSession{\n\t\tValues:    parentSession.Values,\n\t\tvaluesMtx: parentSession.valuesMtx,\n\t}\n\tif isolatedSession.valuesMtx == nil {\n\t\tisolatedSession.valuesMtx = &sync.Mutex{}\n\t}\n\tif isolatedSession.Values == nil {\n\t\tisolatedSession.Values = make(map[string]any)\n\t}\n\n\tctx = setRunCtx(ctx, &runContext{\n\t\tSession:   isolatedSession,\n\t\tRootInput: parentRunCtx.RootInput,\n\t\tRunPath:   parentRunCtx.RunPath,\n\t})\n\n\titer := fa.Run(ctx, input, options...)\n\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\tgo handleFlowAgentEvents(ctx, iter, generator, isolatedSession, parentSession, toAgentNames)\n\n\treturn iterator\n}\n\nfunc resumeFlowAgentWithIsolatedSession(ctx context.Context, fa *flowAgent, info *ResumeInfo,\n\ttoAgentNames []string, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\n\tstate, ok := info.InterruptState.(*deterministicTransferState)\n\tif !ok || state == nil {\n\t\treturn genErrorIter(errors.New(\"invalid interrupt state for flowAgent resume in deterministic transfer\"))\n\t}\n\n\tparentSession := getSession(ctx)\n\tparentRunCtx := getRunCtx(ctx)\n\n\tisolatedSession := &runSession{\n\t\tValues:    parentSession.Values,\n\t\tvaluesMtx: parentSession.valuesMtx,\n\t\tEvents:    state.EventList,\n\t}\n\tif isolatedSession.valuesMtx == nil {\n\t\tisolatedSession.valuesMtx = &sync.Mutex{}\n\t}\n\tif isolatedSession.Values == nil {\n\t\tisolatedSession.Values = make(map[string]any)\n\t}\n\n\tctx = setRunCtx(ctx, &runContext{\n\t\tSession:   isolatedSession,\n\t\tRootInput: parentRunCtx.RootInput,\n\t\tRunPath:   parentRunCtx.RunPath,\n\t})\n\n\titer := fa.Resume(ctx, info, opts...)\n\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\tgo handleFlowAgentEvents(ctx, iter, generator, isolatedSession, parentSession, toAgentNames)\n\n\treturn iterator\n}\n\nfunc handleFlowAgentEvents(ctx context.Context, iter *AsyncIterator[*AgentEvent],\n\tgenerator *AsyncGenerator[*AgentEvent], isolatedSession, parentSession *runSession, toAgentNames []string) {\n\n\tdefer func() {\n\t\tif panicErr := recover(); panicErr != nil {\n\t\t\tgenerator.Send(&AgentEvent{Err: safe.NewPanicErr(panicErr, debug.Stack())})\n\t\t}\n\t\tgenerator.Close()\n\t}()\n\n\tvar lastEvent *AgentEvent\n\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\n\t\tif parentSession != nil && (event.Action == nil || event.Action.Interrupted == nil) {\n\t\t\tcopied := copyAgentEvent(event)\n\t\t\tsetAutomaticClose(copied)\n\t\t\tsetAutomaticClose(event)\n\t\t\tparentSession.addEvent(copied)\n\t\t}\n\n\t\tif event.Action != nil && event.Action.internalInterrupted != nil {\n\t\t\tlastEvent = event\n\t\t\tcontinue\n\t\t}\n\n\t\tgenerator.Send(event)\n\t\tlastEvent = event\n\t}\n\n\tif lastEvent != nil && lastEvent.Action != nil {\n\t\tif lastEvent.Action.internalInterrupted != nil {\n\t\t\tevents := isolatedSession.getEvents()\n\t\t\tstate := &deterministicTransferState{EventList: events}\n\t\t\tcompositeEvent := CompositeInterrupt(ctx, \"deterministic transfer wrapper interrupted\",\n\t\t\t\tstate, lastEvent.Action.internalInterrupted)\n\t\t\tgenerator.Send(compositeEvent)\n\t\t\treturn\n\t\t}\n\n\t\tif lastEvent.Action.Exit {\n\t\t\treturn\n\t\t}\n\t}\n\n\tsendTransferEvents(generator, toAgentNames)\n}\n\nfunc sendTransferEvents(generator *AsyncGenerator[*AgentEvent], toAgentNames []string) {\n\tfor _, toAgentName := range toAgentNames {\n\t\taMsg, tMsg := GenTransferMessages(context.Background(), toAgentName)\n\n\t\taEvent := EventFromMessage(aMsg, nil, schema.Assistant, \"\")\n\t\tgenerator.Send(aEvent)\n\n\t\ttEvent := EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)\n\t\ttEvent.Action = &AgentAction{\n\t\t\tTransferToAgent: &TransferToAgentAction{\n\t\t\t\tDestAgentName: toAgentName,\n\t\t\t},\n\t\t}\n\t\tgenerator.Send(tEvent)\n\t}\n}\n"
  },
  {
    "path": "adk/deterministic_transfer_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype dtTestStore struct {\n\tdata map[string][]byte\n}\n\nfunc newDTTestStore() *dtTestStore {\n\treturn &dtTestStore{data: make(map[string][]byte)}\n}\n\nfunc (s *dtTestStore) Set(_ context.Context, key string, value []byte) error {\n\ts.data[key] = value\n\treturn nil\n}\n\nfunc (s *dtTestStore) Get(_ context.Context, key string) ([]byte, bool, error) {\n\tv, ok := s.data[key]\n\treturn v, ok, nil\n}\n\ntype dtTestAgent struct {\n\tname     string\n\trunFn    func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]\n\tresumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]\n}\n\nfunc (a *dtTestAgent) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *dtTestAgent) Description(_ context.Context) string {\n\treturn a.name + \" description\"\n}\n\nfunc (a *dtTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\treturn a.runFn(ctx, input, options...)\n}\n\nfunc (a *dtTestAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tif a.resumeFn != nil {\n\t\treturn a.resumeFn(ctx, info, opts...)\n\t}\n\treturn a.runFn(ctx, &AgentInput{}, opts...)\n}\n\nfunc TestDeterministicTransferFlowAgentInterruptResume(t *testing.T) {\n\tctx := context.Background()\n\tstore := newDTTestStore()\n\n\tinterruptData := \"interrupt_data\"\n\tvar runCount int\n\n\tinnerAgent := &dtTestAgent{\n\t\tname: \"inner\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\trunCount++\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tgen.Send(EventFromMessage(schema.AssistantMessage(\"before interrupt\", nil), nil, schema.Assistant, \"\"))\n\t\t\t\tintEvent := Interrupt(ctx, interruptData)\n\t\t\t\tintEvent.Action.Interrupted.Data = interruptData\n\t\t\t\tgen.Send(intEvent)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\trunCount++\n\n\t\t\tassert.True(t, info.WasInterrupted, \"innerAgent resumeFn: should be interrupted\")\n\t\t\tassert.True(t, info.IsResumeTarget, \"innerAgent resumeFn: should be resume target\")\n\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tassert.NotNil(t, runCtx, \"innerAgent resumeFn: runCtx should not be nil\")\n\t\t\tassert.NotNil(t, runCtx.Session, \"innerAgent resumeFn: runCtx.Session should not be nil\")\n\n\t\t\tvar agentEvents []*AgentEvent\n\t\t\tfor _, ev := range runCtx.Session.Events {\n\t\t\t\tif ev.AgentEvent != nil {\n\t\t\t\t\tagentEvents = append(agentEvents, ev.AgentEvent)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tassert.Len(t, agentEvents, 1, \"innerAgent resumeFn: should have exactly 1 agent event\")\n\t\t\tif len(agentEvents) == 1 {\n\t\t\t\tev := agentEvents[0]\n\t\t\t\tassert.Equal(t, \"inner\", ev.AgentName, \"innerAgent resumeFn: event should be from inner agent\")\n\t\t\t\tassert.Equal(t, \"before interrupt\", ev.Output.MessageOutput.Message.Content, \"innerAgent resumeFn: event content should be 'before interrupt'\")\n\t\t\t\tassert.Len(t, ev.RunPath, 2, \"innerAgent resumeFn: RunPath should have 2 steps (outer agent, inner agent)\")\n\t\t\t\tif len(ev.RunPath) == 2 {\n\t\t\t\t\tassert.Equal(t, \"outer\", ev.RunPath[0].agentName, \"innerAgent resumeFn: RunPath[0] should be outer agent\")\n\t\t\t\t\tassert.Equal(t, \"inner\", ev.RunPath[1].agentName, \"innerAgent resumeFn: RunPath[1] should be inner agent\")\n\t\t\t\t}\n\t\t\t}\n\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tgen.Send(EventFromMessage(schema.AssistantMessage(\"after resume\", nil), nil, schema.Assistant, \"\"))\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tinnerFlowAgent := toFlowAgent(ctx, innerAgent)\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        innerFlowAgent,\n\t\tToAgentNames: []string{\"next_agent\"},\n\t})\n\n\touterAgent := &dtTestAgent{\n\t\tname: \"outer\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\treturn wrapped.Run(ctx, input, options...)\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.True(t, info.WasInterrupted, \"outerAgent resumeFn: should be interrupted\")\n\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tassert.NotNil(t, runCtx, \"outerAgent resumeFn: runCtx should not be nil\")\n\t\t\tassert.NotNil(t, runCtx.Session, \"outerAgent resumeFn: runCtx.Session should not be nil\")\n\n\t\t\tvar agentEvents []*AgentEvent\n\t\t\tfor _, ev := range runCtx.Session.Events {\n\t\t\t\tif ev.AgentEvent != nil {\n\t\t\t\t\tagentEvents = append(agentEvents, ev.AgentEvent)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tassert.Len(t, agentEvents, 1, \"outerAgent resumeFn: should have exactly 1 agent event\")\n\t\t\tif len(agentEvents) == 1 {\n\t\t\t\tev := agentEvents[0]\n\t\t\t\tassert.Equal(t, \"inner\", ev.AgentName, \"outerAgent resumeFn: event should be from inner agent (preserved original)\")\n\t\t\t\tassert.Equal(t, \"before interrupt\", ev.Output.MessageOutput.Message.Content, \"outerAgent resumeFn: event content should be 'before interrupt'\")\n\t\t\t\tassert.Len(t, ev.RunPath, 2, \"outerAgent resumeFn: RunPath should have 2 steps\")\n\t\t\t\tif len(ev.RunPath) == 2 {\n\t\t\t\t\tassert.Equal(t, \"outer\", ev.RunPath[0].agentName, \"outerAgent resumeFn: RunPath[0] should be outer agent\")\n\t\t\t\t\tassert.Equal(t, \"inner\", ev.RunPath[1].agentName, \"outerAgent resumeFn: RunPath[1] should be inner agent\")\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tra := wrapped.(ResumableAgent)\n\t\t\treturn ra.Resume(ctx, info, opts...)\n\t\t},\n\t}\n\n\touterFlowAgent := toFlowAgent(ctx, outerAgent)\n\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           outerFlowAgent,\n\t\tEnableStreaming: true,\n\t\tCheckPointStore: store,\n\t})\n\n\titer := runner.Run(ctx, []Message{schema.UserMessage(\"test\")}, WithCheckPointID(\"cp1\"))\n\n\tvar events []*AgentEvent\n\tvar interrupted bool\n\tvar interruptEvent *AgentEvent\n\tfor {\n\t\tev, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, ev)\n\t\tif ev.Action != nil && ev.Action.Interrupted != nil {\n\t\t\tinterrupted = true\n\t\t\tinterruptEvent = ev\n\t\t}\n\t}\n\n\tassert.Equal(t, 1, runCount, \"run should have been called once\")\n\tassert.True(t, interrupted, \"should have interrupted\")\n\tassert.Greater(t, len(events), 0, \"should have events\")\n\tif interruptEvent == nil {\n\t\tt.Fatal(\"should have interrupt event\")\n\t}\n\tassert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, \"should have interrupt contexts\")\n\n\t_, exists, err := store.Get(ctx, \"cp1\")\n\tassert.NoError(t, err)\n\tassert.True(t, exists, \"checkpoint should have been saved\")\n\n\tvar hasDeterministicTransferContext bool\n\tfor _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tt.Logf(\"InterruptContext: ID=%s, Info=%v, IsRootCause=%v, Addr=%v\", intCtx.ID, intCtx.Info, intCtx.IsRootCause, intCtx.Address)\n\t\tif intCtx.Info == \"deterministic transfer wrapper interrupted\" {\n\t\t\thasDeterministicTransferContext = true\n\t\t}\n\t\tfor parent := intCtx.Parent; parent != nil; parent = parent.Parent {\n\t\t\tt.Logf(\"  Parent: ID=%s, Info=%v, Addr=%v\", parent.ID, parent.Info, parent.Address)\n\t\t\tif parent.Info == \"deterministic transfer wrapper interrupted\" {\n\t\t\t\thasDeterministicTransferContext = true\n\t\t\t}\n\t\t}\n\t}\n\tassert.True(t, hasDeterministicTransferContext, \"should have deterministic transfer interrupt context\")\n\n\tvar rootCauseID string\n\tfor _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif intCtx.IsRootCause {\n\t\t\trootCauseID = intCtx.ID\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.NotEmpty(t, rootCauseID, \"should have root cause interrupt ID\")\n\n\tresumeIter, err := runner.ResumeWithParams(ctx, \"cp1\", &ResumeParams{\n\t\tTargets: map[string]any{rootCauseID: nil},\n\t})\n\tassert.NoError(t, err)\n\n\tvar resumeEvents []*AgentEvent\n\tvar resumeErr error\n\tvar hasTransfer bool\n\tfor {\n\t\tev, ok := resumeIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\n\t\tif ev.Err != nil {\n\t\t\tresumeErr = ev.Err\n\t\t}\n\t\tif ev.Action != nil && ev.Action.TransferToAgent != nil {\n\t\t\thasTransfer = true\n\t\t}\n\t\tresumeEvents = append(resumeEvents, ev)\n\t}\n\n\tassert.Equal(t, 2, runCount, \"inner agent should be called twice (once for initial, once for resume)\")\n\tassert.NotEmpty(t, resumeEvents, \"should have resume events\")\n\tassert.True(t, hasTransfer, \"should have transfer action after resume\")\n\tassert.Error(t, resumeErr, \"transfer should fail because next_agent doesn't exist\")\n\tassert.Contains(t, resumeErr.Error(), \"next_agent\", \"error should mention the missing agent\")\n}\n\nfunc TestDeterministicTransferRunPathPreserved(t *testing.T) {\n\tctx := context.Background()\n\tstore := newDTTestStore()\n\n\tvar collectedRunPaths [][]RunStep\n\n\tinnerAgent := &dtTestAgent{\n\t\tname: \"inner\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tev := EventFromMessage(schema.AssistantMessage(\"from inner\", nil), nil, schema.Assistant, \"\")\n\t\t\t\tgen.Send(ev)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tinnerFlowAgent := toFlowAgent(ctx, innerAgent)\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        innerFlowAgent,\n\t\tToAgentNames: []string{},\n\t})\n\n\touterAgent := &dtTestAgent{\n\t\tname: \"outer\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tinnerIter := wrapped.Run(ctx, input, options...)\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tfor {\n\t\t\t\t\tev, ok := innerIter.Next()\n\t\t\t\t\tif !ok {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\tcollectedRunPaths = append(collectedRunPaths, ev.RunPath)\n\t\t\t\t\tgen.Send(ev)\n\t\t\t\t}\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\touterFlowAgent := toFlowAgent(ctx, outerAgent)\n\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           outerFlowAgent,\n\t\tEnableStreaming: true,\n\t\tCheckPointStore: store,\n\t})\n\n\titer := runner.Run(ctx, []Message{schema.UserMessage(\"test\")}, WithCheckPointID(\"cp1\"))\n\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassert.NotEmpty(t, collectedRunPaths, \"should have collected RunPaths\")\n\tfor _, rp := range collectedRunPaths {\n\t\tassert.Len(t, rp, 2, \"RunPath should have 2 steps (outer agent, inner agent)\")\n\t\tif len(rp) == 2 {\n\t\t\tassert.Equal(t, \"outer\", rp[0].agentName, \"RunPath[0] should be outer agent\")\n\t\t\tassert.Equal(t, \"inner\", rp[1].agentName, \"RunPath[1] should be inner agent\")\n\t\t}\n\t}\n}\n\nfunc TestDeterministicTransferExitSkipsTransfer(t *testing.T) {\n\tctx := context.Background()\n\tstore := newDTTestStore()\n\n\tinnerAgent := &dtTestAgent{\n\t\tname: \"inner\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tev := EventFromMessage(schema.AssistantMessage(\"inner exits\", nil), nil, schema.Assistant, \"\")\n\t\t\t\tev.Action = &AgentAction{Exit: true}\n\t\t\t\tgen.Send(ev)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tinnerFlowAgent := toFlowAgent(ctx, innerAgent)\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        innerFlowAgent,\n\t\tToAgentNames: []string{\"next_agent\"},\n\t})\n\n\tvar outerSawExit bool\n\tvar transferGenerated bool\n\n\touterAgent := &dtTestAgent{\n\t\tname: \"outer\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tinnerIter := wrapped.Run(ctx, input, options...)\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tfor {\n\t\t\t\t\tev, ok := innerIter.Next()\n\t\t\t\t\tif !ok {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\tif ev.Action != nil && ev.Action.Exit {\n\t\t\t\t\t\touterSawExit = true\n\t\t\t\t\t}\n\t\t\t\t\tif ev.Action != nil && ev.Action.TransferToAgent != nil {\n\t\t\t\t\t\ttransferGenerated = true\n\t\t\t\t\t}\n\t\t\t\t\tgen.Send(ev)\n\t\t\t\t}\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\touterFlowAgent := toFlowAgent(ctx, outerAgent)\n\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           outerFlowAgent,\n\t\tEnableStreaming: true,\n\t\tCheckPointStore: store,\n\t})\n\n\titer := runner.Run(ctx, []Message{schema.UserMessage(\"test\")}, WithCheckPointID(\"cp1\"))\n\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassert.True(t, outerSawExit, \"outer should see exit event from inner\")\n\tassert.False(t, transferGenerated, \"transfer should not be generated when inner exits\")\n}\n\ntype nonFlowTestAgent struct {\n\tname     string\n\trunFn    func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]\n\tresumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]\n}\n\nfunc (a *nonFlowTestAgent) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *nonFlowTestAgent) Description(_ context.Context) string {\n\treturn a.name + \" description\"\n}\n\nfunc (a *nonFlowTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\treturn a.runFn(ctx, input, options...)\n}\n\nfunc (a *nonFlowTestAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tif a.resumeFn != nil {\n\t\treturn a.resumeFn(ctx, info, opts...)\n\t}\n\treturn a.runFn(ctx, &AgentInput{}, opts...)\n}\n\ntype nonResumableTestAgent struct {\n\tname  string\n\trunFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]\n}\n\nfunc (a *nonResumableTestAgent) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *nonResumableTestAgent) Description(_ context.Context) string {\n\treturn a.name + \" description\"\n}\n\nfunc (a *nonResumableTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\treturn a.runFn(ctx, input, options...)\n}\n\nfunc TestDeterministicTransferNonFlowAgent_ExitSkipsTransfer(t *testing.T) {\n\tctx := context.Background()\n\n\tagent := &nonFlowTestAgent{\n\t\tname: \"test_agent\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tev := EventFromMessage(schema.AssistantMessage(\"exiting\", nil), nil, schema.Assistant, \"\")\n\t\t\t\tev.Action = &AgentAction{Exit: true}\n\t\t\t\tgen.Send(ev)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        agent,\n\t\tToAgentNames: []string{\"next_agent\"},\n\t})\n\n\titer := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\n\tvar events []*AgentEvent\n\tvar sawExit bool\n\tvar sawTransfer bool\n\tfor {\n\t\tev, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, ev)\n\t\tif ev.Action != nil && ev.Action.Exit {\n\t\t\tsawExit = true\n\t\t}\n\t\tif ev.Action != nil && ev.Action.TransferToAgent != nil {\n\t\t\tsawTransfer = true\n\t\t}\n\t}\n\n\tassert.True(t, sawExit, \"should see exit event\")\n\tassert.False(t, sawTransfer, \"should NOT see transfer when exit is last event\")\n\tassert.Len(t, events, 1, \"should have exactly 1 event (exit)\")\n}\n\nfunc TestDeterministicTransferNonFlowAgent_AppendsTransfer(t *testing.T) {\n\tctx := context.Background()\n\n\tagent := &nonFlowTestAgent{\n\t\tname: \"test_agent\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tev := EventFromMessage(schema.AssistantMessage(\"normal output\", nil), nil, schema.Assistant, \"\")\n\t\t\t\tgen.Send(ev)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        agent,\n\t\tToAgentNames: []string{\"next_agent\"},\n\t})\n\n\titer := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\n\tvar events []*AgentEvent\n\tvar sawTransfer bool\n\tvar transferTarget string\n\tfor {\n\t\tev, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, ev)\n\t\tif ev.Action != nil && ev.Action.TransferToAgent != nil {\n\t\t\tsawTransfer = true\n\t\t\ttransferTarget = ev.Action.TransferToAgent.DestAgentName\n\t\t}\n\t}\n\n\tassert.True(t, sawTransfer, \"should see transfer event after normal completion\")\n\tassert.Equal(t, \"next_agent\", transferTarget, \"transfer target should be next_agent\")\n\tassert.Greater(t, len(events), 1, \"should have more than 1 event (output + transfer messages)\")\n}\n\nfunc TestDeterministicTransferNonFlowAgent_InterruptSkipsTransfer(t *testing.T) {\n\tctx := context.Background()\n\n\tagent := &nonFlowTestAgent{\n\t\tname: \"test_agent\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tev := &AgentEvent{\n\t\t\t\t\tAction: &AgentAction{\n\t\t\t\t\t\tInterrupted: &InterruptInfo{Data: \"test interrupt\"},\n\t\t\t\t\t},\n\t\t\t\t}\n\t\t\t\tgen.Send(ev)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        agent,\n\t\tToAgentNames: []string{\"next_agent\"},\n\t})\n\n\titer := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\n\tvar events []*AgentEvent\n\tvar sawInterrupt bool\n\tvar sawTransfer bool\n\tfor {\n\t\tev, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, ev)\n\t\tif ev.Action != nil && ev.Action.Interrupted != nil {\n\t\t\tsawInterrupt = true\n\t\t}\n\t\tif ev.Action != nil && ev.Action.TransferToAgent != nil {\n\t\t\tsawTransfer = true\n\t\t}\n\t}\n\n\tassert.True(t, sawInterrupt, \"should see interrupt event\")\n\tassert.False(t, sawTransfer, \"should NOT see transfer when interrupted\")\n\tassert.Len(t, events, 1, \"should have exactly 1 event (interrupt)\")\n}\n\nfunc TestDeterministicTransferNonFlowAgent_Resume(t *testing.T) {\n\tctx := context.Background()\n\n\tvar resumeCalled bool\n\tagent := &nonFlowTestAgent{\n\t\tname: \"test_agent\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tev := EventFromMessage(schema.AssistantMessage(\"from run\", nil), nil, schema.Assistant, \"\")\n\t\t\t\tgen.Send(ev)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tresumeCalled = true\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tev := EventFromMessage(schema.AssistantMessage(\"from resume\", nil), nil, schema.Assistant, \"\")\n\t\t\t\tgen.Send(ev)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        agent,\n\t\tToAgentNames: []string{\"next_agent\"},\n\t})\n\n\tra, ok := wrapped.(ResumableAgent)\n\tassert.True(t, ok, \"wrapped agent should be ResumableAgent\")\n\n\titer := ra.Resume(ctx, &ResumeInfo{WasInterrupted: true})\n\n\tvar events []*AgentEvent\n\tvar sawTransfer bool\n\tfor {\n\t\tev, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, ev)\n\t\tif ev.Action != nil && ev.Action.TransferToAgent != nil {\n\t\t\tsawTransfer = true\n\t\t}\n\t}\n\n\tassert.True(t, resumeCalled, \"resume should have been called on inner agent\")\n\tassert.True(t, sawTransfer, \"should see transfer event after resume completion\")\n}\n\nfunc TestDeterministicTransferFlowAgent_ResumeWithInvalidState(t *testing.T) {\n\tctx := context.Background()\n\n\tinnerAgent := &dtTestAgent{\n\t\tname: \"inner\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tgen.Send(EventFromMessage(schema.AssistantMessage(\"test\", nil), nil, schema.Assistant, \"\"))\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tinnerFlowAgent := toFlowAgent(ctx, innerAgent)\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        innerFlowAgent,\n\t\tToAgentNames: []string{\"next_agent\"},\n\t})\n\n\tra, ok := wrapped.(ResumableAgent)\n\tassert.True(t, ok, \"wrapped flowAgent should be ResumableAgent\")\n\n\titer := ra.Resume(ctx, &ResumeInfo{\n\t\tWasInterrupted: true,\n\t\tInterruptState: nil,\n\t})\n\n\tvar gotError bool\n\tvar errorMsg string\n\tfor {\n\t\tev, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif ev.Err != nil {\n\t\t\tgotError = true\n\t\t\terrorMsg = ev.Err.Error()\n\t\t}\n\t}\n\n\tassert.True(t, gotError, \"should get error for invalid state\")\n\tassert.Contains(t, errorMsg, \"invalid interrupt state\", \"error should mention invalid state\")\n}\n\nfunc TestDeterministicTransferNonResumableAgent(t *testing.T) {\n\tctx := context.Background()\n\n\tagent := &nonResumableTestAgent{\n\t\tname: \"non_resumable\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tev := EventFromMessage(schema.AssistantMessage(\"output\", nil), nil, schema.Assistant, \"\")\n\t\t\t\tgen.Send(ev)\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\twrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{\n\t\tAgent:        agent,\n\t\tToAgentNames: []string{\"next_agent\"},\n\t})\n\n\t_, isResumable := wrapped.(ResumableAgent)\n\tassert.False(t, isResumable, \"wrapped non-resumable agent should NOT be ResumableAgent\")\n\n\tassert.Equal(t, \"non_resumable\", wrapped.Name(ctx), \"Name should delegate to inner agent\")\n\tassert.Equal(t, \"non_resumable description\", wrapped.Description(ctx), \"Description should delegate to inner agent\")\n\n\titer := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\n\tvar sawTransfer bool\n\tfor {\n\t\tev, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif ev.Action != nil && ev.Action.TransferToAgent != nil {\n\t\t\tsawTransfer = true\n\t\t}\n\t}\n\n\tassert.True(t, sawTransfer, \"should see transfer event\")\n}\n"
  },
  {
    "path": "adk/filesystem/backend.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package filesystem provides file system operations.\npackage filesystem\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// FileInfo represents basic file metadata information.\ntype FileInfo struct {\n\t// Path is the path of the file or directory, which can be a filename, relative path, or absolute path.\n\tPath string\n\n\t// IsDir indicates whether the entry is a directory.\n\t// true for directories, false for regular files.\n\tIsDir bool\n\n\t// Size is the file size in bytes.\n\t// For directories, this value may be 0 or platform-dependent.\n\tSize int64\n\n\t// ModifiedAt is the last modification time in ISO 8601 format.\n\t// Example: \"2025-01-15T10:30:00Z\"\n\tModifiedAt string\n}\n\n// GrepMatch represents a single pattern match result.\ntype GrepMatch struct {\n\tContent string\n\n\t// Path is the file path where the match was found.\n\tPath string\n\n\t// Line is the 1-based line number of the match.\n\tLine int\n}\n\n// LsInfoRequest contains parameters for listing file information.\ntype LsInfoRequest struct {\n\t// Path specifies the directory path to list.\n\tPath string\n}\n\n// ReadRequest contains parameters for reading file content.\ntype ReadRequest struct {\n\t// FilePath is the path to the file to be read.\n\tFilePath string\n\n\t// Offset specifies the starting line number (1-based) for reading.\n\t// Line 1 is the first line of the file.\n\t// Use this when the file is too large to read at once.\n\t// Defaults to 1 (start from the first line).\n\t// Values < 1 will be treated as 1.\n\tOffset int\n\n\t// Limit specifies the maximum number of lines to read.\n\t// Use this when the file is too large to read at once.\n\t// Defaults to 2000 if not provided or non-positive (<= 0).\n\tLimit int\n}\n\n// GrepRequest contains parameters for searching file content.\ntype GrepRequest struct {\n\t// ===== Search Parameters =====\n\n\t// Pattern is the search pattern, supports full regular expression syntax.\n\t// Uses ripgrep syntax (not grep). Examples:\n\t//   - \"log.*Error\" matches lines with \"log\" followed by \"Error\"\n\t//   - \"function\\\\s+\\\\w+\" matches \"function\" followed by whitespace and word characters\n\t//   - Literal braces need escaping: \"interface\\\\{\\\\}\" matches \"interface{}\"\n\tPattern string\n\n\t// Path is an optional directory path to limit the search scope.\n\tPath string\n\n\t// ===== File Filtering =====\n\n\t// Glob is an optional pattern to filter the files to be searched.\n\t// It filters by file path, not content. If empty, no files are filtered.\n\t// Supports standard glob wildcards:\n\t//   - `*` matches any characters except path separators.\n\t//   - `**` matches any directories recursively.\n\t//   - `?` matches a single character.\n\t//   - `[abc]` matches one character from the set.\n\tGlob string\n\n\t// FileType is the file type filter, e.g., \"js\", \"py\", \"rust\".\n\t// More efficient than Glob for standard file types.\n\tFileType string\n\n\t// ===== Search Options =====\n\n\t// CaseInsensitive enables case insensitive search.\n\tCaseInsensitive bool\n\n\t// EnableMultiline enables multiline mode where patterns can span lines.\n\t// Default: false (patterns match within single lines only).\n\tEnableMultiline bool\n\n\t// ===== Context Display (Content mode only) =====\n\n\t// AfterLines shows N lines after each match.\n\t// Only applicable when OutputMode is \"content\".\n\t// Values <= 0 are treated as unset.\n\tAfterLines int\n\n\t// BeforeLines shows N lines before each match.\n\t// Only applicable when OutputMode is \"content\".\n\t// Values <= 0 are treated as unset.\n\tBeforeLines int\n}\n\n// GlobInfoRequest contains parameters for glob pattern matching.\ntype GlobInfoRequest struct {\n\t// Pattern is the glob expression used to match file paths.\n\t// It supports standard glob syntax:\n\t//   - `*` matches any characters except path separators.\n\t//   - `**` matches any directories recursively.\n\t//   - `?` matches a single character.\n\t//   - `[abc]` matches one character from the set.\n\tPattern string\n\n\t// Path is the base directory from which to start the search.\n\tPath string\n}\n\n// WriteRequest contains parameters for writing file content.\ntype WriteRequest struct {\n\t// FilePath is the path of the file to write.\n\tFilePath string\n\n\t// Content is the data to be written to the file.\n\tContent string\n}\n\n// EditRequest contains parameters for editing file content.\ntype EditRequest struct {\n\t// FilePath is the path of the file to edit.\n\tFilePath string\n\n\t// OldString is the exact string to be replaced. It must be non-empty and will be matched literally, including whitespace.\n\tOldString string\n\n\t// NewString is the string that will replace OldString.\n\t// It must be different from OldString.\n\t// An empty string can be used to effectively delete OldString.\n\tNewString string\n\n\t// ReplaceAll controls the replacement behavior.\n\t// If true, all occurrences of OldString are replaced.\n\t// If false, the operation fails unless OldString appears exactly once in the file.\n\tReplaceAll bool\n}\n\ntype FileContent struct {\n\tContent string\n}\n\n// Backend is a pluggable, unified file backend protocol interface.\n//\n// All methods use struct-based parameters to allow future extensibility\n// without breaking backward compatibility.\ntype Backend interface {\n\t// LsInfo lists file information under the given path.\n\t//\n\t// Returns:\n\t//   - []FileInfo: List of matching file information\n\t//   - error: Error if the operation fails\n\tLsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error)\n\n\t// Read reads file content with support for line-based offset and limit.\n\t//\n\t// Returns:\n\t//   - string: The file content read\n\t//   - error: Error if file does not exist or read fails\n\tRead(ctx context.Context, req *ReadRequest) (*FileContent, error)\n\n\t// GrepRaw searches for content matching the specified pattern in files.\n\t//\n\t// Returns:\n\t//   - []GrepMatch: List of all matching results\n\t//   - error: Error if the search fails\n\tGrepRaw(ctx context.Context, req *GrepRequest) ([]GrepMatch, error)\n\n\t// GlobInfo returns file information matching the glob pattern.\n\t//\n\t// Returns:\n\t//   - []FileInfo: List of matching file information\n\t//   - error: Error if the pattern is invalid or operation fails\n\tGlobInfo(ctx context.Context, req *GlobInfoRequest) ([]FileInfo, error)\n\n\t// Write creates or updates file content.\n\t//\n\t// Returns:\n\t//   - error: Error if the write operation fails\n\tWrite(ctx context.Context, req *WriteRequest) error\n\n\t// Edit replaces string occurrences in a file.\n\t//\n\t// Returns:\n\t//   - error: Error if file does not exist, OldString is empty, or OldString is not found\n\tEdit(ctx context.Context, req *EditRequest) error\n}\n\n// ExecuteRequest contains parameters for executing a command.\ntype ExecuteRequest struct {\n\tCommand            string // The command to execute\n\tRunInBackendGround bool\n}\n\n// ExecuteResponse contains the response result of command execution.\ntype ExecuteResponse struct {\n\tOutput    string // Command output content\n\tExitCode  *int   // Command exit code\n\tTruncated bool   // Whether the output was truncated\n}\n\ntype Shell interface {\n\tExecute(ctx context.Context, input *ExecuteRequest) (result *ExecuteResponse, err error)\n}\n\ntype StreamingShell interface {\n\tExecuteStreaming(ctx context.Context, input *ExecuteRequest) (result *schema.StreamReader[*ExecuteResponse], err error)\n}\n"
  },
  {
    "path": "adk/filesystem/backend_inmemory.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage filesystem\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"regexp\"\n\t\"strings\"\n\t\"sync\"\n\t\"time\"\n\n\t\"github.com/bmatcuk/doublestar/v4\"\n)\n\ntype fileEntry struct {\n\tcontent    string\n\tmodifiedAt time.Time\n}\n\n// InMemoryBackend is an in-memory implementation of the Backend interface.\n// It stores files in a map and is safe for concurrent use.\ntype InMemoryBackend struct {\n\tmu    sync.RWMutex\n\tfiles map[string]*fileEntry\n}\n\n// NewInMemoryBackend creates a new in-memory backend.\nfunc NewInMemoryBackend() *InMemoryBackend {\n\treturn &InMemoryBackend{\n\t\tfiles: make(map[string]*fileEntry),\n\t}\n}\n\n// LsInfo lists file information under the given path.\nfunc (b *InMemoryBackend) LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error) {\n\tb.mu.RLock()\n\tdefer b.mu.RUnlock()\n\n\t// Normalize path\n\tpath := normalizePath(req.Path)\n\n\tvar result []FileInfo\n\tseen := make(map[string]bool)\n\tdirInfo := make(map[string]*FileInfo)\n\n\tfor filePath, entry := range b.files {\n\t\tnormalizedFilePath := normalizePath(filePath)\n\n\t\t// Check if file is under the given path\n\t\tif path == \"/\" || strings.HasPrefix(normalizedFilePath, path+\"/\") || normalizedFilePath == path {\n\t\t\t// For directory listing, we want to show immediate children\n\t\t\trelativePath := strings.TrimPrefix(normalizedFilePath, path)\n\t\t\trelativePath = strings.TrimPrefix(relativePath, \"/\")\n\n\t\t\tif relativePath == \"\" {\n\t\t\t\t// The path itself is a file\n\t\t\t\tif !seen[normalizedFilePath] {\n\t\t\t\t\tresult = append(result, FileInfo{\n\t\t\t\t\t\tPath:       filepath.Base(normalizedFilePath),\n\t\t\t\t\t\tIsDir:      false,\n\t\t\t\t\t\tSize:       int64(len(entry.content)),\n\t\t\t\t\t\tModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano),\n\t\t\t\t\t})\n\t\t\t\t\tseen[normalizedFilePath] = true\n\t\t\t\t}\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\t// Get the first segment (immediate child)\n\t\t\tparts := strings.SplitN(relativePath, \"/\", 2)\n\t\t\tif len(parts) > 0 {\n\t\t\t\tchildPath := path\n\t\t\t\tif path != \"/\" {\n\t\t\t\t\tchildPath += \"/\"\n\t\t\t\t}\n\t\t\t\tchildPath += parts[0]\n\n\t\t\t\tisDir := len(parts) > 1\n\t\t\t\tif !seen[childPath] {\n\t\t\t\t\tif isDir {\n\t\t\t\t\t\tdirInfo[childPath] = &FileInfo{\n\t\t\t\t\t\t\tPath:       parts[0],\n\t\t\t\t\t\t\tIsDir:      true,\n\t\t\t\t\t\t\tSize:       0,\n\t\t\t\t\t\t\tModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano),\n\t\t\t\t\t\t}\n\t\t\t\t\t} else {\n\t\t\t\t\t\tresult = append(result, FileInfo{\n\t\t\t\t\t\t\tPath:       parts[0],\n\t\t\t\t\t\t\tIsDir:      false,\n\t\t\t\t\t\t\tSize:       int64(len(entry.content)),\n\t\t\t\t\t\t\tModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano),\n\t\t\t\t\t\t})\n\t\t\t\t\t}\n\t\t\t\t\tseen[childPath] = true\n\t\t\t\t} else if isDir {\n\t\t\t\t\tif info, ok := dirInfo[childPath]; ok {\n\t\t\t\t\t\tif entry.modifiedAt.After(mustParseTime(info.ModifiedAt)) {\n\t\t\t\t\t\t\tinfo.ModifiedAt = entry.modifiedAt.Format(time.RFC3339Nano)\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tfor _, info := range dirInfo {\n\t\tresult = append(result, *info)\n\t}\n\n\treturn result, nil\n}\n\nfunc mustParseTime(s string) time.Time {\n\tt, _ := time.Parse(time.RFC3339Nano, s)\n\treturn t\n}\n\n// Read reads file content with offset and limit.\nfunc (b *InMemoryBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) {\n\tb.mu.RLock()\n\tdefer b.mu.RUnlock()\n\n\tfilePath := normalizePath(req.FilePath)\n\n\tentry, exists := b.files[filePath]\n\tif !exists {\n\t\treturn nil, fmt.Errorf(\"file not found: %s\", filePath)\n\t}\n\n\t// Convert 1-based offset to 0-based index; values < 1 default to line 1\n\toffset := req.Offset - 1\n\tif offset < 0 {\n\t\toffset = 0\n\t}\n\tlimit := req.Limit\n\tif limit <= 0 {\n\t\tlimit = 2000\n\t}\n\n\tcontent := entry.content\n\n\t// Fast path: no offset, content fits within limit — return as-is\n\tif offset == 0 {\n\t\tlineCount := strings.Count(content, \"\\n\") + 1\n\t\tif lineCount <= limit {\n\t\t\treturn &FileContent{Content: content}, nil\n\t\t}\n\t}\n\n\t// Skip `offset` lines by scanning for newlines directly\n\tstart := 0\n\tfor i := 0; i < offset; i++ {\n\t\tidx := strings.IndexByte(content[start:], '\\n')\n\t\tif idx == -1 {\n\t\t\t// offset exceeds total lines\n\t\t\treturn &FileContent{}, nil\n\t\t}\n\t\tstart += idx + 1\n\t}\n\n\t// Find the end position after `limit` lines\n\tend := start\n\tfor i := 0; i < limit; i++ {\n\t\tidx := strings.IndexByte(content[end:], '\\n')\n\t\tif idx == -1 {\n\t\t\t// Reached the end of content\n\t\t\treturn &FileContent{Content: content[start:]}, nil\n\t\t}\n\t\tend += idx + 1\n\t}\n\n\t// Trim the trailing newline from the last included line\n\treturn &FileContent{Content: content[start : end-1]}, nil\n}\n\n// GrepRaw returns matches for the given pattern.\nfunc (b *InMemoryBackend) GrepRaw(ctx context.Context, req *GrepRequest) ([]GrepMatch, error) {\n\tb.mu.RLock()\n\tdefer b.mu.RUnlock()\n\n\tif req.Pattern == \"\" {\n\t\treturn nil, fmt.Errorf(\"pattern cannot be empty\")\n\t}\n\n\tre, err := b.compilePattern(req)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tsearchPath := \"/\"\n\tif req.Path != \"\" {\n\t\tsearchPath = normalizePath(req.Path)\n\t}\n\n\tfilteredFiles, err := b.filterFiles(searchPath, req)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tif len(filteredFiles) == 0 {\n\t\treturn []GrepMatch{}, nil\n\t}\n\n\tif len(filteredFiles) == 1 {\n\t\tcollector := newGrepCollector()\n\t\tentry := b.files[filteredFiles[0]]\n\t\tcollector.processFile(filteredFiles[0], entry.content, re, req)\n\t\treturn collector.buildResults(b, req)\n\t}\n\n\tmatches, err := b.grepFilesInParallel(filteredFiles, re, req)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tif req.BeforeLines > 0 || req.AfterLines > 0 {\n\t\tmatches = b.applyContext(matches, req)\n\t}\n\n\treturn matches, nil\n}\n\nfunc (b *InMemoryBackend) grepFilesInParallel(filteredFiles []string, re *regexp.Regexp, req *GrepRequest) ([]GrepMatch, error) {\n\tnumWorkers := len(filteredFiles)\n\tif numWorkers > 10 {\n\t\tnumWorkers = 10\n\t}\n\n\ttype fileTask struct {\n\t\tpath    string\n\t\tcontent string\n\t}\n\n\ttasks := make(chan fileTask, len(filteredFiles))\n\tresults := make(chan []GrepMatch, len(filteredFiles))\n\terrChan := make(chan error, numWorkers)\n\n\tvar wg sync.WaitGroup\n\tfor i := 0; i < numWorkers; i++ {\n\t\twg.Add(1)\n\t\tgo func(workerID int) {\n\t\t\tdefer wg.Done()\n\t\t\tdefer func() {\n\t\t\t\tif r := recover(); r != nil {\n\t\t\t\t\terrChan <- fmt.Errorf(\"worker %d panic: %v\", workerID, r)\n\t\t\t\t}\n\t\t\t}()\n\n\t\t\tcollector := newGrepCollector()\n\t\t\tfor task := range tasks {\n\t\t\t\tfileMatches := collector.findMatches(task.path, task.content, re, req)\n\t\t\t\tif len(fileMatches) > 0 {\n\t\t\t\t\tresults <- fileMatches\n\t\t\t\t}\n\t\t\t}\n\t\t}(i)\n\t}\n\n\tfor _, filePath := range filteredFiles {\n\t\tentry := b.files[filePath]\n\t\ttasks <- fileTask{\n\t\t\tpath:    filePath,\n\t\t\tcontent: entry.content,\n\t\t}\n\t}\n\tclose(tasks)\n\n\tgo func() {\n\t\twg.Wait()\n\t\tclose(results)\n\t\tclose(errChan)\n\t}()\n\n\tvar allMatches []GrepMatch\n\tvar errs []error\n\n\tfor {\n\t\tselect {\n\t\tcase matches, ok := <-results:\n\t\t\tif !ok {\n\t\t\t\tresults = nil\n\t\t\t} else {\n\t\t\t\tallMatches = append(allMatches, matches...)\n\t\t\t}\n\t\tcase err, ok := <-errChan:\n\t\t\tif !ok {\n\t\t\t\terrChan = nil\n\t\t\t} else if err != nil {\n\t\t\t\terrs = append(errs, err)\n\t\t\t}\n\t\t}\n\n\t\tif results == nil && errChan == nil {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tif len(errs) > 0 {\n\t\treturn nil, fmt.Errorf(\"grep failed with %d error(s): %v\", len(errs), errs[0])\n\t}\n\n\treturn allMatches, nil\n}\n\nfunc (b *InMemoryBackend) compilePattern(req *GrepRequest) (*regexp.Regexp, error) {\n\tpattern := req.Pattern\n\tif req.CaseInsensitive {\n\t\tpattern = \"(?i)\" + pattern\n\t}\n\tre, err := regexp.Compile(pattern)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"invalid regex pattern: %w\", err)\n\t}\n\treturn re, nil\n}\n\nfunc (b *InMemoryBackend) filterFiles(searchPath string, req *GrepRequest) ([]string, error) {\n\tvar candidateFiles []string\n\n\tfor filePath := range b.files {\n\t\tnormalizedFilePath := normalizePath(filePath)\n\n\t\tif searchPath != \"/\" && !strings.HasPrefix(normalizedFilePath, searchPath+\"/\") && normalizedFilePath != searchPath {\n\t\t\tcontinue\n\t\t}\n\n\t\tcandidateFiles = append(candidateFiles, normalizedFilePath)\n\t}\n\n\tif req.Glob != \"\" {\n\t\tfiltered, err := b.filterByGlob(candidateFiles, searchPath, req.Glob)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tcandidateFiles = filtered\n\t}\n\n\tif req.FileType != \"\" {\n\t\tcandidateFiles = b.filterByFileType(candidateFiles, req.FileType)\n\t}\n\n\treturn candidateFiles, nil\n}\n\nfunc (b *InMemoryBackend) filterByGlob(files []string, searchPath string, globPattern string) ([]string, error) {\n\tvar result []string\n\n\tfor _, filePath := range files {\n\t\tvar matchPath string\n\t\tif strings.Contains(globPattern, \"/\") || strings.Contains(globPattern, \"**\") {\n\t\t\tif searchPath == \"/\" {\n\t\t\t\tmatchPath = strings.TrimPrefix(filePath, \"/\")\n\t\t\t} else {\n\t\t\t\tmatchPath = strings.TrimPrefix(filePath, searchPath+\"/\")\n\t\t\t}\n\t\t} else {\n\t\t\tmatchPath = filepath.Base(filePath)\n\t\t}\n\n\t\tmatched, err := doublestar.Match(globPattern, matchPath)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"invalid glob pattern: %w\", err)\n\t\t}\n\t\tif matched {\n\t\t\tresult = append(result, filePath)\n\t\t}\n\t}\n\n\treturn result, nil\n}\n\nfunc (b *InMemoryBackend) filterByFileType(files []string, fileType string) []string {\n\tvar result []string\n\n\tfor _, filePath := range files {\n\t\text := strings.TrimPrefix(filepath.Ext(filePath), \".\")\n\t\tif matchFileType(ext, fileType) {\n\t\t\tresult = append(result, filePath)\n\t\t}\n\t}\n\n\treturn result\n}\n\n// matchFileType checks if the file extension matches the given file type.\nfunc matchFileType(ext, fileType string) bool {\n\ttypeMap := map[string][]string{\n\t\t\"ada\":          {\"adb\", \"ads\"},\n\t\t\"agda\":         {\"agda\", \"lagda\"},\n\t\t\"aidl\":         {\"aidl\"},\n\t\t\"amake\":        {\"bp\", \"mk\"},\n\t\t\"asciidoc\":     {\"adoc\", \"asc\", \"asciidoc\"},\n\t\t\"asm\":          {\"S\", \"asm\", \"s\"},\n\t\t\"asp\":          {\"ascx\", \"asp\", \"aspx\"},\n\t\t\"ats\":          {\"ats\", \"dats\", \"hats\", \"sats\"},\n\t\t\"avro\":         {\"avdl\", \"avpr\", \"avsc\"},\n\t\t\"awk\":          {\"awk\"},\n\t\t\"bat\":          {\"bat\"},\n\t\t\"bazel\":        {\"BUILD\", \"bazel\", \"bzl\"},\n\t\t\"bitbake\":      {\"bb\", \"bbappend\", \"bbclass\", \"conf\", \"inc\"},\n\t\t\"c\":            {\"c\", \"h\", \"H\", \"cats\"},\n\t\t\"cabal\":        {\"cabal\"},\n\t\t\"cbor\":         {\"cbor\"},\n\t\t\"ceylon\":       {\"ceylon\"},\n\t\t\"clojure\":      {\"clj\", \"cljc\", \"cljs\", \"cljx\"},\n\t\t\"cmake\":        {\"cmake\"},\n\t\t\"coffeescript\": {\"coffee\"},\n\t\t\"config\":       {\"cfg\", \"conf\", \"config\", \"ini\"},\n\t\t\"coq\":          {\"v\"},\n\t\t\"cpp\":          {\"C\", \"cc\", \"cpp\", \"cxx\", \"c++\", \"h\", \"hh\", \"hpp\", \"hxx\", \"h++\", \"inl\"},\n\t\t\"crystal\":      {\"cr\", \"ecr\"},\n\t\t\"cs\":           {\"cs\"},\n\t\t\"csharp\":       {\"cs\"},\n\t\t\"cshtml\":       {\"cshtml\"},\n\t\t\"css\":          {\"css\", \"scss\", \"sass\", \"less\"},\n\t\t\"csv\":          {\"csv\"},\n\t\t\"cuda\":         {\"cu\", \"cuh\"},\n\t\t\"cython\":       {\"pxd\", \"pxi\", \"pyx\"},\n\t\t\"d\":            {\"d\"},\n\t\t\"dart\":         {\"dart\"},\n\t\t\"devicetree\":   {\"dts\", \"dtsi\"},\n\t\t\"dhall\":        {\"dhall\"},\n\t\t\"diff\":         {\"diff\", \"patch\"},\n\t\t\"docker\":       {\"dockerfile\"},\n\t\t\"go\":           {\"go\"},\n\t\t\"groovy\":       {\"gradle\", \"groovy\"},\n\t\t\"haskell\":      {\"c2hs\", \"cpphs\", \"hs\", \"hsc\", \"lhs\"},\n\t\t\"html\":         {\"ejs\", \"htm\", \"html\"},\n\t\t\"java\":         {\"java\", \"jsp\", \"jspx\", \"properties\"},\n\t\t\"js\":           {\"cjs\", \"js\", \"jsx\", \"mjs\", \"vue\"},\n\t\t\"json\":         {\"json\", \"sarif\"},\n\t\t\"jsonl\":        {\"jsonl\"},\n\t\t\"julia\":        {\"jl\"},\n\t\t\"jupyter\":      {\"ipynb\", \"jpynb\"},\n\t\t\"kotlin\":       {\"kt\", \"kts\"},\n\t\t\"less\":         {\"less\"},\n\t\t\"lua\":          {\"lua\"},\n\t\t\"make\":         {\"mak\", \"mk\"},\n\t\t\"markdown\":     {\"markdown\", \"md\", \"mdown\", \"mdwn\", \"mdx\", \"mkd\", \"mkdn\"},\n\t\t\"md\":           {\"markdown\", \"md\", \"mdown\", \"mdwn\", \"mdx\", \"mkd\", \"mkdn\"},\n\t\t\"matlab\":       {\"m\"},\n\t\t\"ocaml\":        {\"ml\", \"mli\", \"mll\", \"mly\"},\n\t\t\"perl\":         {\"PL\", \"perl\", \"pl\", \"plh\", \"plx\", \"pm\", \"t\"},\n\t\t\"php\":          {\"php\", \"php3\", \"php4\", \"php5\", \"php7\", \"php8\", \"pht\", \"phtml\"},\n\t\t\"python\":       {\"py\", \"pyi\"},\n\t\t\"py\":           {\"py\", \"pyi\"},\n\t\t\"ruby\":         {\"gemspec\", \"rb\", \"rbw\"},\n\t\t\"rust\":         {\"rs\"},\n\t\t\"sass\":         {\"sass\", \"scss\"},\n\t\t\"scala\":        {\"sbt\", \"scala\"},\n\t\t\"sh\":           {\"bash\", \"sh\", \"zsh\"},\n\t\t\"sql\":          {\"psql\", \"sql\"},\n\t\t\"swift\":        {\"swift\"},\n\t\t\"toml\":         {\"toml\"},\n\t\t\"ts\":           {\"cts\", \"mts\", \"ts\", \"tsx\"},\n\t\t\"typescript\":   {\"cts\", \"mts\", \"ts\", \"tsx\"},\n\t\t\"txt\":          {\"txt\"},\n\t\t\"vue\":          {\"vue\"},\n\t\t\"xml\":          {\"dtd\", \"xml\", \"xsd\", \"xsl\", \"xslt\"},\n\t\t\"yaml\":         {\"yaml\", \"yml\"},\n\t\t\"zig\":          {\"zig\"},\n\t}\n\n\tif exts, ok := typeMap[fileType]; ok {\n\t\tfor _, e := range exts {\n\t\t\tif ext == e {\n\t\t\t\treturn true\n\t\t\t}\n\t\t}\n\t}\n\treturn ext == fileType\n}\n\n// applyContext adds context lines around matches.\nfunc (b *InMemoryBackend) applyContext(matches []GrepMatch, req *GrepRequest) []GrepMatch {\n\tif len(matches) == 0 {\n\t\treturn matches\n\t}\n\n\tbeforeLines := 0\n\tafterLines := 0\n\n\tif req.BeforeLines > 0 {\n\t\tbeforeLines = req.BeforeLines\n\t}\n\tif req.AfterLines > 0 {\n\t\tafterLines = req.AfterLines\n\t}\n\n\tif beforeLines <= 0 && afterLines <= 0 {\n\t\treturn matches\n\t}\n\n\t// Group matches by file path for efficient processing\n\tmatchesByFile := make(map[string][]GrepMatch)\n\tfileOrder := make([]string, 0)\n\tseenFiles := make(map[string]bool)\n\n\tfor _, match := range matches {\n\t\tif !seenFiles[match.Path] {\n\t\t\tfileOrder = append(fileOrder, match.Path)\n\t\t\tseenFiles[match.Path] = true\n\t\t}\n\t\tmatchesByFile[match.Path] = append(matchesByFile[match.Path], match)\n\t}\n\n\tvar result []GrepMatch\n\n\t// Process each file once\n\tfor _, filePath := range fileOrder {\n\t\tfileMatches := matchesByFile[filePath]\n\n\t\t// Get file content once per file\n\t\tb.mu.RLock()\n\t\tentry, exists := b.files[filePath]\n\t\tb.mu.RUnlock()\n\n\t\tif !exists {\n\t\t\t// If file doesn't exist, keep original matches\n\t\t\tresult = append(result, fileMatches...)\n\t\t\tcontinue\n\t\t}\n\n\t\tlines := strings.Split(entry.content, \"\\n\")\n\t\tprocessedLines := make(map[int]bool)\n\n\t\t// Process all matches for this file\n\t\tfor _, match := range fileMatches {\n\t\t\tstartLine := match.Line - beforeLines\n\t\t\tif startLine < 1 {\n\t\t\t\tstartLine = 1\n\t\t\t}\n\n\t\t\tendLine := match.Line + afterLines\n\t\t\tif endLine > len(lines) {\n\t\t\t\tendLine = len(lines)\n\t\t\t}\n\n\t\t\tfor lineNum := startLine; lineNum <= endLine; lineNum++ {\n\t\t\t\tif !processedLines[lineNum] {\n\t\t\t\t\tprocessedLines[lineNum] = true\n\t\t\t\t\tresult = append(result, GrepMatch{\n\t\t\t\t\t\tPath:    filePath,\n\t\t\t\t\t\tLine:    lineNum,\n\t\t\t\t\t\tContent: lines[lineNum-1],\n\t\t\t\t\t})\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\treturn result\n}\n\n// GlobInfo returns file info entries matching the glob pattern.\nfunc (b *InMemoryBackend) GlobInfo(ctx context.Context, req *GlobInfoRequest) ([]FileInfo, error) {\n\tb.mu.RLock()\n\tdefer b.mu.RUnlock()\n\n\tbasePath := normalizePath(req.Path)\n\tisAbsolutePattern := strings.HasPrefix(req.Pattern, \"/\")\n\n\tvar result []FileInfo\n\n\tfor filePath, entry := range b.files {\n\t\tnormalizedFilePath := normalizePath(filePath)\n\n\t\tvar matchPath string\n\t\tvar resultPath string\n\n\t\tif isAbsolutePattern {\n\t\t\tmatchPath = normalizedFilePath\n\t\t\tresultPath = normalizedFilePath\n\t\t} else {\n\t\t\tif basePath != \"/\" && !strings.HasPrefix(normalizedFilePath, basePath+\"/\") && normalizedFilePath != basePath {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tif basePath == \"/\" {\n\t\t\t\tmatchPath = strings.TrimPrefix(normalizedFilePath, \"/\")\n\t\t\t} else {\n\t\t\t\tmatchPath = strings.TrimPrefix(normalizedFilePath, basePath+\"/\")\n\t\t\t}\n\t\t\tresultPath = matchPath\n\t\t}\n\n\t\tmatched, err := doublestar.Match(req.Pattern, matchPath)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"invalid glob pattern: %w\", err)\n\t\t}\n\n\t\tif matched {\n\t\t\tresult = append(result, FileInfo{\n\t\t\t\tPath:       resultPath,\n\t\t\t\tIsDir:      false,\n\t\t\t\tSize:       int64(len(entry.content)),\n\t\t\t\tModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano),\n\t\t\t})\n\t\t}\n\t}\n\n\treturn result, nil\n}\n\n// Write creates or overwrites file content.\nfunc (b *InMemoryBackend) Write(ctx context.Context, req *WriteRequest) error {\n\tb.mu.Lock()\n\tdefer b.mu.Unlock()\n\n\tfilePath := normalizePath(req.FilePath)\n\tb.files[filePath] = &fileEntry{\n\t\tcontent:    req.Content,\n\t\tmodifiedAt: time.Now(),\n\t}\n\n\treturn nil\n}\n\n// Edit replaces string occurrences in a file.\nfunc (b *InMemoryBackend) Edit(ctx context.Context, req *EditRequest) error {\n\tb.mu.Lock()\n\tdefer b.mu.Unlock()\n\n\tfilePath := normalizePath(req.FilePath)\n\n\tentry, exists := b.files[filePath]\n\tif !exists {\n\t\treturn fmt.Errorf(\"file not found: %s\", filePath)\n\t}\n\n\tif req.OldString == \"\" {\n\t\treturn fmt.Errorf(\"oldString must be non-empty\")\n\t}\n\n\tcontent := entry.content\n\tif !strings.Contains(content, req.OldString) {\n\t\treturn fmt.Errorf(\"oldString not found in file: %s\", filePath)\n\t}\n\n\tif !req.ReplaceAll {\n\t\tfirstIndex := strings.Index(content, req.OldString)\n\t\tif firstIndex != -1 {\n\t\t\t// Check if there's another occurrence after the first one\n\t\t\tif strings.Contains(content[firstIndex+len(req.OldString):], req.OldString) {\n\t\t\t\treturn fmt.Errorf(\"multiple occurrences of oldString found in file %s, but ReplaceAll is false\", filePath)\n\t\t\t}\n\t\t}\n\t}\n\n\tvar newContent string\n\tif req.ReplaceAll {\n\t\tnewContent = strings.ReplaceAll(content, req.OldString, req.NewString)\n\t} else {\n\t\tnewContent = strings.Replace(content, req.OldString, req.NewString, 1)\n\t}\n\n\tb.files[filePath] = &fileEntry{\n\t\tcontent:    newContent,\n\t\tmodifiedAt: time.Now(),\n\t}\n\n\treturn nil\n}\n\n// normalizePath normalizes a file path by ensuring it starts with \"/\" and removing trailing slashes.\nfunc normalizePath(path string) string {\n\tif path == \"\" {\n\t\treturn \"/\"\n\t}\n\n\t// Ensure path starts with \"/\"\n\tif !strings.HasPrefix(path, \"/\") {\n\t\tpath = \"/\" + path\n\t}\n\n\treturn filepath.Clean(path)\n}\n\ntype grepCollector struct {\n\tallMatches []GrepMatch\n}\n\nfunc newGrepCollector() *grepCollector {\n\treturn &grepCollector{\n\t\tallMatches: []GrepMatch{},\n\t}\n}\n\nfunc (c *grepCollector) processFile(filePath, content string, re *regexp.Regexp, req *GrepRequest) {\n\tfileMatches := c.findMatches(filePath, content, re, req)\n\tif len(fileMatches) > 0 {\n\t\tc.allMatches = append(c.allMatches, fileMatches...)\n\t}\n}\n\nfunc (c *grepCollector) findMatches(filePath, content string, re *regexp.Regexp, req *GrepRequest) []GrepMatch {\n\tif req.EnableMultiline {\n\t\treturn c.findMultilineMatches(filePath, content, re)\n\t}\n\treturn c.findSingleLineMatches(filePath, content, re)\n}\n\nfunc (c *grepCollector) findMultilineMatches(filePath, content string, re *regexp.Regexp) []GrepMatch {\n\tvar fileMatches []GrepMatch\n\tmatches := re.FindAllStringIndex(content, -1)\n\tlines := strings.Split(content, \"\\n\")\n\n\tfor _, match := range matches {\n\t\tmatchStart := match[0]\n\t\tmatchEnd := match[1]\n\t\tstartLineNum := 1 + strings.Count(content[:matchStart], \"\\n\")\n\t\tendLineNum := 1 + strings.Count(content[:matchEnd], \"\\n\")\n\n\t\tfor lineNum := startLineNum; lineNum <= endLineNum && lineNum <= len(lines); lineNum++ {\n\t\t\tfileMatches = append(fileMatches, GrepMatch{\n\t\t\t\tPath:    filePath,\n\t\t\t\tLine:    lineNum,\n\t\t\t\tContent: lines[lineNum-1],\n\t\t\t})\n\t\t}\n\t}\n\treturn fileMatches\n}\n\nfunc (c *grepCollector) findSingleLineMatches(filePath, content string, re *regexp.Regexp) []GrepMatch {\n\tvar fileMatches []GrepMatch\n\tlines := strings.Split(content, \"\\n\")\n\tfor lineNum, line := range lines {\n\t\tif re.MatchString(line) {\n\t\t\tfileMatches = append(fileMatches, GrepMatch{\n\t\t\t\tPath:    filePath,\n\t\t\t\tLine:    lineNum + 1,\n\t\t\t\tContent: line,\n\t\t\t})\n\t\t}\n\t}\n\treturn fileMatches\n}\n\nfunc (c *grepCollector) buildResults(b *InMemoryBackend, req *GrepRequest) ([]GrepMatch, error) {\n\treturn c.buildContentResult(b, req), nil\n}\n\nfunc (c *grepCollector) buildContentResult(b *InMemoryBackend, req *GrepRequest) []GrepMatch {\n\tresults := c.allMatches\n\tif req.BeforeLines > 0 || req.AfterLines > 0 {\n\t\tresults = b.applyContext(c.allMatches, req)\n\t}\n\treturn results\n}\n"
  },
  {
    "path": "adk/filesystem/backend_inmemory_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage filesystem\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"testing\"\n\t\"time\"\n)\n\nfunc TestInMemoryBackend_WriteAndRead(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\t// Test Write\n\terr := backend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/test.txt\",\n\t\tContent:  \"line1\\nline2\\nline3\\nline4\\nline5\",\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"Write failed: %v\", err)\n\t}\n\n\t// Test Read - full content\n\tcontent, err := backend.Read(ctx, &ReadRequest{\n\t\tFilePath: \"/test.txt\",\n\t\tLimit:    100,\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"Read failed: %v\", err)\n\t}\n\texpected := \"line1\\nline2\\nline3\\nline4\\nline5\"\n\tif content.Content != expected {\n\t\tt.Errorf(\"Read content mismatch. Expected: %q, Got: %q\", expected, content.Content)\n\t}\n\n\t// Test Read - with offset and limit\n\tcontent, err = backend.Read(ctx, &ReadRequest{\n\t\tFilePath: \"/test.txt\",\n\t\tOffset:   1,\n\t\tLimit:    2,\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"Read with offset failed: %v\", err)\n\t}\n\texpected = \"line1\\nline2\"\n\tif content.Content != expected {\n\t\tt.Errorf(\"Read with offset content mismatch. Expected: %q, Got: %q\", expected, content.Content)\n\t}\n\n\t// Test Read - non-existent file\n\t_, err = backend.Read(ctx, &ReadRequest{\n\t\tFilePath: \"/nonexistent.txt\",\n\t\tLimit:    10,\n\t})\n\tif err == nil {\n\t\tt.Error(\"Expected error for non-existent file, got nil\")\n\t}\n}\n\nfunc TestInMemoryBackend_LsInfo(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\t// Create some files\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/file1.txt\",\n\t\tContent:  \"content1\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/file2.txt\",\n\t\tContent:  \"content2\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/dir1/file3.txt\",\n\t\tContent:  \"content3\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/dir1/subdir/file4.txt\",\n\t\tContent:  \"content4\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/dir2/file5.txt\",\n\t\tContent:  \"content5\",\n\t})\n\n\t// Test LsInfo - root\n\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\tif err != nil {\n\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t}\n\tif len(infos) != 4 { // file1.txt, file2.txt, dir1, dir2\n\t\tt.Errorf(\"Expected 4 items in root, got %d\", len(infos))\n\t}\n\n\t// Test LsInfo - specific directory\n\tinfos, err = backend.LsInfo(ctx, &LsInfoRequest{Path: \"/dir1\"})\n\tif err != nil {\n\t\tt.Fatalf(\"LsInfo for /dir1 failed: %v\", err)\n\t}\n\tif len(infos) != 2 { // file3.txt, subdir\n\t\tt.Errorf(\"Expected 2 items in /dir1, got %d\", len(infos))\n\t}\n}\n\nfunc TestInMemoryBackend_Edit(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\t// Create a file\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/edit.txt\",\n\t\tContent:  \"hello world\\nhello again\\nhello world\",\n\t})\n\n\t// Test Edit - report error if old string occurs\n\terr := backend.Edit(ctx, &EditRequest{\n\t\tFilePath:   \"/edit.txt\",\n\t\tOldString:  \"hello\",\n\t\tNewString:  \"hi\",\n\t\tReplaceAll: false,\n\t})\n\tif err == nil {\n\t\tt.Fatal(\"should have failed\")\n\t}\n\n\t// Test Edit - replace all occurrences\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/edit2.txt\",\n\t\tContent:  \"hello world\\nhello again\\nhello world\",\n\t})\n\terr = backend.Edit(ctx, &EditRequest{\n\t\tFilePath:   \"/edit2.txt\",\n\t\tOldString:  \"hello\",\n\t\tNewString:  \"hi\",\n\t\tReplaceAll: true,\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"Edit (replace all) failed: %v\", err)\n\t}\n\n\tcontent, _ := backend.Read(ctx, &ReadRequest{\n\t\tFilePath: \"/edit2.txt\",\n\t\tLimit:    100,\n\t})\n\texpected := \"hi world\\nhi again\\nhi world\"\n\tif content.Content != expected {\n\t\tt.Errorf(\"Edit (replace all) content mismatch. Expected: %q, Got: %q\", expected, content.Content)\n\t}\n\n\t// Test Edit - non-existent file\n\terr = backend.Edit(ctx, &EditRequest{\n\t\tFilePath:   \"/nonexistent.txt\",\n\t\tOldString:  \"old\",\n\t\tNewString:  \"new\",\n\t\tReplaceAll: false,\n\t})\n\tif err == nil {\n\t\tt.Error(\"Expected error for non-existent file, got nil\")\n\t}\n\n\t// Test Edit - empty oldString\n\terr = backend.Edit(ctx, &EditRequest{\n\t\tFilePath:   \"/edit.txt\",\n\t\tOldString:  \"\",\n\t\tNewString:  \"new\",\n\t\tReplaceAll: false,\n\t})\n\tif err == nil {\n\t\tt.Error(\"Expected error for empty oldString, got nil\")\n\t}\n}\n\nfunc TestInMemoryBackend_LsInfo_PathIsFilename(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/file1.txt\",\n\t\tContent:  \"content1\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/file2.txt\",\n\t\tContent:  \"content2\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/dir1/file3.txt\",\n\t\tContent:  \"content3\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/dir1/subdir/file4.txt\",\n\t\tContent:  \"content4\",\n\t})\n\n\tt.Run(\"RootDirectory\", func(t *testing.T) {\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif strings.Contains(info.Path, \"/\") {\n\t\t\t\tt.Errorf(\"Path should be filename only, got: %s\", info.Path)\n\t\t\t}\n\t\t\tif info.IsDir {\n\t\t\t\tif info.Path != \"dir1\" {\n\t\t\t\t\tt.Errorf(\"Expected directory name 'dir1', got: %s\", info.Path)\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tif info.Path != \"file1.txt\" && info.Path != \"file2.txt\" {\n\t\t\t\t\tt.Errorf(\"Expected filename 'file1.txt' or 'file2.txt', got: %s\", info.Path)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"Subdirectory\", func(t *testing.T) {\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/dir1\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif strings.Contains(info.Path, \"/\") {\n\t\t\t\tt.Errorf(\"Path should be filename only, got: %s\", info.Path)\n\t\t\t}\n\t\t\tif info.IsDir {\n\t\t\t\tif info.Path != \"subdir\" {\n\t\t\t\t\tt.Errorf(\"Expected directory name 'subdir', got: %s\", info.Path)\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tif info.Path != \"file3.txt\" {\n\t\t\t\t\tt.Errorf(\"Expected filename 'file3.txt', got: %s\", info.Path)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"NestedSubdirectory\", func(t *testing.T) {\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/dir1/subdir\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos))\n\t\t}\n\n\t\tinfo := infos[0]\n\t\tif info.Path != \"file4.txt\" {\n\t\t\tt.Errorf(\"Expected filename 'file4.txt', got: %s\", info.Path)\n\t\t}\n\t\tif strings.Contains(info.Path, \"/\") {\n\t\t\tt.Errorf(\"Path should be filename only, got: %s\", info.Path)\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_GlobInfo(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\t// Create some files\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/file1.txt\",\n\t\tContent:  \"content1\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/file2.py\",\n\t\tContent:  \"content2\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/dir1/file3.txt\",\n\t\tContent:  \"content3\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/dir1/file4.py\",\n\t\tContent:  \"content4\",\n\t})\n\n\t// Test GlobInfo - match .txt files in root only\n\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\tPattern: \"*.txt\",\n\t\tPath:    \"/\",\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t}\n\tif len(infos) != 1 { // only file1.txt in root\n\t\tt.Errorf(\"Expected 1 .txt file in root, got %d\", len(infos))\n\t}\n\tif infos[0].Path != \"file1.txt\" {\n\t\tt.Errorf(\"Expected relative path 'file1.txt', got %s\", infos[0].Path)\n\t}\n\n\t// Test GlobInfo - match all .py files in dir1\n\tinfos, err = backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\tPattern: \"*.py\",\n\t\tPath:    \"/dir1\",\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"GlobInfo for /dir1 failed: %v\", err)\n\t}\n\tif len(infos) != 1 { // file4.py\n\t\tt.Errorf(\"Expected 1 .py file in /dir1, got %d\", len(infos))\n\t}\n\tif infos[0].Path != \"file4.py\" {\n\t\tt.Errorf(\"Expected relative path 'file4.py', got %s\", infos[0].Path)\n\t}\n}\n\nfunc TestInMemoryBackend_GlobInfo_RelativePath(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/Users/bytedance/Desktop/github/eino/file1.go\",\n\t\tContent:  \"content1\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/Users/bytedance/Desktop/github/openai-go/paginationmanual_test.go\",\n\t\tContent:  \"content2\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/Users/bytedance/Desktop/github/openai-go/paginationauto_test.go\",\n\t\tContent:  \"content3\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/Users/bytedance/Desktop/other/test.go\",\n\t\tContent:  \"content4\",\n\t})\n\n\tt.Run(\"GlobFromRootWithPattern\", func(t *testing.T) {\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"**/*.go\",\n\t\t\tPath:    \"/Users/bytedance/Desktop/github\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 3 {\n\t\t\tt.Fatalf(\"Expected 3 .go files, got %d\", len(infos))\n\t\t}\n\n\t\texpectedPaths := map[string]bool{\n\t\t\t\"eino/file1.go\":                      false,\n\t\t\t\"openai-go/paginationmanual_test.go\": false,\n\t\t\t\"openai-go/paginationauto_test.go\":   false,\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif _, exists := expectedPaths[info.Path]; exists {\n\t\t\t\texpectedPaths[info.Path] = true\n\t\t\t} else {\n\t\t\t\tt.Errorf(\"Unexpected path: %s\", info.Path)\n\t\t\t}\n\t\t}\n\n\t\tfor path, found := range expectedPaths {\n\t\t\tif !found {\n\t\t\t\tt.Errorf(\"Expected path not found: %s\", path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"GlobFromSubdirectory\", func(t *testing.T) {\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"*.go\",\n\t\t\tPath:    \"/Users/bytedance/Desktop/github/openai-go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 2 {\n\t\t\tt.Fatalf(\"Expected 2 .go files, got %d\", len(infos))\n\t\t}\n\n\t\texpectedPaths := map[string]bool{\n\t\t\t\"paginationmanual_test.go\": false,\n\t\t\t\"paginationauto_test.go\":   false,\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif _, exists := expectedPaths[info.Path]; exists {\n\t\t\t\texpectedPaths[info.Path] = true\n\t\t\t} else {\n\t\t\t\tt.Errorf(\"Unexpected path: %s\", info.Path)\n\t\t\t}\n\t\t}\n\n\t\tfor path, found := range expectedPaths {\n\t\t\tif !found {\n\t\t\t\tt.Errorf(\"Expected path not found: %s\", path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"GlobFromRootWithAbsolutePattern\", func(t *testing.T) {\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"/Users/bytedance/Desktop/github/**/*.go\",\n\t\t\tPath:    \"/\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\texpected := map[string]bool{\n\t\t\t\"/Users/bytedance/Desktop/github/eino/file1.go\":                      false,\n\t\t\t\"/Users/bytedance/Desktop/github/openai-go/paginationmanual_test.go\": false,\n\t\t\t\"/Users/bytedance/Desktop/github/openai-go/paginationauto_test.go\":   false,\n\t\t}\n\t\tfor _, info := range infos {\n\t\t\tif _, ok := expected[info.Path]; ok {\n\t\t\t\texpected[info.Path] = true\n\t\t\t}\n\t\t}\n\t\tfor path, found := range expected {\n\t\t\tif !found {\n\t\t\t\tt.Errorf(\"Expected absolute path not found: %s\", path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"GlobRecursiveWithRelativePattern\", func(t *testing.T) {\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"**/*.go\",\n\t\t\tPath:    \"/Users/bytedance/Desktop/github\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 3 {\n\t\t\tt.Fatalf(\"Expected 3 .go files with ** pattern, got %d\", len(infos))\n\t\t}\n\n\t\texpected := map[string]bool{\n\t\t\t\"eino/file1.go\":                      false,\n\t\t\t\"openai-go/paginationmanual_test.go\": false,\n\t\t\t\"openai-go/paginationauto_test.go\":   false,\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif _, ok := expected[info.Path]; ok {\n\t\t\t\texpected[info.Path] = true\n\t\t\t} else {\n\t\t\t\tt.Errorf(\"Unexpected path: %s\", info.Path)\n\t\t\t}\n\t\t}\n\n\t\tfor path, found := range expected {\n\t\t\tif !found {\n\t\t\t\tt.Errorf(\"Expected relative path not found: %s\", path)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_GlobInfo_RecursivePattern(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/project/src/main.go\",\n\t\tContent:  \"main\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/project/src/utils/helper.go\",\n\t\tContent:  \"helper\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/project/src/utils/deep/nested.go\",\n\t\tContent:  \"nested\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/project/test/test.go\",\n\t\tContent:  \"test\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/project/README.md\",\n\t\tContent:  \"readme\",\n\t})\n\n\tt.Run(\"DoubleStarMatchesAllSubdirectories\", func(t *testing.T) {\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"**/*.go\",\n\t\t\tPath:    \"/project\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 4 {\n\t\t\tt.Fatalf(\"Expected 4 .go files, got %d\", len(infos))\n\t\t}\n\n\t\texpected := map[string]bool{\n\t\t\t\"src/main.go\":              false,\n\t\t\t\"src/utils/helper.go\":      false,\n\t\t\t\"src/utils/deep/nested.go\": false,\n\t\t\t\"test/test.go\":             false,\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif _, ok := expected[info.Path]; ok {\n\t\t\t\texpected[info.Path] = true\n\t\t\t} else {\n\t\t\t\tt.Errorf(\"Unexpected path: %s\", info.Path)\n\t\t\t}\n\t\t}\n\n\t\tfor path, found := range expected {\n\t\t\tif !found {\n\t\t\t\tt.Errorf(\"Expected path not found: %s\", path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"DoubleStarInMiddleOfPattern\", func(t *testing.T) {\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"src/**/*.go\",\n\t\t\tPath:    \"/project\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 3 {\n\t\t\tt.Fatalf(\"Expected 3 .go files under src/, got %d\", len(infos))\n\t\t}\n\n\t\texpected := map[string]bool{\n\t\t\t\"src/main.go\":              false,\n\t\t\t\"src/utils/helper.go\":      false,\n\t\t\t\"src/utils/deep/nested.go\": false,\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif _, ok := expected[info.Path]; ok {\n\t\t\t\texpected[info.Path] = true\n\t\t\t} else {\n\t\t\t\tt.Errorf(\"Unexpected path: %s\", info.Path)\n\t\t\t}\n\t\t}\n\n\t\tfor path, found := range expected {\n\t\t\tif !found {\n\t\t\t\tt.Errorf(\"Expected path not found: %s\", path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"DoubleStarAtEnd\", func(t *testing.T) {\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"src/**\",\n\t\t\tPath:    \"/project\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 3 {\n\t\t\tt.Fatalf(\"Expected 3 files under src/, got %d\", len(infos))\n\t\t}\n\n\t\texpected := map[string]bool{\n\t\t\t\"src/main.go\":              false,\n\t\t\t\"src/utils/helper.go\":      false,\n\t\t\t\"src/utils/deep/nested.go\": false,\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif _, ok := expected[info.Path]; ok {\n\t\t\t\texpected[info.Path] = true\n\t\t\t}\n\t\t}\n\n\t\tfor path, found := range expected {\n\t\t\tif !found {\n\t\t\t\tt.Errorf(\"Expected path not found: %s\", path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"AbsolutePatternWithDoubleStarRecursive\", func(t *testing.T) {\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"/project/**/*.go\",\n\t\t\tPath:    \"/\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 4 {\n\t\t\tt.Fatalf(\"Expected 4 .go files, got %d\", len(infos))\n\t\t}\n\n\t\texpected := map[string]bool{\n\t\t\t\"/project/src/main.go\":              false,\n\t\t\t\"/project/src/utils/helper.go\":      false,\n\t\t\t\"/project/src/utils/deep/nested.go\": false,\n\t\t\t\"/project/test/test.go\":             false,\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif _, ok := expected[info.Path]; ok {\n\t\t\t\texpected[info.Path] = true\n\t\t\t} else {\n\t\t\t\tt.Errorf(\"Unexpected path: %s\", info.Path)\n\t\t\t}\n\t\t}\n\n\t\tfor path, found := range expected {\n\t\t\tif !found {\n\t\t\t\tt.Errorf(\"Expected absolute path not found: %s\", path)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_Concurrent(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\t// Test concurrent writes and reads\n\tdone := make(chan bool)\n\tfor i := 0; i < 10; i++ {\n\t\tgo func(n int) {\n\t\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\t\tFilePath: \"/concurrent.txt\",\n\t\t\t\tContent:  \"content\",\n\t\t\t})\n\t\t\tbackend.Read(ctx, &ReadRequest{\n\t\t\t\tFilePath: \"/concurrent.txt\",\n\t\t\t\tLimit:    10,\n\t\t\t})\n\t\t\tdone <- true\n\t\t}(i)\n\t}\n\n\tfor i := 0; i < 10; i++ {\n\t\t<-done\n\t}\n}\n\nfunc TestInMemoryBackend_LsInfo_FileInfoMetadata(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tt.Run(\"FileMetadata\", func(t *testing.T) {\n\t\tcontent := \"hello world\"\n\t\terr := backend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tContent:  content,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Write failed: %v\", err)\n\t\t}\n\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos))\n\t\t}\n\n\t\tinfo := infos[0]\n\t\tif info.Path != \"test.txt\" {\n\t\t\tt.Errorf(\"Expected path test.txt, got %s\", info.Path)\n\t\t}\n\t\tif info.IsDir {\n\t\t\tt.Error(\"Expected IsDir to be false for file\")\n\t\t}\n\t\tif info.Size != int64(len(content)) {\n\t\t\tt.Errorf(\"Expected size %d, got %d\", len(content), info.Size)\n\t\t}\n\t\tif info.ModifiedAt == \"\" {\n\t\t\tt.Error(\"Expected ModifiedAt to be non-empty\")\n\t\t}\n\t\t_, err = time.Parse(time.RFC3339Nano, info.ModifiedAt)\n\t\tif err != nil {\n\t\t\tt.Errorf(\"ModifiedAt is not valid RFC3339 format: %v\", err)\n\t\t}\n\t})\n\n\tt.Run(\"DirectoryMetadata\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\terr := backend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/dir1/file1.txt\",\n\t\t\tContent:  \"content1\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Write failed: %v\", err)\n\t\t}\n\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 directory, got %d\", len(infos))\n\t\t}\n\n\t\tinfo := infos[0]\n\t\tif info.Path != \"dir1\" {\n\t\t\tt.Errorf(\"Expected path dir1, got %s\", info.Path)\n\t\t}\n\t\tif !info.IsDir {\n\t\t\tt.Error(\"Expected IsDir to be true for directory\")\n\t\t}\n\t\tif info.Size != 0 {\n\t\t\tt.Errorf(\"Expected size 0 for directory, got %d\", info.Size)\n\t\t}\n\t\tif info.ModifiedAt == \"\" {\n\t\t\tt.Error(\"Expected ModifiedAt to be non-empty for directory\")\n\t\t}\n\t})\n\n\tt.Run(\"MixedFilesAndDirectories\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/file1.txt\",\n\t\t\tContent:  \"content1\",\n\t\t})\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/dir1/file2.txt\",\n\t\t\tContent:  \"content2\",\n\t\t})\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/dir1/subdir/file3.txt\",\n\t\t\tContent:  \"content3\",\n\t\t})\n\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 2 {\n\t\t\tt.Fatalf(\"Expected 2 items (file1.txt, dir1), got %d\", len(infos))\n\t\t}\n\n\t\tfileCount := 0\n\t\tdirCount := 0\n\t\tfor _, info := range infos {\n\t\t\tif info.IsDir {\n\t\t\t\tdirCount++\n\t\t\t\tif info.Path != \"dir1\" {\n\t\t\t\t\tt.Errorf(\"Expected directory path dir1, got %s\", info.Path)\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tfileCount++\n\t\t\t\tif info.Path != \"file1.txt\" {\n\t\t\t\t\tt.Errorf(\"Expected file path file1.txt, got %s\", info.Path)\n\t\t\t\t}\n\t\t\t\tif info.Size != int64(len(\"content1\")) {\n\t\t\t\t\tt.Errorf(\"Expected file size %d, got %d\", len(\"content1\"), info.Size)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tif fileCount != 1 {\n\t\t\tt.Errorf(\"Expected 1 file, got %d\", fileCount)\n\t\t}\n\t\tif dirCount != 1 {\n\t\t\tt.Errorf(\"Expected 1 directory, got %d\", dirCount)\n\t\t}\n\t})\n\n\tt.Run(\"SubdirectoryListing\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/dir1/file1.txt\",\n\t\t\tContent:  \"short\",\n\t\t})\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/dir1/subdir/file2.txt\",\n\t\t\tContent:  \"longer content here\",\n\t\t})\n\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/dir1\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 2 {\n\t\t\tt.Fatalf(\"Expected 2 items (file1.txt, subdir), got %d\", len(infos))\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif info.Path == \"file1.txt\" {\n\t\t\t\tif info.IsDir {\n\t\t\t\t\tt.Error(\"Expected file1.txt to be a file\")\n\t\t\t\t}\n\t\t\t\tif info.Size != int64(len(\"short\")) {\n\t\t\t\t\tt.Errorf(\"Expected size %d, got %d\", len(\"short\"), info.Size)\n\t\t\t\t}\n\t\t\t} else if info.Path == \"subdir\" {\n\t\t\t\tif !info.IsDir {\n\t\t\t\t\tt.Error(\"Expected subdir to be a directory\")\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tt.Errorf(\"Unexpected path: %s\", info.Path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"DirectoryModifiedAtUsesLatestFile\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/dir1/file1.txt\",\n\t\t\tContent:  \"content1\",\n\t\t})\n\t\ttime.Sleep(10 * time.Millisecond)\n\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/dir1/file2.txt\",\n\t\t\tContent:  \"content2\",\n\t\t})\n\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 directory, got %d\", len(infos))\n\t\t}\n\n\t\tdirInfo := infos[0]\n\t\tif !dirInfo.IsDir {\n\t\t\tt.Fatal(\"Expected directory\")\n\t\t}\n\n\t\tdirModTime, _ := time.Parse(time.RFC3339Nano, dirInfo.ModifiedAt)\n\n\t\tsubInfos, _ := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/dir1\"})\n\t\tvar latestFileTime time.Time\n\t\tfor _, info := range subInfos {\n\t\t\tfileTime, _ := time.Parse(time.RFC3339Nano, info.ModifiedAt)\n\t\t\tif fileTime.After(latestFileTime) {\n\t\t\t\tlatestFileTime = fileTime\n\t\t\t}\n\t\t}\n\n\t\tif !dirModTime.Equal(latestFileTime) && dirModTime.Before(latestFileTime) {\n\t\t\tt.Logf(\"Directory mod time: %v, Latest file time: %v\", dirModTime, latestFileTime)\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_GlobInfo_FileInfoMetadata(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tt.Run(\"BasicMetadata\", func(t *testing.T) {\n\t\tcontent := \"test content\"\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tContent:  content,\n\t\t})\n\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"*.txt\",\n\t\t\tPath:    \"/\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos))\n\t\t}\n\n\t\tinfo := infos[0]\n\t\tif info.Path != \"test.txt\" {\n\t\t\tt.Errorf(\"Expected path test.txt, got %s\", info.Path)\n\t\t}\n\t\tif info.IsDir {\n\t\t\tt.Error(\"Expected IsDir to be false\")\n\t\t}\n\t\tif info.Size != int64(len(content)) {\n\t\t\tt.Errorf(\"Expected size %d, got %d\", len(content), info.Size)\n\t\t}\n\t\tif info.ModifiedAt == \"\" {\n\t\t\tt.Error(\"Expected ModifiedAt to be non-empty\")\n\t\t}\n\t})\n\n\tt.Run(\"MultipleFilesMetadata\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/file1.txt\",\n\t\t\tContent:  \"short\",\n\t\t})\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/file2.txt\",\n\t\t\tContent:  \"much longer content\",\n\t\t})\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/file3.py\",\n\t\t\tContent:  \"python\",\n\t\t})\n\n\t\tinfos, err := backend.GlobInfo(ctx, &GlobInfoRequest{\n\t\t\tPattern: \"*.txt\",\n\t\t\tPath:    \"/\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GlobInfo failed: %v\", err)\n\t\t}\n\n\t\tif len(infos) != 2 {\n\t\t\tt.Fatalf(\"Expected 2 .txt files, got %d\", len(infos))\n\t\t}\n\n\t\tfor _, info := range infos {\n\t\t\tif info.IsDir {\n\t\t\t\tt.Errorf(\"Expected IsDir to be false for %s\", info.Path)\n\t\t\t}\n\t\t\tif info.Size <= 0 {\n\t\t\t\tt.Errorf(\"Expected positive size for %s, got %d\", info.Path, info.Size)\n\t\t\t}\n\t\t\tif info.ModifiedAt == \"\" {\n\t\t\t\tt.Errorf(\"Expected ModifiedAt to be non-empty for %s\", info.Path)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_WriteAndEdit_ModifiedAt(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"WriteUpdatesModifiedAt\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbeforeWrite := time.Now()\n\t\ttime.Sleep(1 * time.Millisecond)\n\n\t\terr := backend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tContent:  \"initial content\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Write failed: %v\", err)\n\t\t}\n\n\t\ttime.Sleep(1 * time.Millisecond)\n\t\tafterWrite := time.Now()\n\n\t\tinfos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\t\tif len(infos) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos))\n\t\t}\n\n\t\tmodTime, err := time.Parse(time.RFC3339Nano, infos[0].ModifiedAt)\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Failed to parse ModifiedAt: %v\", err)\n\t\t}\n\n\t\tif modTime.Before(beforeWrite) || modTime.After(afterWrite) {\n\t\t\tt.Errorf(\"ModifiedAt %v should be between %v and %v\", modTime, beforeWrite, afterWrite)\n\t\t}\n\t})\n\n\tt.Run(\"EditUpdatesModifiedAt\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\terr := backend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/edit.txt\",\n\t\t\tContent:  \"hello world\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Write failed: %v\", err)\n\t\t}\n\n\t\tinfos1, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\t\tif len(infos1) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos1))\n\t\t}\n\t\tmodTime1, err := time.Parse(time.RFC3339Nano, infos1[0].ModifiedAt)\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Failed to parse ModifiedAt: %v\", err)\n\t\t}\n\n\t\ttime.Sleep(10 * time.Millisecond)\n\n\t\terr = backend.Edit(ctx, &EditRequest{\n\t\t\tFilePath:   \"/edit.txt\",\n\t\t\tOldString:  \"hello\",\n\t\t\tNewString:  \"hi\",\n\t\t\tReplaceAll: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Edit failed: %v\", err)\n\t\t}\n\n\t\tinfos2, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\t\tif len(infos2) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos2))\n\t\t}\n\t\tmodTime2, err := time.Parse(time.RFC3339Nano, infos2[0].ModifiedAt)\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Failed to parse ModifiedAt: %v\", err)\n\t\t}\n\n\t\tif !modTime2.After(modTime1) {\n\t\t\tt.Errorf(\"ModifiedAt should be updated after edit. Before: %v, After: %v\", modTime1, modTime2)\n\t\t}\n\t})\n\n\tt.Run(\"OverwriteUpdatesModifiedAt\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\terr := backend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/overwrite.txt\",\n\t\t\tContent:  \"original\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Write failed: %v\", err)\n\t\t}\n\n\t\tinfos1, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\t\tif len(infos1) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos1))\n\t\t}\n\t\tmodTime1, err := time.Parse(time.RFC3339Nano, infos1[0].ModifiedAt)\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Failed to parse ModifiedAt: %v\", err)\n\t\t}\n\n\t\ttime.Sleep(10 * time.Millisecond)\n\n\t\terr = backend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/overwrite.txt\",\n\t\t\tContent:  \"new content\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Write failed: %v\", err)\n\t\t}\n\n\t\tinfos2, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\t\tif len(infos2) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos2))\n\t\t}\n\t\tmodTime2, err := time.Parse(time.RFC3339Nano, infos2[0].ModifiedAt)\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Failed to parse ModifiedAt: %v\", err)\n\t\t}\n\n\t\tif !modTime2.After(modTime1) {\n\t\t\tt.Errorf(\"ModifiedAt should be updated after overwrite. Before: %v, After: %v\", modTime1, modTime2)\n\t\t}\n\t})\n\n\tt.Run(\"SizeUpdatesAfterEdit\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\terr := backend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/size.txt\",\n\t\t\tContent:  \"hello\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Write failed: %v\", err)\n\t\t}\n\n\t\tinfos1, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\t\tif len(infos1) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos1))\n\t\t}\n\t\tsize1 := infos1[0].Size\n\n\t\terr = backend.Edit(ctx, &EditRequest{\n\t\t\tFilePath:   \"/size.txt\",\n\t\t\tOldString:  \"hello\",\n\t\t\tNewString:  \"hello world\",\n\t\t\tReplaceAll: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Edit failed: %v\", err)\n\t\t}\n\n\t\tinfos2, err := backend.LsInfo(ctx, &LsInfoRequest{Path: \"/\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"LsInfo failed: %v\", err)\n\t\t}\n\t\tif len(infos2) != 1 {\n\t\t\tt.Fatalf(\"Expected 1 file, got %d\", len(infos2))\n\t\t}\n\t\tsize2 := infos2[0].Size\n\n\t\tif size2 <= size1 {\n\t\t\tt.Errorf(\"Size should increase after edit. Before: %d, After: %d\", size1, size2)\n\t\t}\n\t\tif size2 != int64(len(\"hello world\")) {\n\t\t\tt.Errorf(\"Expected size %d, got %d\", len(\"hello world\"), size2)\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_Read_EdgeCases(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/test.txt\",\n\t\tContent:  \"line1\\nline2\\nline3\",\n\t})\n\n\tt.Run(\"negative offset should be treated as zero\", func(t *testing.T) {\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tOffset:   -5,\n\t\t\tLimit:    2,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Read failed: %v\", err)\n\t\t}\n\t\texpected := \"line1\\nline2\"\n\t\tif content.Content != expected {\n\t\t\tt.Errorf(\"Expected: %q, Got: %q\", expected, content.Content)\n\t\t}\n\t})\n\n\tt.Run(\"offset exceeds file length\", func(t *testing.T) {\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tOffset:   100,\n\t\t\tLimit:    10,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Read failed: %v\", err)\n\t\t}\n\t\tif content.Content != \"\" {\n\t\t\tt.Errorf(\"Expected empty content, got: %q\", content.Content)\n\t\t}\n\t})\n\n\tt.Run(\"zero or negative limit should use default 200\", func(t *testing.T) {\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tOffset:   0,\n\t\t\tLimit:    0,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Read failed: %v\", err)\n\t\t}\n\t\tlines := strings.Split(content.Content, \"\\n\")\n\t\tif len(lines) != 3 {\n\t\t\tt.Errorf(\"Expected 3 lines, got %d\", len(lines))\n\t\t}\n\t})\n\n\tt.Run(\"limit exceeds remaining lines\", func(t *testing.T) {\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tOffset:   1,\n\t\t\tLimit:    100,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Read failed: %v\", err)\n\t\t}\n\t\tlines := strings.Split(content.Content, \"\\n\")\n\t\tif len(lines) != 3 {\n\t\t\tt.Errorf(\"Expected 3 lines, got %d\", len(lines))\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_Edit_EdgeCases(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tt.Run(\"edit non-existent file\", func(t *testing.T) {\n\t\terr := backend.Edit(ctx, &EditRequest{\n\t\t\tFilePath:  \"/nonexistent.txt\",\n\t\t\tOldString: \"old\",\n\t\t\tNewString: \"new\",\n\t\t})\n\t\tif err == nil {\n\t\t\tt.Error(\"Expected error for non-existent file\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"not found\") {\n\t\t\tt.Errorf(\"Expected 'not found' error, got: %v\", err)\n\t\t}\n\t})\n\n\tt.Run(\"empty oldString\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tContent:  \"content\",\n\t\t})\n\n\t\terr := backend.Edit(ctx, &EditRequest{\n\t\t\tFilePath:  \"/test.txt\",\n\t\t\tOldString: \"\",\n\t\t\tNewString: \"new\",\n\t\t})\n\t\tif err == nil {\n\t\t\tt.Error(\"Expected error for empty oldString\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"non-empty\") {\n\t\t\tt.Errorf(\"Expected 'non-empty' error, got: %v\", err)\n\t\t}\n\t})\n\n\tt.Run(\"oldString not found\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tContent:  \"hello world\",\n\t\t})\n\n\t\terr := backend.Edit(ctx, &EditRequest{\n\t\t\tFilePath:  \"/test.txt\",\n\t\t\tOldString: \"notfound\",\n\t\t\tNewString: \"new\",\n\t\t})\n\t\tif err == nil {\n\t\t\tt.Error(\"Expected error when oldString not found\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"not found in file\") {\n\t\t\tt.Errorf(\"Expected 'not found in file' error, got: %v\", err)\n\t\t}\n\t})\n\n\tt.Run(\"multiple occurrences with ReplaceAll false\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tContent:  \"foo bar foo baz\",\n\t\t})\n\n\t\terr := backend.Edit(ctx, &EditRequest{\n\t\t\tFilePath:   \"/test.txt\",\n\t\t\tOldString:  \"foo\",\n\t\t\tNewString:  \"FOO\",\n\t\t\tReplaceAll: false,\n\t\t})\n\t\tif err == nil {\n\t\t\tt.Error(\"Expected error for multiple occurrences with ReplaceAll=false\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"multiple occurrences\") {\n\t\t\tt.Errorf(\"Expected 'multiple occurrences' error, got: %v\", err)\n\t\t}\n\t})\n\n\tt.Run(\"single occurrence with ReplaceAll false\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tContent:  \"foo bar baz\",\n\t\t})\n\n\t\terr := backend.Edit(ctx, &EditRequest{\n\t\t\tFilePath:   \"/test.txt\",\n\t\t\tOldString:  \"foo\",\n\t\t\tNewString:  \"FOO\",\n\t\t\tReplaceAll: false,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Edit failed: %v\", err)\n\t\t}\n\n\t\tcontent, _ := backend.Read(ctx, &ReadRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tLimit:    100,\n\t\t})\n\t\tif !strings.Contains(content.Content, \"FOO\") {\n\t\t\tt.Error(\"Expected content to contain 'FOO'\")\n\t\t}\n\t})\n\n\tt.Run(\"ReplaceAll replaces all occurrences\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tContent:  \"foo bar foo baz foo\",\n\t\t})\n\n\t\terr := backend.Edit(ctx, &EditRequest{\n\t\t\tFilePath:   \"/test.txt\",\n\t\t\tOldString:  \"foo\",\n\t\t\tNewString:  \"FOO\",\n\t\t\tReplaceAll: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"Edit failed: %v\", err)\n\t\t}\n\n\t\tcontent, _ := backend.Read(ctx, &ReadRequest{\n\t\t\tFilePath: \"/test.txt\",\n\t\t\tLimit:    100,\n\t\t})\n\t\tif strings.Contains(content.Content, \"foo\") {\n\t\t\tt.Error(\"Expected all 'foo' to be replaced\")\n\t\t}\n\t\tfooCount := strings.Count(content.Content, \"FOO\")\n\t\tif fooCount != 3 {\n\t\t\tt.Errorf(\"Expected 3 occurrences of 'FOO', got %d\", fooCount)\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_NormalizePath(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tt.Run(\"paths are normalized on write\", func(t *testing.T) {\n\t\ttestCases := []struct {\n\t\t\tinputPath      string\n\t\t\tnormalizedPath string\n\t\t}{\n\t\t\t{\"test.txt\", \"/test.txt\"},\n\t\t\t{\"/test.txt\", \"/test.txt\"},\n\t\t\t{\"//test.txt\", \"/test.txt\"},\n\t\t\t{\"/dir//file.txt\", \"/dir/file.txt\"},\n\t\t\t{\"/dir/../file.txt\", \"/file.txt\"},\n\t\t}\n\n\t\tfor _, tc := range testCases {\n\t\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\t\tFilePath: tc.inputPath,\n\t\t\t\tContent:  \"content\",\n\t\t\t})\n\n\t\t\tcontent, err := backend.Read(ctx, &ReadRequest{\n\t\t\t\tFilePath: tc.normalizedPath,\n\t\t\t\tLimit:    10,\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tt.Errorf(\"Failed to read normalized path %s (from %s): %v\", tc.normalizedPath, tc.inputPath, err)\n\t\t\t}\n\t\t\tif !strings.Contains(content.Content, \"content\") {\n\t\t\t\tt.Errorf(\"Content not found for normalized path %s (from %s)\", tc.normalizedPath, tc.inputPath)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_MatchFileType(t *testing.T) {\n\ttestCases := []struct {\n\t\text      string\n\t\tfileType string\n\t\texpected bool\n\t}{\n\t\t{\"go\", \"go\", true},\n\t\t{\"py\", \"python\", true},\n\t\t{\"py\", \"py\", true},\n\t\t{\"js\", \"js\", true},\n\t\t{\"ts\", \"typescript\", true},\n\t\t{\"ts\", \"ts\", true},\n\t\t{\"cpp\", \"cpp\", true},\n\t\t{\"c\", \"c\", true},\n\t\t{\"h\", \"c\", true},\n\t\t{\"md\", \"markdown\", true},\n\t\t{\"txt\", \"txt\", true},\n\t\t{\"go\", \"python\", false},\n\t\t{\"js\", \"typescript\", false},\n\t\t{\"unknown\", \"go\", false},\n\t}\n\n\tfor _, tc := range testCases {\n\t\tt.Run(fmt.Sprintf(\"%s matches %s\", tc.ext, tc.fileType), func(t *testing.T) {\n\t\t\tresult := matchFileType(tc.ext, tc.fileType)\n\t\t\tif result != tc.expected {\n\t\t\t\tt.Errorf(\"matchFileType(%q, %q) = %v, expected %v\", tc.ext, tc.fileType, result, tc.expected)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestInMemoryBackend_GrepRaw(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/test.go\",\n\t\tContent:  \"package main\\nfunc main() {\\n\\tlog.Error(\\\"error\\\")\\n\\tfmt.Println(\\\"hello\\\")\\n}\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/test.py\",\n\t\tContent:  \"def hello():\\n    print('error')\\n    print('world')\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/dir/file.go\",\n\t\tContent:  \"package test\\nfunc TestError() {\\n\\tlog.Error(\\\"test error\\\")\\n}\",\n\t})\n\n\tt.Run(\"basic pattern search\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"error\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 3 {\n\t\t\tt.Errorf(\"Expected 2 matches, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"empty pattern error\", func(t *testing.T) {\n\t\t_, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"\",\n\t\t})\n\t\tif err == nil {\n\t\t\tt.Error(\"Expected error for empty pattern\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"cannot be empty\") {\n\t\t\tt.Errorf(\"Expected 'cannot be empty' error, got: %v\", err)\n\t\t}\n\t})\n\n\tt.Run(\"invalid regex pattern\", func(t *testing.T) {\n\t\t_, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"[invalid\",\n\t\t})\n\t\tif err == nil {\n\t\t\tt.Error(\"Expected error for invalid regex\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"invalid regex\") {\n\t\t\tt.Errorf(\"Expected 'invalid regex' error, got: %v\", err)\n\t\t}\n\t})\n\n\tt.Run(\"case sensitive search\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"Error\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 3 {\n\t\t\tt.Errorf(\"Expected 2 matches, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"case insensitive search\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:         \"ERROR\",\n\t\t\tCaseInsensitive: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) < 3 {\n\t\t\tt.Errorf(\"Expected at least 2 matches, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"filter by file type\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:  \"error\",\n\t\t\tFileType: \"go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif !strings.HasSuffix(match.Path, \".go\") {\n\t\t\t\tt.Errorf(\"Expected only .go files, got: %s\", match.Path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"filter by glob pattern\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"Error\",\n\t\t\tGlob:    \"*.go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif !strings.HasSuffix(match.Path, \".go\") {\n\t\t\t\tt.Errorf(\"Expected only .go files, got: %s\", match.Path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"invalid glob pattern\", func(t *testing.T) {\n\t\t_, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"error\",\n\t\t\tGlob:    \"[invalid\",\n\t\t})\n\t\tif err == nil {\n\t\t\tt.Error(\"Expected error for invalid glob pattern\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"invalid glob\") {\n\t\t\tt.Errorf(\"Expected 'invalid glob' error, got: %v\", err)\n\t\t}\n\t})\n\n\tt.Run(\"search in specific path\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"Error\",\n\t\t\tPath:    \"/dir\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif !strings.HasPrefix(match.Path, \"/dir\") {\n\t\t\t\tt.Errorf(\"Expected matches only from /dir, got: %s\", match.Path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"search with non-existent path\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"error\",\n\t\t\tPath:    \"/nonexistent\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 0 {\n\t\t\tt.Errorf(\"Expected 0 matches for non-existent path, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"regex pattern matching\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"log\\\\..*Error\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) < 1 {\n\t\t\tt.Errorf(\"Expected at least 1 match, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"no matches found\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"nonexistent_pattern_xyz\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 0 {\n\t\t\tt.Errorf(\"Expected 0 matches, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"match line numbers\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:  \"log\\\\.Error\",\n\t\t\tFileType: \"go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif match.Line <= 0 {\n\t\t\t\tt.Errorf(\"Expected positive line number, got %d\", match.Line)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"match content is returned\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"package main\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) < 1 {\n\t\t\tt.Fatal(\"Expected at least 1 match\")\n\t\t}\n\t\tfound := false\n\t\tfor _, match := range matches {\n\t\t\tif strings.Contains(match.Content, \"package main\") {\n\t\t\t\tfound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tif !found {\n\t\t\tt.Error(\"Expected match content to contain 'package main'\")\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_GrepRaw_WithContext(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/context.txt\",\n\t\tContent:  \"line1\\nline2\\ntarget line\\nline4\\nline5\\nline6\",\n\t})\n\n\tt.Run(\"with before context\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:     \"target\",\n\t\t\tBeforeLines: 2,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) < 3 {\n\t\t\tt.Errorf(\"Expected at least 3 matches (2 before + target), got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"with after context\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:    \"target\",\n\t\t\tAfterLines: 2,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) < 3 {\n\t\t\tt.Errorf(\"Expected at least 3 matches (target + 2 after), got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"with both before and after context\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:     \"target\",\n\t\t\tBeforeLines: 1,\n\t\t\tAfterLines:  1,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) < 3 {\n\t\t\tt.Errorf(\"Expected at least 3 matches (1 before + target + 1 after), got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"context at file boundaries\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/boundary.txt\",\n\t\t\tContent:  \"first line target\\nsecond line\",\n\t\t})\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:     \"target\",\n\t\t\tPath:        \"/boundary.txt\",\n\t\t\tBeforeLines: 5,\n\t\t\tAfterLines:  5,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) == 0 {\n\t\t\tt.Error(\"Expected at least 1 match\")\n\t\t}\n\t})\n\n\tt.Run(\"zero context lines\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:     \"target\",\n\t\t\tBeforeLines: 0,\n\t\t\tAfterLines:  0,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) < 1 {\n\t\t\tt.Error(\"Expected at least 1 match\")\n\t\t}\n\t})\n\n\tt.Run(\"negative context lines treated as zero\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:     \"target\",\n\t\t\tBeforeLines: -5,\n\t\t\tAfterLines:  -5,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) < 1 {\n\t\t\tt.Error(\"Expected at least 1 match\")\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_GrepRaw_Multiline(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/multiline.txt\",\n\t\tContent:  \"start\\nmiddle line\\nend\",\n\t})\n\n\tt.Run(\"single line mode (default)\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:         \"start.*end\",\n\t\t\tEnableMultiline: false,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 0 {\n\t\t\tt.Errorf(\"Expected 0 matches in single-line mode, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"multiline mode enabled\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:         \"start[\\\\s\\\\S]*end\",\n\t\t\tEnableMultiline: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) == 0 {\n\t\t\tt.Error(\"Expected matches in multiline mode\")\n\t\t}\n\t})\n\n\tt.Run(\"multiline with multiple matches\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/multiline2.txt\",\n\t\t\tContent:  \"block1 start\\nblock1 middle\\nblock1 end\\n\\nblock2 start\\nblock2 end\",\n\t\t})\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:         \"start[\\\\s\\\\S]*?end\",\n\t\t\tPath:            \"/multiline2.txt\",\n\t\t\tEnableMultiline: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) == 0 {\n\t\t\tt.Error(\"Expected matches in multiline mode\")\n\t\t}\n\t})\n\n\tt.Run(\"multiline with multiple matches v2\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/multiline3.txt\",\n\t\t\tContent: `\nconst a = 1;\nfunction calculateTotal(\n  items,\n  discount\n) {\n  return items.reduce((sum, item) => sum + item.price, 0);\n}\n\nconst b = 2;\n\n/*\n * This is a comment\n * spanning multiple lines\n */\n\nclass UserService {\n  constructor(db) {\n    this.db = db;\n  }\n}\n`,\n\t\t})\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:         \"function calculateTotal\\\\([^\\\\)]*\\\\)\",\n\t\t\tPath:            \"/multiline3.txt\",\n\t\t\tEnableMultiline: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) == 0 {\n\t\t\tt.Error(\"Expected matches in multiline mode\")\n\t\t}\n\n\t\tfoundLastLine := false\n\t\tfor _, match := range matches {\n\t\t\tif match.Line == 6 && strings.Contains(match.Content, \") {\") {\n\t\t\t\tfoundLastLine = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tif !foundLastLine {\n\t\t\tt.Error(\"Expected to find line 5 with ') {' in content\")\n\t\t\tfor _, match := range matches {\n\t\t\t\tt.Logf(\"Line %d: %s\", match.Line, match.Content)\n\t\t\t}\n\t\t}\n\t})\n\n}\n\nfunc TestInMemoryBackend_GrepRaw_EmptyFiles(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tt.Run(\"search in empty file\", func(t *testing.T) {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/empty.txt\",\n\t\t\tContent:  \"\",\n\t\t})\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"anything\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 0 {\n\t\t\tt.Errorf(\"Expected 0 matches in empty file, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"search with no files\", func(t *testing.T) {\n\t\temptyBackend := NewInMemoryBackend()\n\t\tmatches, err := emptyBackend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"anything\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 0 {\n\t\t\tt.Errorf(\"Expected 0 matches with no files, got %d\", len(matches))\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_GrepRaw_SpecialCharacters(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/special.txt\",\n\t\tContent:  \"interface{}\\nmap[string]int\\nfunc() error\\n$variable\\n*pointer\",\n\t})\n\n\tt.Run(\"match curly braces\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"interface\\\\{\\\\}\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 1 {\n\t\t\tt.Errorf(\"Expected 1 match, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"match square brackets\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"map\\\\[.*\\\\]\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 1 {\n\t\t\tt.Errorf(\"Expected 1 match, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"match parentheses\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"func\\\\(\\\\)\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 1 {\n\t\t\tt.Errorf(\"Expected 1 match, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"match dollar sign\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"\\\\$variable\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 1 {\n\t\t\tt.Errorf(\"Expected 1 match, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"match asterisk\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"\\\\*pointer\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 1 {\n\t\t\tt.Errorf(\"Expected 1 match, got %d\", len(matches))\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_GrepRaw_Concurrent(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tfor i := 0; i < 10; i++ {\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: fmt.Sprintf(\"/file%d.txt\", i),\n\t\t\tContent:  fmt.Sprintf(\"content%d with error message\", i),\n\t\t})\n\t}\n\n\tt.Run(\"concurrent grep operations\", func(t *testing.T) {\n\t\tdone := make(chan bool)\n\t\tfor i := 0; i < 10; i++ {\n\t\t\tgo func() {\n\t\t\t\t_, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\t\t\tPattern: \"error\",\n\t\t\t\t})\n\t\t\t\tif err != nil {\n\t\t\t\t\tt.Errorf(\"Concurrent GrepRaw failed: %v\", err)\n\t\t\t\t}\n\t\t\t\tdone <- true\n\t\t\t}()\n\t\t}\n\n\t\tfor i := 0; i < 10; i++ {\n\t\t\t<-done\n\t\t}\n\t})\n\n\tt.Run(\"parallel file processing\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tfor i := 0; i < 100; i++ {\n\t\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\t\tFilePath: fmt.Sprintf(\"/large/file%d.go\", i),\n\t\t\t\tContent:  fmt.Sprintf(\"package main\\nimport \\\"log\\\"\\nfunc test%d() {\\n\\tlog.Error(\\\"error %d\\\")\\n}\", i, i),\n\t\t\t})\n\t\t}\n\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:  \"log\\\\.Error\",\n\t\t\tFileType: \"go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 100 {\n\t\t\tt.Errorf(\"Expected 100 matches, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"single file no parallelism\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: \"/single.txt\",\n\t\t\tContent:  \"error line 1\\nerror line 2\\nerror line 3\",\n\t\t})\n\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"error\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 3 {\n\t\t\tt.Errorf(\"Expected 3 matches, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"empty files list\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"anything\",\n\t\t\tPath:    \"/nonexistent\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) != 0 {\n\t\t\tt.Errorf(\"Expected 0 matches, got %d\", len(matches))\n\t\t}\n\t})\n\n\tt.Run(\"concurrent operations are safe\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tfor i := 0; i < 20; i++ {\n\t\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\t\tFilePath: fmt.Sprintf(\"/concurrent/file%d.txt\", i),\n\t\t\t\tContent:  fmt.Sprintf(\"line1\\nline2\\npattern%d\\nline4\", i),\n\t\t\t})\n\t\t}\n\n\t\tdone := make(chan error, 5)\n\t\tfor i := 0; i < 5; i++ {\n\t\t\tgo func(id int) {\n\t\t\t\t_, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\t\t\tPattern: \"pattern\\\\d+\",\n\t\t\t\t})\n\t\t\t\tdone <- err\n\t\t\t}(i)\n\t\t}\n\n\t\tfor i := 0; i < 5; i++ {\n\t\t\tif err := <-done; err != nil {\n\t\t\t\tt.Errorf(\"Concurrent operation %d failed: %v\", i, err)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc BenchmarkInMemoryBackend_GrepRaw(b *testing.B) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tfor i := 0; i < 100; i++ {\n\t\tcontent := fmt.Sprintf(`package main\n\nimport (\n\t\"fmt\"\n\t\"log\"\n)\n\nfunc process%d() error {\n\tlog.Error(\"processing error %d\")\n\tfmt.Println(\"hello world\")\n\treturn nil\n}\n\nfunc calculate%d(x, y int) int {\n\treturn x + y\n}\n`, i, i, i)\n\t\tbackend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: fmt.Sprintf(\"/project/src/file%d.go\", i),\n\t\t\tContent:  content,\n\t\t})\n\t}\n\n\tb.Run(\"parallel_grep\", func(b *testing.B) {\n\t\tfor i := 0; i < b.N; i++ {\n\t\t\t_, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\t\tPattern:  \"log\\\\.Error\",\n\t\t\t\tFileType: \"go\",\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tb.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t\t}\n\t\t}\n\t})\n\n\tb.Run(\"with_glob_filter\", func(b *testing.B) {\n\t\tfor i := 0; i < b.N; i++ {\n\t\t\t_, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\t\tPattern: \"Error\",\n\t\t\t\tGlob:    \"**/*.go\",\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tb.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t\t}\n\t\t}\n\t})\n\n\tb.Run(\"case_insensitive\", func(b *testing.B) {\n\t\tfor i := 0; i < b.N; i++ {\n\t\t\t_, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\t\tPattern:         \"ERROR\",\n\t\t\t\tCaseInsensitive: true,\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tb.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_GrepRaw_ComplexScenarios(t *testing.T) {\n\tbackend := NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/project/src/main.go\",\n\t\tContent:  \"package main\\nimport \\\"log\\\"\\nfunc main() {\\n\\tlog.Error(\\\"error\\\")\\n}\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/project/src/utils/helper.go\",\n\t\tContent:  \"package utils\\nfunc Helper() error {\\n\\treturn nil\\n}\",\n\t})\n\tbackend.Write(ctx, &WriteRequest{\n\t\tFilePath: \"/project/test/main_test.go\",\n\t\tContent:  \"package main\\nimport \\\"testing\\\"\\nfunc TestMain(t *testing.T) {\\n}\",\n\t})\n\n\tt.Run(\"combine path and file type filters\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:  \"package\",\n\t\t\tPath:     \"/project/src\",\n\t\t\tFileType: \"go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif !strings.HasPrefix(match.Path, \"/project/src\") {\n\t\t\t\tt.Errorf(\"Expected path to start with /project/src, got: %s\", match.Path)\n\t\t\t}\n\t\t\tif !strings.HasSuffix(match.Path, \".go\") {\n\t\t\t\tt.Errorf(\"Expected .go file, got: %s\", match.Path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"complex regex with case insensitive\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern:         \"func\\\\s+\\\\w+\",\n\t\t\tCaseInsensitive: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tif len(matches) == 0 {\n\t\t\tt.Error(\"Expected at least 1 match for function declarations\")\n\t\t}\n\t})\n\n\tt.Run(\"glob with directory structure\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"package\",\n\t\t\tGlob:    \"*_test.go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif !strings.HasSuffix(match.Path, \"_test.go\") {\n\t\t\t\tt.Errorf(\"Expected test file, got: %s\", match.Path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"glob with recursive pattern\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"package\",\n\t\t\tGlob:    \"**/*.go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif !strings.HasSuffix(match.Path, \".go\") {\n\t\t\t\tt.Errorf(\"Expected .go file, got: %s\", match.Path)\n\t\t\t}\n\t\t}\n\t\tif len(matches) == 0 {\n\t\t\tt.Error(\"Expected at least 1 match for **/*.go pattern\")\n\t\t}\n\t})\n\n\tt.Run(\"glob with path prefix\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"package\",\n\t\t\tGlob:    \"src/**/*.go\",\n\t\t\tPath:    \"/project\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif !strings.HasPrefix(match.Path, \"/project/src\") {\n\t\t\t\tt.Errorf(\"Expected path to start with /project/src, got: %s\", match.Path)\n\t\t\t}\n\t\t\tif !strings.HasSuffix(match.Path, \".go\") {\n\t\t\t\tt.Errorf(\"Expected .go file, got: %s\", match.Path)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"glob simple filename pattern\", func(t *testing.T) {\n\t\tmatches, err := backend.GrepRaw(ctx, &GrepRequest{\n\t\t\tPattern: \"package\",\n\t\t\tGlob:    \"main.go\",\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"GrepRaw failed: %v\", err)\n\t\t}\n\t\tfor _, match := range matches {\n\t\t\tif filepath.Base(match.Path) != \"main.go\" {\n\t\t\t\tt.Errorf(\"Expected filename 'main.go', got: %s\", match.Path)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestInMemoryBackend_Read_Scenarios(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"empty file returns empty content\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/empty.txt\", Content: \"\"})\n\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: \"/empty.txt\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t\t}\n\t\tif content.Content != \"\" {\n\t\t\tt.Errorf(\"expected empty content, got %q\", content.Content)\n\t\t}\n\t})\n\n\tt.Run(\"single-line file without newline\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/single.txt\", Content: \"hello\"})\n\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: \"/single.txt\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t\t}\n\t\tif content.Content != \"hello\" {\n\t\t\tt.Errorf(\"expected %q, got %q\", \"hello\", content.Content)\n\t\t}\n\t})\n\n\tt.Run(\"offset 0 and offset 1 both start from first line\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/f.txt\", Content: \"a\\nb\\nc\"})\n\n\t\tc0, _ := backend.Read(ctx, &ReadRequest{FilePath: \"/f.txt\", Offset: 0, Limit: 1})\n\t\tc1, _ := backend.Read(ctx, &ReadRequest{FilePath: \"/f.txt\", Offset: 1, Limit: 1})\n\t\tif c0.Content != c1.Content {\n\t\t\tt.Errorf(\"Offset=0 (%q) and Offset=1 (%q) should return the same first line\", c0.Content, c1.Content)\n\t\t}\n\t\tif c0.Content != \"a\" {\n\t\t\tt.Errorf(\"expected first line %q, got %q\", \"a\", c0.Content)\n\t\t}\n\t})\n\n\tt.Run(\"file with trailing newline preserves trailing empty line\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/trail.txt\", Content: \"line1\\nline2\\n\"})\n\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: \"/trail.txt\"})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t\t}\n\t\tif content.Content != \"line1\\nline2\\n\" {\n\t\t\tt.Errorf(\"expected %q, got %q\", \"line1\\nline2\\n\", content.Content)\n\t\t}\n\t\tlines := strings.Split(content.Content, \"\\n\")\n\t\tif len(lines) != 3 { // [\"line1\", \"line2\", \"\"]\n\t\t\tt.Errorf(\"expected 3 elements from split, got %d\", len(lines))\n\t\t}\n\t})\n\n\tt.Run(\"offset exactly at last line\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/f.txt\", Content: \"a\\nb\\nc\"})\n\n\t\t// Offset=3 (1-based) → last line \"c\"\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: \"/f.txt\", Offset: 3, Limit: 10})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t\t}\n\t\tif content.Content != \"c\" {\n\t\t\tt.Errorf(\"expected %q, got %q\", \"c\", content.Content)\n\t\t}\n\t})\n\n\tt.Run(\"offset one beyond last line returns empty\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/f.txt\", Content: \"a\\nb\\nc\"})\n\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: \"/f.txt\", Offset: 4, Limit: 10})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t\t}\n\t\tif content.Content != \"\" {\n\t\t\tt.Errorf(\"expected empty content, got %q\", content.Content)\n\t\t}\n\t})\n\n\tt.Run(\"limit=1 reads exactly one line\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/f.txt\", Content: \"a\\nb\\nc\"})\n\n\t\tfor i, expected := range []string{\"a\", \"b\", \"c\"} {\n\t\t\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: \"/f.txt\", Offset: i + 1, Limit: 1})\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"line %d: unexpected error: %v\", i+1, err)\n\t\t\t}\n\t\t\tif content.Content != expected {\n\t\t\t\tt.Errorf(\"line %d: expected %q, got %q\", i+1, expected, content.Content)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"sliding window reads consecutive ranges correctly\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/f.txt\", Content: \"l1\\nl2\\nl3\\nl4\\nl5\"})\n\n\t\ttests := []struct {\n\t\t\toffset   int\n\t\t\tlimit    int\n\t\t\texpected string\n\t\t}{\n\t\t\t{1, 2, \"l1\\nl2\"},\n\t\t\t{2, 2, \"l2\\nl3\"},\n\t\t\t{3, 2, \"l3\\nl4\"},\n\t\t\t{4, 2, \"l4\\nl5\"},\n\t\t\t{5, 2, \"l5\"},\n\t\t}\n\t\tfor _, tt := range tests {\n\t\t\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: \"/f.txt\", Offset: tt.offset, Limit: tt.limit})\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"offset=%d limit=%d: unexpected error: %v\", tt.offset, tt.limit, err)\n\t\t\t}\n\t\t\tif content.Content != tt.expected {\n\t\t\t\tt.Errorf(\"offset=%d limit=%d: expected %q, got %q\", tt.offset, tt.limit, tt.expected, content.Content)\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"file with only newlines\", func(t *testing.T) {\n\t\tbackend := NewInMemoryBackend()\n\t\tbackend.Write(ctx, &WriteRequest{FilePath: \"/newlines.txt\", Content: \"\\n\\n\\n\"})\n\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: \"/newlines.txt\", Offset: 2, Limit: 1})\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t\t}\n\t\t// Line 2 is an empty string between two newlines\n\t\tif content.Content != \"\" {\n\t\t\tt.Errorf(\"expected empty line content, got %q\", content.Content)\n\t\t}\n\t})\n}\n"
  },
  {
    "path": "adk/flow.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"log\"\n\t\"runtime/debug\"\n\t\"strings\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/compose\"\n\ticb \"github.com/cloudwego/eino/internal/callbacks\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype HistoryEntry struct {\n\tIsUserInput bool\n\tAgentName   string\n\tMessage     Message\n}\n\ntype HistoryRewriter func(ctx context.Context, entries []*HistoryEntry) ([]Message, error)\n\ntype flowAgent struct {\n\tAgent\n\n\tsubAgents   []*flowAgent\n\tparentAgent *flowAgent\n\n\tdisallowTransferToParent bool\n\thistoryRewriter          HistoryRewriter\n\n\tcheckPointStore compose.CheckPointStore\n}\n\nfunc (a *flowAgent) deepCopy() *flowAgent {\n\tret := &flowAgent{\n\t\tAgent:                    a.Agent,\n\t\tsubAgents:                make([]*flowAgent, 0, len(a.subAgents)),\n\t\tparentAgent:              a.parentAgent,\n\t\tdisallowTransferToParent: a.disallowTransferToParent,\n\t\thistoryRewriter:          a.historyRewriter,\n\t\tcheckPointStore:          a.checkPointStore,\n\t}\n\n\tfor _, sa := range a.subAgents {\n\t\tret.subAgents = append(ret.subAgents, sa.deepCopy())\n\t}\n\treturn ret\n}\n\n// SetSubAgents sets sub-agents for the given agent and returns the updated agent.\nfunc SetSubAgents(ctx context.Context, agent Agent, subAgents []Agent) (ResumableAgent, error) {\n\treturn setSubAgents(ctx, agent, subAgents)\n}\n\ntype AgentOption func(options *flowAgent)\n\n// WithDisallowTransferToParent prevents a sub-agent from transferring to its parent.\nfunc WithDisallowTransferToParent() AgentOption {\n\treturn func(fa *flowAgent) {\n\t\tfa.disallowTransferToParent = true\n\t}\n}\n\n// WithHistoryRewriter sets a rewriter to transform conversation history.\nfunc WithHistoryRewriter(h HistoryRewriter) AgentOption {\n\treturn func(fa *flowAgent) {\n\t\tfa.historyRewriter = h\n\t}\n}\n\nfunc toFlowAgent(ctx context.Context, agent Agent, opts ...AgentOption) *flowAgent {\n\tvar fa *flowAgent\n\tvar ok bool\n\tif fa, ok = agent.(*flowAgent); !ok {\n\t\tfa = &flowAgent{Agent: agent}\n\t} else {\n\t\tfa = fa.deepCopy()\n\t}\n\tfor _, opt := range opts {\n\t\topt(fa)\n\t}\n\n\tif fa.historyRewriter == nil {\n\t\tfa.historyRewriter = buildDefaultHistoryRewriter(agent.Name(ctx))\n\t}\n\n\treturn fa\n}\n\n// AgentWithOptions wraps an agent with flow-specific options and returns it.\nfunc AgentWithOptions(ctx context.Context, agent Agent, opts ...AgentOption) Agent {\n\treturn toFlowAgent(ctx, agent, opts...)\n}\n\nfunc setSubAgents(ctx context.Context, agent Agent, subAgents []Agent) (*flowAgent, error) {\n\tfa := toFlowAgent(ctx, agent)\n\n\tif len(fa.subAgents) > 0 {\n\t\treturn nil, errors.New(\"agent's sub-agents has already been set\")\n\t}\n\n\tif onAgent, ok_ := fa.Agent.(OnSubAgents); ok_ {\n\t\terr := onAgent.OnSetSubAgents(ctx, subAgents)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tfor _, s := range subAgents {\n\t\tfsa := toFlowAgent(ctx, s)\n\n\t\tif fsa.parentAgent != nil {\n\t\t\treturn nil, errors.New(\"agent has already been set as a sub-agent of another agent\")\n\t\t}\n\n\t\tfsa.parentAgent = fa\n\t\tif onAgent, ok__ := fsa.Agent.(OnSubAgents); ok__ {\n\t\t\terr := onAgent.OnSetAsSubAgent(ctx, agent)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\tif fsa.disallowTransferToParent {\n\t\t\t\terr = onAgent.OnDisallowTransferToParent(ctx)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tfa.subAgents = append(fa.subAgents, fsa)\n\t}\n\n\treturn fa, nil\n}\n\nfunc (a *flowAgent) getAgent(ctx context.Context, name string) *flowAgent {\n\tfor _, subAgent := range a.subAgents {\n\t\tif subAgent.Name(ctx) == name {\n\t\t\treturn subAgent\n\t\t}\n\t}\n\n\tif a.parentAgent != nil && a.parentAgent.Name(ctx) == name {\n\t\treturn a.parentAgent\n\t}\n\n\treturn nil\n}\n\nfunc rewriteMessage(msg Message, agentName string) Message {\n\tvar sb strings.Builder\n\tsb.WriteString(\"For context:\")\n\tif msg.Role == schema.Assistant {\n\t\tif msg.Content != \"\" {\n\t\t\tsb.WriteString(fmt.Sprintf(\" [%s] said: %s.\", agentName, msg.Content))\n\t\t}\n\t\tif len(msg.ToolCalls) > 0 {\n\t\t\tfor i := range msg.ToolCalls {\n\t\t\t\tf := msg.ToolCalls[i].Function\n\t\t\t\tsb.WriteString(fmt.Sprintf(\" [%s] called tool: `%s` with arguments: %s.\",\n\t\t\t\t\tagentName, f.Name, f.Arguments))\n\t\t\t}\n\t\t}\n\t} else if msg.Role == schema.Tool && msg.Content != \"\" {\n\t\tsb.WriteString(fmt.Sprintf(\" [%s] `%s` tool returned result: %s.\",\n\t\t\tagentName, msg.ToolName, msg.Content))\n\t}\n\n\trewritten := schema.UserMessage(sb.String())\n\tif msg.MultiContent != nil {\n\t\trewritten.MultiContent = append([]schema.ChatMessagePart{}, msg.MultiContent...)\n\t}\n\tif msg.UserInputMultiContent != nil {\n\t\trewritten.UserInputMultiContent = append([]schema.MessageInputPart{}, msg.UserInputMultiContent...)\n\t}\n\n\t// Convert AssistantGenMultiContent to UserInputMultiContent, since the role changes to User.\n\t// Reasoning parts have no user input equivalent and are dropped.\n\tfor _, part := range msg.AssistantGenMultiContent {\n\t\tswitch part.Type {\n\t\tcase schema.ChatMessagePartTypeText:\n\t\t\trewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{\n\t\t\t\tType:  part.Type,\n\t\t\t\tText:  part.Text,\n\t\t\t\tExtra: part.Extra,\n\t\t\t})\n\t\tcase schema.ChatMessagePartTypeImageURL:\n\t\t\tif part.Image != nil {\n\t\t\t\trewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{\n\t\t\t\t\tType:  part.Type,\n\t\t\t\t\tImage: &schema.MessageInputImage{MessagePartCommon: part.Image.MessagePartCommon},\n\t\t\t\t\tExtra: part.Extra,\n\t\t\t\t})\n\t\t\t}\n\t\tcase schema.ChatMessagePartTypeAudioURL:\n\t\t\tif part.Audio != nil {\n\t\t\t\trewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{\n\t\t\t\t\tType:  part.Type,\n\t\t\t\t\tAudio: &schema.MessageInputAudio{MessagePartCommon: part.Audio.MessagePartCommon},\n\t\t\t\t\tExtra: part.Extra,\n\t\t\t\t})\n\t\t\t}\n\t\tcase schema.ChatMessagePartTypeVideoURL:\n\t\t\tif part.Video != nil {\n\t\t\t\trewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{\n\t\t\t\t\tType:  part.Type,\n\t\t\t\t\tVideo: &schema.MessageInputVideo{MessagePartCommon: part.Video.MessagePartCommon},\n\t\t\t\t\tExtra: part.Extra,\n\t\t\t\t})\n\t\t\t}\n\t\t}\n\t}\n\n\treturn rewritten\n}\n\nfunc genMsg(entry *HistoryEntry, agentName string) (Message, error) {\n\tmsg := entry.Message\n\tif entry.AgentName != agentName {\n\t\tmsg = rewriteMessage(msg, entry.AgentName)\n\t}\n\n\treturn msg, nil\n}\n\nfunc (ai *AgentInput) deepCopy() *AgentInput {\n\tcopied := &AgentInput{\n\t\tMessages:        make([]Message, len(ai.Messages)),\n\t\tEnableStreaming: ai.EnableStreaming,\n\t}\n\n\tcopy(copied.Messages, ai.Messages)\n\n\treturn copied\n}\n\nfunc (a *flowAgent) genAgentInput(ctx context.Context, runCtx *runContext, skipTransferMessages bool) (*AgentInput, error) {\n\tinput := runCtx.RootInput.deepCopy()\n\n\tevents := runCtx.Session.getEvents()\n\thistoryEntries := make([]*HistoryEntry, 0)\n\n\tfor _, m := range input.Messages {\n\t\thistoryEntries = append(historyEntries, &HistoryEntry{\n\t\t\tIsUserInput: true,\n\t\t\tMessage:     m,\n\t\t})\n\t}\n\n\tfor _, event := range events {\n\t\tif skipTransferMessages && event.Action != nil && event.Action.TransferToAgent != nil {\n\t\t\t// If skipTransferMessages is true and the event contain transfer action, the message in this event won't be appended to history entries.\n\t\t\tif event.Output != nil &&\n\t\t\t\tevent.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Tool &&\n\t\t\t\tlen(historyEntries) > 0 {\n\t\t\t\t// If the skipped message's role is Tool, remove the previous history entry as it's also a transfer message(from ChatModelAgent and GenTransferMessages).\n\t\t\t\thistoryEntries = historyEntries[:len(historyEntries)-1]\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\n\t\tmsg, err := getMessageFromWrappedEvent(event)\n\t\tif err != nil {\n\t\t\tvar retryErr *WillRetryError\n\t\t\tif errors.As(err, &retryErr) {\n\t\t\t\tlog.Printf(\"failed to get message from event, but will retry: %v\", err)\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\treturn nil, err\n\t\t}\n\n\t\tif msg == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\thistoryEntries = append(historyEntries, &HistoryEntry{\n\t\t\tAgentName: event.AgentName,\n\t\t\tMessage:   msg,\n\t\t})\n\t}\n\n\tmessages, err := a.historyRewriter(ctx, historyEntries)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tinput.Messages = messages\n\n\treturn input, nil\n}\n\nfunc buildDefaultHistoryRewriter(agentName string) HistoryRewriter {\n\treturn func(ctx context.Context, entries []*HistoryEntry) ([]Message, error) {\n\t\tmessages := make([]Message, 0, len(entries))\n\t\tvar err error\n\t\tfor _, entry := range entries {\n\t\t\tmsg := entry.Message\n\t\t\tif !entry.IsUserInput {\n\t\t\t\tmsg, err = genMsg(entry, agentName)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, fmt.Errorf(\"gen agent input failed: %w\", err)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tif msg != nil {\n\t\t\t\tmessages = append(messages, msg)\n\t\t\t}\n\t\t}\n\n\t\treturn messages, nil\n\t}\n}\n\nfunc (a *flowAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tagentName := a.Name(ctx)\n\n\tvar runCtx *runContext\n\tctx, runCtx = initRunCtx(ctx, agentName, input)\n\tctx = AppendAddressSegment(ctx, AddressSegmentAgent, agentName)\n\n\to := getCommonOptions(nil, opts...)\n\n\tprocessedInput, err := a.genAgentInput(ctx, runCtx, o.skipTransferMessages)\n\tif err != nil {\n\t\tcbInput := &AgentCallbackInput{Input: input}\n\t\tctx = callbacks.OnStart(ctx, cbInput)\n\t\treturn wrapIterWithOnEnd(ctx, genErrorIter(err))\n\t}\n\n\tctxForSubAgents := ctx\n\n\tagentType := getAgentType(a.Agent)\n\tctx = initAgentCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)\n\tcbInput := &AgentCallbackInput{Input: processedInput}\n\tctx = callbacks.OnStart(ctx, cbInput)\n\n\tinput = processedInput\n\n\tif wf, ok := a.Agent.(*workflowAgent); ok {\n\t\treturn wrapIterWithOnEnd(ctx, wf.Run(ctx, input, filterCallbackHandlersForNestedAgents(agentName, opts)...))\n\t}\n\n\taIter := a.Agent.Run(ctx, input, filterOptions(agentName, opts)...)\n\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo a.run(ctx, ctxForSubAgents, runCtx, aIter, generator, opts...)\n\n\treturn iterator\n}\n\nfunc (a *flowAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tagentName := a.Name(ctx)\n\n\tctx, info = buildResumeInfo(ctx, agentName, info)\n\n\tctxForSubAgents := ctx\n\n\tagentType := getAgentType(a.Agent)\n\tctx = initAgentCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)\n\tcbInput := &AgentCallbackInput{ResumeInfo: info}\n\tctx = callbacks.OnStart(ctx, cbInput)\n\n\tif info.WasInterrupted {\n\t\tra, ok := a.Agent.(ResumableAgent)\n\t\tif !ok {\n\t\t\treturn wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf(\"failed to resume agent: agent '%s' is an interrupt point \"+\n\t\t\t\t\"but is not a ResumableAgent\", agentName)))\n\t\t}\n\t\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\t\tif _, ok := ra.(*workflowAgent); ok {\n\t\t\tfilteredOpts := filterCallbackHandlersForNestedAgents(agentName, opts)\n\t\t\taIter := ra.Resume(ctx, info, filteredOpts...)\n\t\t\treturn wrapIterWithOnEnd(ctx, aIter)\n\t\t}\n\t\taIter := ra.Resume(ctx, info, opts...)\n\t\tgo a.run(ctx, ctxForSubAgents, getRunCtx(ctxForSubAgents), aIter, generator, opts...)\n\t\treturn iterator\n\t}\n\n\tnextAgentName, err := getNextResumeAgent(ctx, info)\n\tif err != nil {\n\t\treturn wrapIterWithOnEnd(ctx, genErrorIter(err))\n\t}\n\n\tsubAgent := a.getAgent(ctxForSubAgents, nextAgentName)\n\tif subAgent == nil {\n\t\t// the inner agent wrapped by flowAgent may be ANY agent, including flowAgent,\n\t\t// AgentWithDeterministicTransferTo, or any other custom agent user defined,\n\t\t// or any combinations of the above in any order,\n\t\t// that ultimately wraps the flowAgent with sub-agents\n\t\t// We need to go through these wrappers to reach the flowAgent with sub-agents.\n\t\tif len(a.subAgents) == 0 {\n\t\t\tif ra, ok := a.Agent.(ResumableAgent); ok {\n\t\t\t\t// Use ctx (callback-enriched) instead of ctxForSubAgents here.\n\t\t\t\t// This is the inner agent that flowAgent wraps (e.g., supervisorContainer),\n\t\t\t\t// not a sub-agent. The callback context from OnStart should be propagated\n\t\t\t\t// to ensure unified tracing for container patterns.\n\t\t\t\treturn wrapIterWithOnEnd(ctx, ra.Resume(ctx, info, opts...))\n\t\t\t}\n\t\t}\n\t\treturn wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf(\"failed to resume agent: agent '%s' not found from flowAgent '%s'\", nextAgentName, agentName)))\n\t}\n\n\treturn wrapIterWithOnEnd(ctx, subAgent.Resume(ctxForSubAgents, info, opts...))\n}\n\ntype DeterministicTransferConfig struct {\n\tAgent        Agent\n\tToAgentNames []string\n}\n\nfunc (a *flowAgent) run(\n\tctx context.Context,\n\tctxForSubAgents context.Context,\n\trunCtx *runContext,\n\taIter *AsyncIterator[*AgentEvent],\n\tgenerator *AsyncGenerator[*AgentEvent],\n\topts ...AgentRunOption) {\n\n\tcbIter, cbGen := NewAsyncIteratorPair[*AgentEvent]()\n\n\tcbOutput := &AgentCallbackOutput{Events: cbIter}\n\ticb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyAgentCallbackOutput), callbacks.TimingOnEnd, false)\n\n\tdefer func() {\n\t\tpanicErr := recover()\n\t\tif panicErr != nil {\n\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\tgenerator.Send(&AgentEvent{Err: e})\n\t\t}\n\n\t\tcbGen.Close()\n\t\tgenerator.Close()\n\t}()\n\n\tvar lastAction *AgentAction\n\tfor {\n\t\tevent, ok := aIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\n\t\t// RunPath ownership: the eino framework sets RunPath exactly once.\n\t\t// If event.RunPath is already set (e.g., by agentTool), we don't modify it.\n\t\t// If event.RunPath is nil/empty, we set it to the current runCtx.RunPath.\n\t\t// This ensures RunPath is set exactly once and not duplicated.\n\t\tif len(event.RunPath) == 0 {\n\t\t\tevent.AgentName = a.Name(ctx)\n\t\t\tevent.RunPath = runCtx.RunPath\n\t\t}\n\t\t// Recording policy: exact RunPath match (non-interrupt) indicates events belonging to this agent execution.\n\t\t// This prevents parent recording of child/tool-internal emissions.\n\t\tif (event.Action == nil || event.Action.Interrupted == nil) && exactRunPathMatch(runCtx.RunPath, event.RunPath) {\n\t\t\t// copy the event so that the copied event's stream is exclusive for any potential consumer\n\t\t\t// copy before adding to session because once added to session it's stream could be consumed by genAgentInput at any time\n\t\t\t// interrupt action are not added to session, because ALL information contained in it\n\t\t\t// is either presented to end-user, or made available to agents through other means\n\t\t\tcopied := copyAgentEvent(event)\n\t\t\tsetAutomaticClose(copied)\n\t\t\tsetAutomaticClose(event)\n\t\t\trunCtx.Session.addEvent(copied)\n\t\t}\n\t\t// Action gating uses exact run-path match as well:\n\t\t// only actions originating from this agent execution (not child/tool runs)\n\t\t// should influence parent control flow (exit/transfer/interrupt).\n\t\tif exactRunPathMatch(runCtx.RunPath, event.RunPath) {\n\t\t\tlastAction = event.Action\n\t\t}\n\t\tcopied := copyAgentEvent(event)\n\t\tsetAutomaticClose(copied)\n\t\tsetAutomaticClose(event)\n\t\tcbGen.Send(copied)\n\t\tgenerator.Send(event)\n\t}\n\n\tvar destName string\n\tif lastAction != nil {\n\t\tif lastAction.Interrupted != nil {\n\t\t\treturn\n\t\t}\n\t\tif lastAction.Exit {\n\t\t\treturn\n\t\t}\n\n\t\tif lastAction.TransferToAgent != nil {\n\t\t\tdestName = lastAction.TransferToAgent.DestAgentName\n\t\t}\n\t}\n\n\t// handle transferring to another agent\n\tif destName != \"\" {\n\t\tagentToRun := a.getAgent(ctxForSubAgents, destName)\n\t\tif agentToRun == nil {\n\t\t\te := fmt.Errorf(\"transfer failed: agent '%s' not found when transferring from '%s'\",\n\t\t\t\tdestName, a.Name(ctxForSubAgents))\n\t\t\tgenerator.Send(&AgentEvent{Err: e})\n\t\t\treturn\n\t\t}\n\n\t\tsubAIter := agentToRun.Run(ctxForSubAgents, nil /*subagents get input from runCtx*/, opts...)\n\t\tfor {\n\t\t\tsubEvent, ok_ := subAIter.Next()\n\t\t\tif !ok_ {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tsetAutomaticClose(subEvent)\n\t\t\tgenerator.Send(subEvent)\n\t\t}\n\t}\n}\n\nfunc exactRunPathMatch(aPath, bPath []RunStep) bool {\n\tif len(aPath) != len(bPath) {\n\t\treturn false\n\t}\n\tfor i := range aPath {\n\t\tif !aPath[i].Equals(bPath[i]) {\n\t\t\treturn false\n\t\t}\n\t}\n\treturn true\n}\n\nfunc wrapIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*AgentEvent]) *AsyncIterator[*AgentEvent] {\n\tcbIter, cbGen := NewAsyncIteratorPair[*AgentEvent]()\n\tcbOutput := &AgentCallbackOutput{Events: cbIter}\n\ticb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyAgentCallbackOutput), callbacks.TimingOnEnd, false)\n\n\toutIter, outGen := NewAsyncIteratorPair[*AgentEvent]()\n\tgo func() {\n\t\tdefer func() {\n\t\t\tcbGen.Close()\n\t\t\toutGen.Close()\n\t\t}()\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tcopied := copyAgentEvent(event)\n\t\t\tcbGen.Send(copied)\n\t\t\toutGen.Send(event)\n\t\t}\n\t}()\n\treturn outIter\n}\n"
  },
  {
    "path": "adk/flow_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc strPtr(s string) *string { return &s }\n\nfunc TestRewriteMessage(t *testing.T) {\n\timageCommon := schema.MessagePartCommon{URL: strPtr(\"http://img.example.com\")}\n\taudioCommon := schema.MessagePartCommon{URL: strPtr(\"http://audio.example.com\")}\n\tvideoCommon := schema.MessagePartCommon{URL: strPtr(\"http://video.example.com\")}\n\n\tmsg := &schema.Message{\n\t\tRole:    schema.Assistant,\n\t\tContent: \"hello\",\n\t\tMultiContent: []schema.ChatMessagePart{\n\t\t\t{Type: schema.ChatMessagePartTypeText, Text: \"legacy\"},\n\t\t},\n\t\tUserInputMultiContent: []schema.MessageInputPart{\n\t\t\t{Type: schema.ChatMessagePartTypeText, Text: \"pre-existing\"},\n\t\t},\n\t\tAssistantGenMultiContent: []schema.MessageOutputPart{\n\t\t\t{Type: schema.ChatMessagePartTypeText, Text: \"gen-text\", Extra: map[string]any{\"k\": \"v\"}},\n\t\t\t{Type: schema.ChatMessagePartTypeImageURL, Image: &schema.MessageOutputImage{MessagePartCommon: imageCommon}},\n\t\t\t{Type: schema.ChatMessagePartTypeAudioURL, Audio: &schema.MessageOutputAudio{MessagePartCommon: audioCommon}},\n\t\t\t{Type: schema.ChatMessagePartTypeVideoURL, Video: &schema.MessageOutputVideo{MessagePartCommon: videoCommon}},\n\t\t\t{Type: schema.ChatMessagePartTypeReasoning, Reasoning: &schema.MessageOutputReasoning{Text: \"secret thoughts\"}},\n\t\t},\n\t}\n\n\trewritten := rewriteMessage(msg, \"OtherAgent\")\n\n\tassert.Equal(t, schema.User, rewritten.Role)\n\n\t// MultiContent: copied, not shared\n\tassert.Equal(t, msg.MultiContent, rewritten.MultiContent)\n\trewritten.MultiContent[0].Text = \"mutated\"\n\tassert.Equal(t, \"legacy\", msg.MultiContent[0].Text)\n\n\t// UserInputMultiContent: pre-existing entry copied, AssistantGenMultiContent appended (reasoning dropped)\n\tassert.Len(t, rewritten.UserInputMultiContent, 5) // 1 pre-existing + 4 converted (text/image/audio/video)\n\n\t// pre-existing entry is not shared\n\trewritten.UserInputMultiContent[0].Text = \"mutated\"\n\tassert.Equal(t, \"pre-existing\", msg.UserInputMultiContent[0].Text)\n\n\t// text conversion\n\tassert.Equal(t, schema.ChatMessagePartTypeText, rewritten.UserInputMultiContent[1].Type)\n\tassert.Equal(t, \"gen-text\", rewritten.UserInputMultiContent[1].Text)\n\tassert.Equal(t, map[string]any{\"k\": \"v\"}, rewritten.UserInputMultiContent[1].Extra)\n\n\t// image conversion\n\tassert.Equal(t, schema.ChatMessagePartTypeImageURL, rewritten.UserInputMultiContent[2].Type)\n\tassert.Equal(t, imageCommon, rewritten.UserInputMultiContent[2].Image.MessagePartCommon)\n\n\t// audio conversion\n\tassert.Equal(t, schema.ChatMessagePartTypeAudioURL, rewritten.UserInputMultiContent[3].Type)\n\tassert.Equal(t, audioCommon, rewritten.UserInputMultiContent[3].Audio.MessagePartCommon)\n\n\t// video conversion\n\tassert.Equal(t, schema.ChatMessagePartTypeVideoURL, rewritten.UserInputMultiContent[4].Type)\n\tassert.Equal(t, videoCommon, rewritten.UserInputMultiContent[4].Video.MessagePartCommon)\n\n\t// reasoning is dropped; AssistantGenMultiContent is not set on rewritten message\n\tassert.Empty(t, rewritten.AssistantGenMultiContent)\n}\n\n// TestTransferToAgent tests the TransferToAgent functionality\nfunc TestTransferToAgent(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create mock models for parent and child agents\n\tparentModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tchildModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// Set up expectations for the parent model\n\t// First call: parent model generates a message with TransferToAgent tool call\n\tparentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"I'll transfer this to the child agent\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      TransferToAgentToolName,\n\t\t\t\t\t\tArguments: `{\"agent_name\": \"ChildAgent\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\n\t// Set up expectations for the child model\n\t// Second call: child model generates a response\n\tchildModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"Hello from child agent\", nil), nil).\n\t\tTimes(1)\n\n\t// Both models should implement WithTools\n\tparentModel.EXPECT().WithTools(gomock.Any()).Return(parentModel, nil).AnyTimes()\n\tchildModel.EXPECT().WithTools(gomock.Any()).Return(childModel, nil).AnyTimes()\n\n\t// Create parent agent\n\tparentAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"ParentAgent\",\n\t\tDescription: \"Parent agent that will transfer to child\",\n\t\tInstruction: \"You are a parent agent.\",\n\t\tModel:       parentModel,\n\t})\n\tassert.NoError(t, err)\n\tassert.NotNil(t, parentAgent)\n\n\t// Create child agent\n\tchildAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"ChildAgent\",\n\t\tDescription: \"Child agent that handles specific tasks\",\n\t\tInstruction: \"You are a child agent.\",\n\t\tModel:       childModel,\n\t})\n\tassert.NoError(t, err)\n\tassert.NotNil(t, childAgent)\n\n\t// Set up parent-child relationship\n\tflowAgent, err := SetSubAgents(ctx, parentAgent, []Agent{childAgent})\n\tassert.NoError(t, err)\n\tassert.NotNil(t, flowAgent)\n\n\tassert.NotNil(t, parentAgent.subAgents)\n\tassert.NotNil(t, childAgent.parentAgent)\n\n\t// Run the parent agent\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Please transfer this to the child agent\"),\n\t\t},\n\t}\n\tctx, _ = initRunCtx(ctx, flowAgent.Name(ctx), input)\n\titerator := flowAgent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// First event: parent model output with tool call\n\tevent1, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event1)\n\tassert.Nil(t, event1.Err)\n\tassert.NotNil(t, event1.Output)\n\tassert.NotNil(t, event1.Output.MessageOutput)\n\tassert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role)\n\n\t// Second event: tool output (TransferToAgent)\n\tevent2, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event2)\n\tassert.Nil(t, event2.Err)\n\tassert.NotNil(t, event2.Output)\n\tassert.NotNil(t, event2.Output.MessageOutput)\n\tassert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role)\n\n\t// Verify the action is TransferToAgent\n\tassert.NotNil(t, event2.Action)\n\tassert.NotNil(t, event2.Action.TransferToAgent)\n\tassert.Equal(t, \"ChildAgent\", event2.Action.TransferToAgent.DestAgentName)\n\n\t// Third event: child model output\n\tevent3, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event3)\n\tassert.Nil(t, event3.Err)\n\tassert.NotNil(t, event3.Output)\n\tassert.NotNil(t, event3.Output.MessageOutput)\n\tassert.Equal(t, schema.Assistant, event3.Output.MessageOutput.Role)\n\n\t// Verify the message content from child agent\n\tmsg := event3.Output.MessageOutput.Message\n\tassert.NotNil(t, msg)\n\tassert.Equal(t, \"Hello from child agent\", msg.Content)\n\n\t// No more events\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestTransferToAgentWithDesignatedCallback(t *testing.T) {\n\tctx := context.Background()\n\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tparentModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tchildModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tparentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"I'll transfer this to the child agent\",\n\t\t\t[]schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      TransferToAgentToolName,\n\t\t\t\t\t\tArguments: `{\"agent_name\": \"ChildAgent\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}), nil).\n\t\tTimes(1)\n\n\tchildModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.AssistantMessage(\"Hello from child agent\", nil), nil).\n\t\tTimes(1)\n\n\tparentModel.EXPECT().WithTools(gomock.Any()).Return(parentModel, nil).AnyTimes()\n\tchildModel.EXPECT().WithTools(gomock.Any()).Return(childModel, nil).AnyTimes()\n\n\tparentAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"ParentAgent\",\n\t\tDescription: \"Parent agent that will transfer to child\",\n\t\tInstruction: \"You are a parent agent.\",\n\t\tModel:       parentModel,\n\t})\n\tassert.NoError(t, err)\n\n\tchildAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"ChildAgent\",\n\t\tDescription: \"Child agent that handles specific tasks\",\n\t\tInstruction: \"You are a child agent.\",\n\t\tModel:       childModel,\n\t})\n\tassert.NoError(t, err)\n\n\tflowAgent, err := SetSubAgents(ctx, parentAgent, []Agent{childAgent})\n\tassert.NoError(t, err)\n\n\tvar childCallbackCount int\n\tvar mu sync.Mutex\n\n\thandler := callbacks.NewHandlerBuilder().OnStartFn(\n\t\tfunc(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component == ComponentOfAgent && info.Name == \"ChildAgent\" {\n\t\t\t\tmu.Lock()\n\t\t\t\tchildCallbackCount++\n\t\t\t\tmu.Unlock()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).Build()\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Please transfer this to the child agent\"),\n\t\t},\n\t}\n\tctx, _ = initRunCtx(ctx, flowAgent.Name(ctx), input)\n\titerator := flowAgent.Run(ctx, input, WithCallbacks(handler).DesignateAgent(\"ChildAgent\"))\n\n\tfor {\n\t\t_, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassert.Equal(t, 1, childCallbackCount, \"designated callback for ChildAgent should fire exactly once during transfer\")\n}\n"
  },
  {
    "path": "adk/handler.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// InvokableToolCallEndpoint is the function signature for invoking a tool synchronously.\n// Middleware authors implement wrappers around this endpoint to add custom behavior.\ntype InvokableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error)\n\n// StreamableToolCallEndpoint is the function signature for invoking a tool with streaming output.\n// Middleware authors implement wrappers around this endpoint to add custom behavior.\ntype StreamableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error)\n\ntype EnhancedInvokableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error)\n\ntype EnhancedStreamableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error)\n\n// ToolContext provides metadata about the tool being wrapped.\ntype ToolContext struct {\n\tName   string\n\tCallID string\n}\n\n// ModelContext contains context information passed to WrapModel.\ntype ModelContext struct {\n\t// Tools contains the current tool list configured for the agent.\n\t// This is populated at request time with the tools that will be sent to the model.\n\tTools []*schema.ToolInfo\n\n\t// ModelRetryConfig contains the retry configuration for the model.\n\t// This is populated at request time from the agent's ModelRetryConfig.\n\t// Used by EventSenderModelWrapper to wrap stream errors appropriately.\n\tModelRetryConfig *ModelRetryConfig\n}\n\n// ChatModelAgentContext contains runtime information passed to handlers before each ChatModelAgent run.\n// Handlers can modify Instruction, Tools, and ReturnDirectly to customize agent behavior.\n//\n// This type is specific to ChatModelAgent. Other agent types may define their own context types.\ntype ChatModelAgentContext struct {\n\t// Instruction is the current instruction for the Agent execution.\n\t// It includes the instruction configured for the agent, additional instructions appended by framework\n\t// and AgentMiddleware, and modifications applied by previous BeforeAgent handlers.\n\t// The finalized instruction after all BeforeAgent handlers are then passed to GenModelInput,\n\t// to be (optionally) formatted with SessionValues and converted to system message.\n\tInstruction string\n\n\t// Tools are the raw tools (without any wrapper or tool middleware) currently configured for the Agent execution.\n\t// They includes tools passed in AgentConfig, implicit tools added by framework such as transfer / exit tools,\n\t// and other tools already added by middlewares.\n\tTools []tool.BaseTool\n\n\t// ReturnDirectly is the set of tool names currently configured to cause the Agent to return directly.\n\t// This is based on the return directly map configured for the agent, plus any modifications\n\t// by previous BeforeAgent handlers.\n\tReturnDirectly map[string]bool\n}\n\n// ChatModelAgentMiddleware defines the interface for customizing ChatModelAgent behavior.\n//\n// IMPORTANT: This interface is specifically designed for ChatModelAgent and agents built\n// on top of it (e.g., DeepAgent).\n//\n// Why ChatModelAgentMiddleware instead of AgentMiddleware?\n//\n// AgentMiddleware is a struct type, which has inherent limitations:\n//   - Struct types are closed: users cannot add new methods to extend functionality\n//   - The framework only recognizes AgentMiddleware's fixed fields, so even if users\n//     embed AgentMiddleware in a custom struct and add methods, the framework cannot\n//     call those methods (config.Middlewares is []AgentMiddleware, not a user type)\n//   - Callbacks in AgentMiddleware only return error, cannot return modified context\n//\n// ChatModelAgentMiddleware is an interface type, which is open for extension:\n//   - Users can implement custom handlers with arbitrary internal state and methods\n//   - Hook methods return (context.Context, ..., error) for direct context propagation\n//   - Wrapper methods (WrapToolCall, WrapModel) enable context propagation through the\n//     wrapped endpoint chain: wrappers can pass modified context to the next wrapper\n//   - Configuration is centralized in struct fields rather than scattered in closures\n//\n// ChatModelAgentMiddleware vs AgentMiddleware:\n//   - Use AgentMiddleware for simple, static additions (extra instruction/tools)\n//   - Use ChatModelAgentMiddleware for dynamic behavior, context modification, or call wrapping\n//   - AgentMiddleware is kept for backward compatibility with existing users\n//   - Both can be used together; see AgentMiddleware documentation for execution order\n//\n// Use *BaseChatModelAgentMiddleware as an embedded struct to provide default no-op\n// implementations for all methods.\ntype ChatModelAgentMiddleware interface {\n\t// BeforeAgent is called before each agent run, allowing modification of\n\t// the agent's instruction and tools configuration.\n\tBeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error)\n\n\t// BeforeModelRewriteState is called before each model invocation.\n\t// The returned state is persisted to the agent's internal state and passed to the model.\n\t// The returned context is propagated to the model call and subsequent handlers.\n\t//\n\t// The ChatModelAgentState struct provides access to:\n\t//   - Messages: the conversation history\n\t//\n\t// The ModelContext struct provides read-only access to:\n\t//   - Tools: the current tool list that will be sent to the model\n\tBeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)\n\n\t// AfterModelRewriteState is called after each model invocation.\n\t// The input state includes the model's response as the last message.\n\t// The returned state is persisted to the agent's internal state.\n\t//\n\t// The ChatModelAgentState struct provides access to:\n\t//   - Messages: the conversation history including the model's response\n\t//\n\t// The ModelContext struct provides read-only access to:\n\t//   - Tools: the current tool list that was sent to the model\n\tAfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)\n\n\t// WrapInvokableToolCall wraps a tool's synchronous execution with custom behavior.\n\t// Return the input endpoint unchanged and nil error if no wrapping is needed.\n\t//\n\t// This method is only called for tools that implement InvokableTool.\n\t// If a tool only implements StreamableTool, this method will not be called for that tool.\n\t//\n\t// This method is called at request time when the tool is about to be executed.\n\t// The tCtx parameter provides metadata about the tool:\n\t//   - Name: The name of the tool being wrapped\n\t//   - CallID: The unique identifier for this specific tool call\n\tWrapInvokableToolCall(ctx context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error)\n\n\t// WrapStreamableToolCall wraps a tool's streaming execution with custom behavior.\n\t// Return the input endpoint unchanged and nil error if no wrapping is needed.\n\t//\n\t// This method is only called for tools that implement StreamableTool.\n\t// If a tool only implements InvokableTool, this method will not be called for that tool.\n\t//\n\t// This method is called at request time when the tool is about to be executed.\n\t// The tCtx parameter provides metadata about the tool:\n\t//   - Name: The name of the tool being wrapped\n\t//   - CallID: The unique identifier for this specific tool call\n\tWrapStreamableToolCall(ctx context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error)\n\n\t// WrapEnhancedInvokableToolCall wraps an enhanced tool's synchronous execution with custom behavior.\n\t// Return the input endpoint unchanged and nil error if no wrapping is needed.\n\t//\n\t// This method is only called for tools that implement EnhancedInvokableTool.\n\t// If a tool only implements EnhancedStreamableTool, this method will not be called for that tool.\n\t//\n\t// This method is called at request time when the tool is about to be executed.\n\t// The tCtx parameter provides metadata about the tool:\n\t//   - Name: The name of the tool being wrapped\n\t//   - CallID: The unique identifier for this specific tool call\n\tWrapEnhancedInvokableToolCall(ctx context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error)\n\n\t// WrapEnhancedStreamableToolCall wraps an enhanced tool's streaming execution with custom behavior.\n\t// Return the input endpoint unchanged and nil error if no wrapping is needed.\n\t//\n\t// This method is only called for tools that implement EnhancedStreamableTool.\n\t// If a tool only implements EnhancedInvokableTool, this method will not be called for that tool.\n\t//\n\t// This method is called at request time when the tool is about to be executed.\n\t// The tCtx parameter provides metadata about the tool:\n\t//   - Name: The name of the tool being wrapped\n\t//   - CallID: The unique identifier for this specific tool call\n\tWrapEnhancedStreamableToolCall(ctx context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error)\n\n\t// WrapModel wraps a chat model with custom behavior.\n\t// Return the input model unchanged and nil error if no wrapping is needed.\n\t//\n\t// This method is called at request time when the model is about to be invoked.\n\t// Note: The parameter is BaseChatModel (not ToolCallingChatModel) because wrappers\n\t// only need to intercept Generate/Stream calls. Tool binding (WithTools) is handled\n\t// separately by the framework and does not flow through user wrappers.\n\t//\n\t// The mc parameter contains the current tool configuration:\n\t//   - Tools: The tool infos that will be sent to the model\n\tWrapModel(ctx context.Context, m model.BaseChatModel, mc *ModelContext) (model.BaseChatModel, error)\n}\n\n// BaseChatModelAgentMiddleware provides default no-op implementations for ChatModelAgentMiddleware.\n// Embed *BaseChatModelAgentMiddleware in custom handlers to only override the methods you need.\n//\n// Example:\n//\n//\ttype MyHandler struct {\n//\t\t*adk.BaseChatModelAgentMiddleware\n//\t\t// custom fields\n//\t}\n//\n//\tfunc (h *MyHandler) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {\n//\t\t// custom logic\n//\t\treturn ctx, state, nil\n//\t}\ntype BaseChatModelAgentMiddleware struct{}\n\nfunc (b *BaseChatModelAgentMiddleware) WrapInvokableToolCall(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) (InvokableToolCallEndpoint, error) {\n\treturn endpoint, nil\n}\n\nfunc (b *BaseChatModelAgentMiddleware) WrapStreamableToolCall(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) (StreamableToolCallEndpoint, error) {\n\treturn endpoint, nil\n}\n\nfunc (b *BaseChatModelAgentMiddleware) WrapEnhancedInvokableToolCall(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) (EnhancedInvokableToolCallEndpoint, error) {\n\treturn endpoint, nil\n}\n\nfunc (b *BaseChatModelAgentMiddleware) WrapEnhancedStreamableToolCall(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) (EnhancedStreamableToolCallEndpoint, error) {\n\treturn endpoint, nil\n}\n\nfunc (b *BaseChatModelAgentMiddleware) WrapModel(_ context.Context, m model.BaseChatModel, _ *ModelContext) (model.BaseChatModel, error) {\n\treturn m, nil\n}\n\nfunc (b *BaseChatModelAgentMiddleware) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\treturn ctx, runCtx, nil\n}\n\nfunc (b *BaseChatModelAgentMiddleware) BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\treturn ctx, state, nil\n}\n\nfunc (b *BaseChatModelAgentMiddleware) AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\treturn ctx, state, nil\n}\n\n// SetRunLocalValue sets a key-value pair that persists for the duration of the current agent Run() invocation.\n// The value is scoped to this specific execution and is not shared across different Run() calls or agent instances.\n//\n// Values stored here are compatible with interrupt/resume cycles - they will be serialized and restored\n// when the agent is resumed. For custom types, you must register them using schema.RegisterName[T]()\n// in an init() function to ensure proper serialization.\n//\n// This function can only be called from within a ChatModelAgentMiddleware during agent execution.\n// Returns an error if called outside of an agent execution context.\nfunc SetRunLocalValue(ctx context.Context, key string, value any) error {\n\terr := compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tif st.Extra == nil {\n\t\t\tst.Extra = make(map[string]any)\n\t\t}\n\t\tst.Extra[key] = value\n\t\treturn nil\n\t})\n\tif err != nil {\n\t\treturn fmt.Errorf(\"SetRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w\", err)\n\t}\n\treturn nil\n}\n\n// GetRunLocalValue retrieves a value that was set during the current agent Run() invocation.\n// The value is scoped to this specific execution and is not shared across different Run() calls or agent instances.\n//\n// Values stored via SetRunLocalValue are compatible with interrupt/resume cycles - they will be serialized\n// and restored when the agent is resumed. For custom types, you must register them using schema.RegisterName[T]()\n// in an init() function to ensure proper serialization.\n//\n// This function can only be called from within a ChatModelAgentMiddleware during agent execution.\n// Returns the value and true if found, or nil and false if not found or if called outside of an agent execution context.\nfunc GetRunLocalValue(ctx context.Context, key string) (any, bool, error) {\n\tvar val any\n\tvar found bool\n\terr := compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tif st.Extra != nil {\n\t\t\tval, found = st.Extra[key]\n\t\t}\n\t\treturn nil\n\t})\n\tif err != nil {\n\t\treturn nil, false, fmt.Errorf(\"GetRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w\", err)\n\t}\n\treturn val, found, nil\n}\n\n// DeleteRunLocalValue removes a value that was set during the current agent Run() invocation.\n//\n// This function can only be called from within a ChatModelAgentMiddleware during agent execution.\n// Returns an error if called outside of an agent execution context.\nfunc DeleteRunLocalValue(ctx context.Context, key string) error {\n\terr := compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tif st.Extra != nil {\n\t\t\tdelete(st.Extra, key)\n\t\t}\n\t\treturn nil\n\t})\n\tif err != nil {\n\t\treturn fmt.Errorf(\"DeleteRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w\", err)\n\t}\n\treturn nil\n}\n\n// SendEvent sends a custom AgentEvent to the event stream during agent execution.\n// This allows ChatModelAgentMiddleware implementations to emit custom events that will be\n// received by the caller iterating over the agent's event stream.\n//\n// This function can only be called from within a ChatModelAgentMiddleware during agent execution.\n// Returns an error if called outside of an agent execution context.\nfunc SendEvent(ctx context.Context, event *AgentEvent) error {\n\texecCtx := getChatModelAgentExecCtx(ctx)\n\tif execCtx == nil || execCtx.generator == nil {\n\t\treturn fmt.Errorf(\"SendEvent failed: must be called within a ChatModelAgent Run() or Resume() execution context\")\n\t}\n\texecCtx.generator.Send(event)\n\treturn nil\n}\n"
  },
  {
    "path": "adk/handler_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype testInstructionHandler struct {\n\t*BaseChatModelAgentMiddleware\n\ttext string\n}\n\nfunc (h *testInstructionHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\tif runCtx.Instruction == \"\" {\n\t\trunCtx.Instruction = h.text\n\t} else if h.text != \"\" {\n\t\trunCtx.Instruction = runCtx.Instruction + \"\\n\" + h.text\n\t}\n\treturn ctx, runCtx, nil\n}\n\ntype testInstructionFuncHandler struct {\n\t*BaseChatModelAgentMiddleware\n\tfn func(ctx context.Context, instruction string) (context.Context, string, error)\n}\n\nfunc (h *testInstructionFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\tnewCtx, newInstruction, err := h.fn(ctx, runCtx.Instruction)\n\tif err != nil {\n\t\treturn ctx, runCtx, err\n\t}\n\trunCtx.Instruction = newInstruction\n\treturn newCtx, runCtx, nil\n}\n\ntype testToolsHandler struct {\n\t*BaseChatModelAgentMiddleware\n\ttools []tool.BaseTool\n}\n\nfunc (h *testToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\trunCtx.Tools = append(runCtx.Tools, h.tools...)\n\treturn ctx, runCtx, nil\n}\n\ntype testToolsFuncHandler struct {\n\t*BaseChatModelAgentMiddleware\n\tfn func(ctx context.Context, tools []tool.BaseTool, returnDirectly map[string]bool) (context.Context, []tool.BaseTool, map[string]bool, error)\n}\n\nfunc (h *testToolsFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\tnewCtx, newTools, newReturnDirectly, err := h.fn(ctx, runCtx.Tools, runCtx.ReturnDirectly)\n\tif err != nil {\n\t\treturn ctx, runCtx, err\n\t}\n\trunCtx.Tools = newTools\n\trunCtx.ReturnDirectly = newReturnDirectly\n\treturn newCtx, runCtx, nil\n}\n\ntype testBeforeAgentHandler struct {\n\t*BaseChatModelAgentMiddleware\n\tfn func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error)\n}\n\nfunc (h *testBeforeAgentHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\treturn h.fn(ctx, runCtx)\n}\n\ntype testBeforeModelRewriteStateHandler struct {\n\t*BaseChatModelAgentMiddleware\n\tfn func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)\n}\n\nfunc (h *testBeforeModelRewriteStateHandler) BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\treturn h.fn(ctx, state, mc)\n}\n\ntype testAfterModelRewriteStateHandler struct {\n\t*BaseChatModelAgentMiddleware\n\tfn func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)\n}\n\nfunc (h *testAfterModelRewriteStateHandler) AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\treturn h.fn(ctx, state, mc)\n}\n\ntype testToolWrapperHandler struct {\n\t*BaseChatModelAgentMiddleware\n\twrapInvokableFn  func(context.Context, InvokableToolCallEndpoint, *ToolContext) InvokableToolCallEndpoint\n\twrapStreamableFn func(context.Context, StreamableToolCallEndpoint, *ToolContext) StreamableToolCallEndpoint\n}\n\nfunc (h *testToolWrapperHandler) WrapInvokableToolCall(ctx context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error) {\n\tif h.wrapInvokableFn != nil {\n\t\treturn h.wrapInvokableFn(ctx, endpoint, tCtx), nil\n\t}\n\treturn endpoint, nil\n}\n\nfunc (h *testToolWrapperHandler) WrapStreamableToolCall(ctx context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error) {\n\tif h.wrapStreamableFn != nil {\n\t\treturn h.wrapStreamableFn(ctx, endpoint, tCtx), nil\n\t}\n\treturn endpoint, nil\n}\n\ntype testModelWrapperHandler struct {\n\t*BaseChatModelAgentMiddleware\n\tfn func(context.Context, model.BaseChatModel, *ModelContext) model.BaseChatModel\n}\n\nfunc (h *testModelWrapperHandler) WrapModel(ctx context.Context, m model.BaseChatModel, mc *ModelContext) (model.BaseChatModel, error) {\n\treturn h.fn(ctx, m, mc), nil\n}\n\nfunc newTestInvokableToolCallWrapper(beforeFn, afterFn func()) func(context.Context, InvokableToolCallEndpoint, *ToolContext) InvokableToolCallEndpoint {\n\treturn func(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) InvokableToolCallEndpoint {\n\t\treturn func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\t\t\tif beforeFn != nil {\n\t\t\t\tbeforeFn()\n\t\t\t}\n\t\t\tresult, err := endpoint(ctx, argumentsInJSON, opts...)\n\t\t\tif afterFn != nil {\n\t\t\t\tafterFn()\n\t\t\t}\n\t\t\treturn result, err\n\t\t}\n\t}\n}\n\nfunc newResultModifyingInvokableToolCallWrapper(modifyFn func(string) string) func(context.Context, InvokableToolCallEndpoint, *ToolContext) InvokableToolCallEndpoint {\n\treturn func(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) InvokableToolCallEndpoint {\n\t\treturn func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\t\t\tresult, err := endpoint(ctx, argumentsInJSON, opts...)\n\t\t\tif err == nil && modifyFn != nil {\n\t\t\t\tresult = modifyFn(result)\n\t\t\t}\n\t\t\treturn result, err\n\t\t}\n\t}\n}\n\nfunc newTestStreamableToolCallWrapper(beforeFn, afterFn func()) func(context.Context, StreamableToolCallEndpoint, *ToolContext) StreamableToolCallEndpoint {\n\treturn func(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) StreamableToolCallEndpoint {\n\t\treturn func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\t\t\tif beforeFn != nil {\n\t\t\t\tbeforeFn()\n\t\t\t}\n\t\t\tresult, err := endpoint(ctx, argumentsInJSON, opts...)\n\t\t\tif afterFn != nil {\n\t\t\t\tafterFn()\n\t\t\t}\n\t\t\treturn result, err\n\t\t}\n\t}\n}\n\nfunc TestHandlerExecutionOrder(t *testing.T) {\n\tt.Run(\"MultipleInstructionHandlersPipeline\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar capturedInstruction string\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) {\n\t\t\t\tif len(msgs) > 0 && msgs[0].Role == schema.System {\n\t\t\t\t\tcapturedInstruction = msgs[0].Content\n\t\t\t\t}\n\t\t\t\treturn schema.AssistantMessage(\"response\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tInstruction: \"Base instruction.\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testInstructionHandler{text: \"Handler 1 addition.\"},\n\t\t\t\t&testInstructionHandler{text: \"Handler 2 addition.\"},\n\t\t\t\t&testInstructionFuncHandler{fn: func(ctx context.Context, instruction string) (context.Context, string, error) {\n\t\t\t\t\treturn ctx, instruction + \"\\nHandler 3 dynamic.\", nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Contains(t, capturedInstruction, \"Base instruction.\")\n\t\tassert.Contains(t, capturedInstruction, \"Handler 1 addition.\")\n\t\tassert.Contains(t, capturedInstruction, \"Handler 2 addition.\")\n\t\tassert.Contains(t, capturedInstruction, \"Handler 3 dynamic.\")\n\t})\n\n\tt.Run(\"MiddlewaresBeforeHandlers\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar capturedInstruction string\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) {\n\t\t\t\tif len(msgs) > 0 && msgs[0].Role == schema.System {\n\t\t\t\t\tcapturedInstruction = msgs[0].Content\n\t\t\t\t}\n\t\t\t\treturn schema.AssistantMessage(\"response\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tInstruction: \"Base.\",\n\t\t\tModel:       cm,\n\t\t\tMiddlewares: []AgentMiddleware{\n\t\t\t\t{AdditionalInstruction: \"Middleware instruction.\"},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testInstructionHandler{text: \"Handler instruction.\"},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tmiddlewareIdx := len(capturedInstruction) - len(\"Middleware instruction.\") - len(\"\\nHandler instruction.\")\n\t\thandlerIdx := len(capturedInstruction) - len(\"Handler instruction.\")\n\t\tassert.True(t, middlewareIdx < handlerIdx, \"Middleware should be applied before Handler\")\n\t})\n}\n\nfunc TestToolsHandlerCombinations(t *testing.T) {\n\tt.Run(\"MultipleToolsHandlersAppend\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttool1 := &fakeToolForTest{tarCount: 1}\n\t\ttool2 := &fakeToolForTest{tarCount: 2}\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tvar capturedToolCount int\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\t\tfunc(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\toptions := model.GetCommonOptions(&model.Options{}, opts...)\n\t\t\t\tcapturedToolCount = len(options.Tools)\n\t\t\t\treturn schema.AssistantMessage(\"response\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{tool1},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testToolsHandler{tools: []tool.BaseTool{tool2}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 2, capturedToolCount)\n\t})\n\n\tt.Run(\"ToolsFuncCanRemoveTools\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttool1 := &namedTool{name: \"tool1\"}\n\t\ttool2 := &namedTool{name: \"tool2\"}\n\t\ttool3 := &namedTool{name: \"tool3\"}\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tvar capturedToolNames []string\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\t\tfunc(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\toptions := model.GetCommonOptions(&model.Options{}, opts...)\n\t\t\t\tfor _, t := range options.Tools {\n\t\t\t\t\tcapturedToolNames = append(capturedToolNames, t.Name)\n\t\t\t\t}\n\t\t\t\treturn schema.AssistantMessage(\"response\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{tool1, tool2, tool3},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testToolsFuncHandler{fn: func(ctx context.Context, tools []tool.BaseTool, returnDirectly map[string]bool) (context.Context, []tool.BaseTool, map[string]bool, error) {\n\t\t\t\t\tfiltered := make([]tool.BaseTool, 0)\n\t\t\t\t\tfor _, t := range tools {\n\t\t\t\t\t\tinfo, _ := t.Info(ctx)\n\t\t\t\t\t\tif info.Name != \"tool2\" {\n\t\t\t\t\t\t\tfiltered = append(filtered, t)\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\treturn ctx, filtered, returnDirectly, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Contains(t, capturedToolNames, \"tool1\")\n\t\tassert.NotContains(t, capturedToolNames, \"tool2\")\n\t\tassert.Contains(t, capturedToolNames, \"tool3\")\n\t})\n\n\tt.Run(\"ReturnDirectlyModification\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttool1 := &namedTool{name: \"tool1\"}\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Using tool\", []schema.ToolCall{\n\t\t\t\t{ID: \"call1\", Function: schema.FunctionCall{Name: \"tool1\", Arguments: \"{}\"}},\n\t\t\t}), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{tool1},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testToolsFuncHandler{fn: func(ctx context.Context, tools []tool.BaseTool, returnDirectly map[string]bool) (context.Context, []tool.BaseTool, map[string]bool, error) {\n\t\t\t\t\tfor _, t := range tools {\n\t\t\t\t\t\tinfo, _ := t.Info(ctx)\n\t\t\t\t\t\tif info.Name == \"tool1\" {\n\t\t\t\t\t\t\treturnDirectly[info.Name] = true\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\treturn ctx, tools, returnDirectly, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\teventCount := 0\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\teventCount++\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.Message != nil &&\n\t\t\t\tevent.Output.MessageOutput.Message.Role == schema.Tool {\n\t\t\t\tassert.Equal(t, \"tool1 result\", event.Output.MessageOutput.Message.Content)\n\t\t\t}\n\t\t}\n\t\tassert.Equal(t, 2, eventCount)\n\t})\n\n\tt.Run(\"DynamicToolCanBeCalledByModel\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tdynamicToolCalled := false\n\t\tdynamicTool := &callableTool{\n\t\t\tname: \"dynamic_tool\",\n\t\t\tinvokeFn: func() {\n\t\t\t\tdynamicToolCalled = true\n\t\t\t},\n\t\t}\n\t\tinfo, _ := dynamicTool.Info(ctx)\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Using dynamic tool\", []schema.ToolCall{\n\t\t\t\t{ID: \"call1\", Function: schema.FunctionCall{Name: info.Name, Arguments: \"{}\"}},\n\t\t\t}), nil).Times(1)\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testToolsHandler{tools: []tool.BaseTool{dynamicTool}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.True(t, dynamicToolCalled, \"Dynamic tool should have been called\")\n\t})\n}\n\nfunc TestMessageRewriteHandlers(t *testing.T) {\n\tt.Run(\"BeforeModelRewriteStatePipeline\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar capturedMsgCount int\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) {\n\t\t\t\tcapturedMsgCount = len(msgs)\n\t\t\t\treturn schema.AssistantMessage(\"response\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tInstruction: \"instruction\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\tstate.Messages = append(state.Messages, schema.UserMessage(\"injected1\"))\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\tstate.Messages = append(state.Messages, schema.UserMessage(\"injected2\"))\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"original\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 4, capturedMsgCount)\n\t})\n\n\tt.Run(\"AfterModelRewriteState\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tafterCalled := false\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testAfterModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\tafterCalled = true\n\t\t\t\t\tassert.True(t, len(state.Messages) > 0)\n\t\t\t\t\tlastMsg := state.Messages[len(state.Messages)-1]\n\t\t\t\t\tassert.Equal(t, schema.Assistant, lastMsg.Role)\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.True(t, afterCalled)\n\t})\n}\n\nfunc TestToolCallWrapperHandlers(t *testing.T) {\n\tt.Run(\"MultipleToolWrappersPipeline\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttestTool := &namedTool{name: \"test_tool\"}\n\t\tinfo, _ := testTool.Info(ctx)\n\n\t\tvar callOrder []string\n\t\tvar mu sync.Mutex\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Using tool\", []schema.ToolCall{\n\t\t\t\t{ID: \"call1\", Function: schema.FunctionCall{Name: info.Name, Arguments: \"{}\"}},\n\t\t\t}), nil).Times(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testToolWrapperHandler{wrapInvokableFn: newTestInvokableToolCallWrapper(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper1-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper1-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t)},\n\t\t\t\t&testToolWrapperHandler{wrapInvokableFn: newTestInvokableToolCallWrapper(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper2-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper2-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t)},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, []string{\"wrapper2-before\", \"wrapper1-before\", \"wrapper1-after\", \"wrapper2-after\"}, callOrder)\n\t})\n\n\tt.Run(\"StreamingToolWrappersPipeline\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttestTool := &streamingNamedTool{name: \"streaming_tool\"}\n\t\tinfo, _ := testTool.Info(ctx)\n\n\t\tvar callOrder []string\n\t\tvar mu sync.Mutex\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"Using tool\", []schema.ToolCall{\n\t\t\t\t\t{ID: \"call1\", Function: schema.FunctionCall{Name: info.Name, Arguments: \"{}\"}},\n\t\t\t\t}),\n\t\t\t}), nil).Times(1)\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"done\", nil),\n\t\t\t}), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testToolWrapperHandler{wrapStreamableFn: newTestStreamableToolCallWrapper(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper1-stream-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper1-stream-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t)},\n\t\t\t\t&testToolWrapperHandler{wrapStreamableFn: newTestStreamableToolCallWrapper(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper2-stream-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper2-stream-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t)},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true, CheckPointStore: newBridgeStore()})\n\t\titer := r.Run(ctx, []Message{schema.UserMessage(\"test\")})\n\n\t\tvar hasStreamingToolResult bool\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.IsStreaming &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Tool {\n\t\t\t\thasStreamingToolResult = true\n\t\t\t\tfor {\n\t\t\t\t\t_, err := event.Output.MessageOutput.MessageStream.Recv()\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tassert.True(t, hasStreamingToolResult, \"Should have streaming tool result\")\n\t\tassert.Equal(t, []string{\"wrapper2-stream-before\", \"wrapper1-stream-before\", \"wrapper1-stream-after\", \"wrapper2-stream-after\"}, callOrder,\n\t\t\t\"Streaming wrappers should be called in correct order\")\n\t})\n\n\tt.Run(\"ToolWrapperCanModifyResult\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttestTool := &namedTool{name: \"test_tool\"}\n\t\tinfo, _ := testTool.Info(ctx)\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Using tool\", []schema.ToolCall{\n\t\t\t\t{ID: \"call1\", Function: schema.FunctionCall{Name: info.Name, Arguments: \"{}\"}},\n\t\t\t}), nil).Times(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testToolWrapperHandler{wrapInvokableFn: newResultModifyingInvokableToolCallWrapper(func(result string) string {\n\t\t\t\t\treturn \"modified: \" + result\n\t\t\t\t})},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.Message != nil &&\n\t\t\t\tevent.Output.MessageOutput.Message.Role == schema.Tool {\n\t\t\t\tassert.Equal(t, \"modified: test_tool result\", event.Output.MessageOutput.Message.Content)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestToolContextFunctions(t *testing.T) {\n\tt.Run(\"ModelContextToolsField\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttestTool := &namedTool{name: \"base_tool\"}\n\t\tinfo, _ := testTool.Info(ctx)\n\n\t\tvar wrapperSeenTools []*schema.ToolInfo\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testModelWrapperHandler{\n\t\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\t\tfn: func(_ context.Context, m model.BaseChatModel, mc *ModelContext) model.BaseChatModel {\n\t\t\t\t\t\treturn &toolChainingTestModel{\n\t\t\t\t\t\t\tinner: m,\n\t\t\t\t\t\t\tmc:    mc,\n\t\t\t\t\t\t\twrapFn: func(ctx context.Context, opts []model.Option) []model.Option {\n\t\t\t\t\t\t\t\twrapperSeenTools = mc.Tools\n\t\t\t\t\t\t\t\treturn opts\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Len(t, wrapperSeenTools, 1, \"Wrapper should see base tool\")\n\t\tassert.Equal(t, info.Name, wrapperSeenTools[0].Name, \"Wrapper should see base_tool\")\n\t})\n}\n\ntype toolChainingTestModel struct {\n\tinner  model.BaseChatModel\n\tmc     *ModelContext\n\twrapFn func(ctx context.Context, opts []model.Option) []model.Option\n}\n\nfunc (m *toolChainingTestModel) Generate(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tif m.wrapFn != nil {\n\t\topts = m.wrapFn(ctx, opts)\n\t}\n\treturn m.inner.Generate(ctx, msgs, opts...)\n}\n\nfunc (m *toolChainingTestModel) Stream(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tif m.wrapFn != nil {\n\t\topts = m.wrapFn(ctx, opts)\n\t}\n\treturn m.inner.Stream(ctx, msgs, opts...)\n}\n\nfunc (m *toolChainingTestModel) BindTools(tools []*schema.ToolInfo) error {\n\treturn nil\n}\n\nfunc TestContextPropagation(t *testing.T) {\n\tt.Run(\"ContextPassedThroughBeforeModelHandlers\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttype ctxKey string\n\t\tconst key1 ctxKey = \"key1\"\n\t\tconst key2 ctxKey = \"key2\"\n\n\t\tvar handler2ReceivedValue1 interface{}\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\treturn context.WithValue(ctx, key1, \"value1\"), state, nil\n\t\t\t\t}},\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\thandler2ReceivedValue1 = ctx.Value(key1)\n\t\t\t\t\treturn context.WithValue(ctx, key2, \"value2\"), state, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, \"value1\", handler2ReceivedValue1, \"Handler 2 should receive context value set by Handler 1\")\n\t})\n\n\tt.Run(\"BeforeAgentContextPropagation\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttype ctxKey string\n\t\tconst key1 ctxKey = \"key1\"\n\n\t\tvar handler2ReceivedValue interface{}\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\t\t\t\t\treturn context.WithValue(ctx, key1, \"value1\"), runCtx, nil\n\t\t\t\t}},\n\t\t\t\t&testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\t\t\t\t\thandler2ReceivedValue = ctx.Value(key1)\n\t\t\t\t\treturn ctx, runCtx, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, \"value1\", handler2ReceivedValue, \"Handler 2 should receive context value set by Handler 1 during BeforeAgent\")\n\t})\n}\n\nfunc TestCustomHandler(t *testing.T) {\n\tt.Run(\"CustomHandlerWithState\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tcustomHandler := &countingHandler{}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers:    []ChatModelAgentMiddleware{customHandler},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 1, customHandler.beforeAgentCount)\n\t\tassert.Equal(t, 1, customHandler.beforeModelCount)\n\t\tassert.Equal(t, 1, customHandler.afterModelCount)\n\t})\n}\n\nfunc TestHandlerErrorHandling(t *testing.T) {\n\tt.Run(\"BeforeAgentErrorStopsRun\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\t\t\t\t\treturn ctx, runCtx, assert.AnError\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{\n\t\t\tMessages: []*schema.Message{schema.UserMessage(\"test\")},\n\t\t})\n\n\t\tvar gotErr error\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Err != nil {\n\t\t\t\tgotErr = event.Err\n\t\t\t}\n\t\t}\n\n\t\tassert.Error(t, gotErr)\n\t\tassert.Contains(t, gotErr.Error(), \"BeforeAgent failed\")\n\t})\n}\n\ntype namedTool struct {\n\tname string\n}\n\nfunc (t *namedTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: t.name, Desc: t.name + \" description\"}, nil\n}\n\nfunc (t *namedTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn t.name + \" result\", nil\n}\n\ntype streamingNamedTool struct {\n\tname string\n}\n\nfunc (t *streamingNamedTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: t.name, Desc: t.name + \" description\"}, nil\n}\n\nfunc (t *streamingNamedTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn t.name + \" result\", nil\n}\n\nfunc (t *streamingNamedTool) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) {\n\treturn schema.StreamReaderFromArray([]string{t.name + \" stream result\"}), nil\n}\n\ntype callableTool struct {\n\tname     string\n\tinvokeFn func()\n}\n\nfunc (t *callableTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: t.name, Desc: t.name + \" description\"}, nil\n}\n\nfunc (t *callableTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\tif t.invokeFn != nil {\n\t\tt.invokeFn()\n\t}\n\treturn t.name + \" result\", nil\n}\n\ntype countingHandler struct {\n\t*BaseChatModelAgentMiddleware\n\tbeforeAgentCount int\n\tbeforeModelCount int\n\tafterModelCount  int\n\tmu               sync.Mutex\n}\n\nfunc (h *countingHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\th.mu.Lock()\n\th.beforeAgentCount++\n\th.mu.Unlock()\n\treturn ctx, runCtx, nil\n}\n\nfunc (h *countingHandler) BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\th.mu.Lock()\n\th.beforeModelCount++\n\th.mu.Unlock()\n\treturn ctx, state, nil\n}\n\nfunc (h *countingHandler) AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\th.mu.Lock()\n\th.afterModelCount++\n\th.mu.Unlock()\n\treturn ctx, state, nil\n}\n\nfunc newTestModelWrapperFn(beforeFn, afterFn func()) func(context.Context, model.BaseChatModel, *ModelContext) model.BaseChatModel {\n\treturn func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel {\n\t\treturn &testWrappedModel{\n\t\t\tinner:    m,\n\t\t\tbeforeFn: beforeFn,\n\t\t\tafterFn:  afterFn,\n\t\t}\n\t}\n}\n\ntype testWrappedModel struct {\n\tinner    model.BaseChatModel\n\tbeforeFn func()\n\tafterFn  func()\n}\n\nfunc (m *testWrappedModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tif m.beforeFn != nil {\n\t\tm.beforeFn()\n\t}\n\tresult, err := m.inner.Generate(ctx, input, opts...)\n\tif m.afterFn != nil {\n\t\tm.afterFn()\n\t}\n\treturn result, err\n}\n\nfunc (m *testWrappedModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tif m.beforeFn != nil {\n\t\tm.beforeFn()\n\t}\n\tresult, err := m.inner.Stream(ctx, input, opts...)\n\tif m.afterFn != nil {\n\t\tm.afterFn()\n\t}\n\treturn result, err\n}\n\nfunc TestModelWrapperHandlers(t *testing.T) {\n\tt.Run(\"MultipleModelWrappersPipeline\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar callOrder []string\n\t\tvar mu sync.Mutex\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testModelWrapperHandler{fn: newTestModelWrapperFn(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper1-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper1-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t)},\n\t\t\t\t&testModelWrapperHandler{fn: newTestModelWrapperFn(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper2-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper2-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t)},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, []string{\"wrapper1-before\", \"wrapper2-before\", \"wrapper2-after\", \"wrapper1-after\"}, callOrder)\n\t})\n\n\tt.Run(\"ModelWrapperBeforeAfterCallOrder\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar callOrder []string\n\t\tvar mu sync.Mutex\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\tmu.Lock()\n\t\t\t\tcallOrder = append(callOrder, \"model-generate\")\n\t\t\t\tmu.Unlock()\n\t\t\t\treturn schema.AssistantMessage(\"original response\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testModelWrapperHandler{fn: newTestModelWrapperFn(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t)},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, []string{\"wrapper-before\", \"model-generate\", \"wrapper-after\"}, callOrder)\n\t})\n\n\tt.Run(\"ModelWrapperWithTools\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttestTool := &namedTool{name: \"test_tool\"}\n\t\tinfo, _ := testTool.Info(ctx)\n\n\t\tvar callOrder []string\n\t\tvar mu sync.Mutex\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\tmu.Lock()\n\t\t\t\tcallOrder = append(callOrder, \"model-call\")\n\t\t\t\tmu.Unlock()\n\t\t\t\treturn schema.AssistantMessage(\"Using tool\", []schema.ToolCall{\n\t\t\t\t\t{ID: \"call1\", Function: schema.FunctionCall{Name: info.Name, Arguments: \"{}\"}},\n\t\t\t\t}), nil\n\t\t\t}).Times(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\tmu.Lock()\n\t\t\t\tcallOrder = append(callOrder, \"model-call\")\n\t\t\t\tmu.Unlock()\n\t\t\t\treturn schema.AssistantMessage(\"done\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testModelWrapperHandler{fn: newTestModelWrapperFn(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tcallOrder = append(callOrder, \"wrapper-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t)},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, []string{\n\t\t\t\"wrapper-before\", \"model-call\", \"wrapper-after\",\n\t\t\t\"wrapper-before\", \"model-call\", \"wrapper-after\",\n\t\t}, callOrder)\n\t})\n}\n\ntype simpleChatModelWithoutCallbacks struct {\n\tgenerateFn func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error)\n\tstreamFn   func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error)\n}\n\nfunc (m *simpleChatModelWithoutCallbacks) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tif m.generateFn != nil {\n\t\treturn m.generateFn(ctx, input, opts...)\n\t}\n\treturn schema.AssistantMessage(\"default response\", nil), nil\n}\n\nfunc (m *simpleChatModelWithoutCallbacks) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tif m.streamFn != nil {\n\t\treturn m.streamFn(ctx, input, opts...)\n\t}\n\treturn schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(\"default response\", nil)}), nil\n}\n\nfunc (m *simpleChatModelWithoutCallbacks) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\treturn m, nil\n}\n\nfunc newInputModifyingWrapperFn(inputPrefix string) func(context.Context, model.BaseChatModel, *ModelContext) model.BaseChatModel {\n\treturn func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel {\n\t\treturn &inputOutputModifyingModel{\n\t\t\tinner:       m,\n\t\t\tinputPrefix: inputPrefix,\n\t\t}\n\t}\n}\n\ntype inputOutputModifyingModel struct {\n\tinner       model.BaseChatModel\n\tinputPrefix string\n}\n\nfunc (m *inputOutputModifyingModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tmodifiedMessages := make([]*schema.Message, len(input))\n\tfor i, msg := range input {\n\t\tif msg.Role == schema.User {\n\t\t\tmodifiedMessages[i] = schema.UserMessage(m.inputPrefix + msg.Content)\n\t\t} else {\n\t\t\tmodifiedMessages[i] = msg\n\t\t}\n\t}\n\treturn m.inner.Generate(ctx, modifiedMessages, opts...)\n}\n\nfunc (m *inputOutputModifyingModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tmodifiedMessages := make([]*schema.Message, len(input))\n\tfor i, msg := range input {\n\t\tif msg.Role == schema.User {\n\t\t\tmodifiedMessages[i] = schema.UserMessage(m.inputPrefix + msg.Content)\n\t\t} else {\n\t\t\tmodifiedMessages[i] = msg\n\t\t}\n\t}\n\treturn m.inner.Stream(ctx, modifiedMessages, opts...)\n}\n\nfunc TestModelWrapper_InputModification(t *testing.T) {\n\tt.Run(\"ModelWrapperModifiesInput_Generate\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar modelReceivedInput []*schema.Message\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\tmodelReceivedInput = input\n\t\t\t\treturn schema.AssistantMessage(\"original response\", nil), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testModelWrapperHandler{fn: newInputModifyingWrapperFn(\"[WRAPPER]\")},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test input\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.NotNil(t, modelReceivedInput)\n\t\tassert.True(t, len(modelReceivedInput) > 0)\n\t\tfound := false\n\t\tfor _, msg := range modelReceivedInput {\n\t\t\tif msg.Content == \"[WRAPPER]test input\" {\n\t\t\t\tfound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, found, \"Model should receive wrapper-modified input\")\n\t})\n\n\tt.Run(\"ModelWrapperModifiesInput_Stream\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar modelReceivedInput []*schema.Message\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\t\t\t\tmodelReceivedInput = input\n\t\t\t\treturn schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\t\tschema.AssistantMessage(\"chunk1\", nil),\n\t\t\t\t\tschema.AssistantMessage(\"chunk2\", nil),\n\t\t\t\t}), nil\n\t\t\t}).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testModelWrapperHandler{fn: newInputModifyingWrapperFn(\"[WRAPPER]\")},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true, CheckPointStore: newBridgeStore()})\n\t\titer := r.Run(ctx, []Message{schema.UserMessage(\"test input\")})\n\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.IsStreaming &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Assistant {\n\t\t\t\tfor {\n\t\t\t\t\t_, err := event.Output.MessageOutput.MessageStream.Recv()\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tassert.NotNil(t, modelReceivedInput)\n\t\tassert.True(t, len(modelReceivedInput) > 0)\n\t\tfound := false\n\t\tfor _, msg := range modelReceivedInput {\n\t\t\tif msg.Content == \"[WRAPPER]test input\" {\n\t\t\t\tfound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, found, \"Model should receive wrapper-modified input\")\n\t})\n}\n\nfunc TestRunLocalValueFunctions(t *testing.T) {\n\tt.Run(\"SetAndGetRunLocalValue\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar capturedValue any\n\t\tvar capturedFound bool\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\terr := SetRunLocalValue(ctx, \"test_key\", \"test_value\")\n\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t\t&testAfterModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\tval, found, err := GetRunLocalValue(ctx, \"test_key\")\n\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\tcapturedValue = val\n\t\t\t\t\tcapturedFound = found\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.True(t, capturedFound, \"Value should be found\")\n\t\tassert.Equal(t, \"test_value\", capturedValue, \"Value should match what was set\")\n\t})\n\n\tt.Run(\"DeleteRunLocalValue\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar valueAfterDelete any\n\t\tvar foundAfterDelete bool\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\terr := SetRunLocalValue(ctx, \"delete_key\", \"delete_value\")\n\t\t\t\t\tassert.NoError(t, err)\n\n\t\t\t\t\terr = DeleteRunLocalValue(ctx, \"delete_key\")\n\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t\t&testAfterModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\tval, found, err := GetRunLocalValue(ctx, \"delete_key\")\n\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\tvalueAfterDelete = val\n\t\t\t\t\tfoundAfterDelete = found\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.False(t, foundAfterDelete, \"Value should not be found after deletion\")\n\t\tassert.Nil(t, valueAfterDelete, \"Value should be nil after deletion\")\n\t})\n\n\tt.Run(\"GetNonExistentKey\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tvar capturedValue any\n\t\tvar capturedFound bool\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\tval, found, err := GetRunLocalValue(ctx, \"non_existent_key\")\n\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\tcapturedValue = val\n\t\t\t\t\tcapturedFound = found\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.False(t, capturedFound, \"Non-existent key should not be found\")\n\t\tassert.Nil(t, capturedValue, \"Non-existent key should return nil value\")\n\t})\n\n\tt.Run(\"RunLocalValueOutsideContext\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\terr := SetRunLocalValue(ctx, \"key\", \"value\")\n\t\tassert.Error(t, err, \"SetRunLocalValue should fail outside agent context\")\n\t\tassert.Contains(t, err.Error(), \"SetRunLocalValue failed\")\n\n\t\t_, _, err = GetRunLocalValue(ctx, \"key\")\n\t\tassert.Error(t, err, \"GetRunLocalValue should fail outside agent context\")\n\t\tassert.Contains(t, err.Error(), \"GetRunLocalValue failed\")\n\n\t\terr = DeleteRunLocalValue(ctx, \"key\")\n\t\tassert.Error(t, err, \"DeleteRunLocalValue should fail outside agent context\")\n\t\tassert.Contains(t, err.Error(), \"DeleteRunLocalValue failed\")\n\t})\n\n\tt.Run(\"RunLocalValuePersistsAcrossModelCalls\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttestTool := &namedTool{name: \"test_tool\"}\n\t\tinfo, _ := testTool.Info(ctx)\n\n\t\tvar firstCallValue any\n\t\tvar secondCallValue any\n\t\tcallCount := 0\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Using tool\", []schema.ToolCall{\n\t\t\t\t{ID: \"call1\", Function: schema.FunctionCall{Name: info.Name, Arguments: \"{}\"}},\n\t\t\t}), nil).Times(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\tcallCount++\n\t\t\t\t\tif callCount == 1 {\n\t\t\t\t\t\terr := SetRunLocalValue(ctx, \"persist_key\", \"persist_value\")\n\t\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\t\tval, _, _ := GetRunLocalValue(ctx, \"persist_key\")\n\t\t\t\t\t\tfirstCallValue = val\n\t\t\t\t\t} else {\n\t\t\t\t\t\tval, _, _ := GetRunLocalValue(ctx, \"persist_key\")\n\t\t\t\t\t\tsecondCallValue = val\n\t\t\t\t\t}\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, \"persist_value\", firstCallValue, \"First call should set value\")\n\t\tassert.Equal(t, \"persist_value\", secondCallValue, \"Value should persist to second model call\")\n\t})\n}\n\nfunc TestHandlerErrorPropagation(t *testing.T) {\n\tt.Run(\"BeforeModelRewriteStateErrorStopsRun\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\treturn ctx, state, assert.AnError\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\n\t\tvar gotErr error\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Err != nil {\n\t\t\t\tgotErr = event.Err\n\t\t\t}\n\t\t}\n\n\t\tassert.Error(t, gotErr)\n\t})\n\n\tt.Run(\"AfterModelRewriteStateErrorStopsRun\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"response\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testAfterModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\treturn ctx, state, assert.AnError\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\n\t\tvar gotErr error\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Err != nil {\n\t\t\t\tgotErr = event.Err\n\t\t\t}\n\t\t}\n\n\t\tassert.Error(t, gotErr)\n\t})\n\n\tt.Run(\"MultipleHandlersFirstErrorStops\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tsecondHandlerCalled := false\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\treturn ctx, state, assert.AnError\n\t\t\t\t}},\n\t\t\t\t&testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {\n\t\t\t\t\tsecondHandlerCalled = true\n\t\t\t\t\treturn ctx, state, nil\n\t\t\t\t}},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.False(t, secondHandlerCalled, \"Second handler should not be called after first handler error\")\n\t})\n}\n\nfunc TestToolContextInWrappers(t *testing.T) {\n\tt.Run(\"ToolContextHasCorrectNameAndCallID\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\ttestTool := &namedTool{name: \"context_test_tool\"}\n\t\tinfo, _ := testTool.Info(ctx)\n\n\t\tvar capturedToolName string\n\t\tvar capturedCallID string\n\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Using tool\", []schema.ToolCall{\n\t\t\t\t{ID: \"test_call_id_123\", Function: schema.FunctionCall{Name: info.Name, Arguments: \"{}\"}},\n\t\t\t}), nil).Times(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"done\", nil), nil).Times(1)\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       cm,\n\t\t\tToolsConfig: ToolsConfig{\n\t\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\t\tTools: []tool.BaseTool{testTool},\n\t\t\t\t},\n\t\t\t},\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\t&testToolWrapperHandler{\n\t\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\t\twrapInvokableFn: func(_ context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) InvokableToolCallEndpoint {\n\t\t\t\t\t\tcapturedToolName = tCtx.Name\n\t\t\t\t\t\tcapturedCallID = tCtx.CallID\n\t\t\t\t\t\treturn endpoint\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\titer := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}})\n\t\tfor {\n\t\t\t_, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, \"context_test_tool\", capturedToolName, \"ToolContext should have correct tool name\")\n\t\tassert.Equal(t, \"test_call_id_123\", capturedCallID, \"ToolContext should have correct call ID\")\n\t})\n}\n"
  },
  {
    "path": "adk/instruction.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\n\t\"github.com/cloudwego/eino/adk/internal\"\n)\n\nconst (\n\tTransferToAgentInstruction = `Available other agents: %s\n\nDecision rule:\n- If you're best suited for the question according to your description: ANSWER\n- If another agent is better according its description: CALL '%s' function with their agent name\n\nWhen transferring: OUTPUT ONLY THE FUNCTION CALL`\n\n\tTransferToAgentInstructionChinese = `可用的其他 agent：%s\n\n决策规则：\n- 如果根据你的职责描述，你最适合回答这个问题：ANSWER\n- 如果根据其职责描述，另一个 agent 更适合：调用 %s 函数，并传入该 agent 的名称\n\n当进行移交时：只输出函数调用，不要输出其他任何内容`\n\n\tagentDescriptionTpl        = \"\\n- Agent name: %s\\n  Agent description: %s\"\n\tagentDescriptionTplChinese = \"\\n- Agent 名字: %s\\n  Agent 描述: %s\"\n)\n\nfunc genTransferToAgentInstruction(ctx context.Context, agents []Agent) string {\n\ttpl := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: agentDescriptionTpl,\n\t\tChinese: agentDescriptionTplChinese,\n\t})\n\tinstruction := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: TransferToAgentInstruction,\n\t\tChinese: TransferToAgentInstructionChinese,\n\t})\n\n\tvar sb strings.Builder\n\tfor _, agent := range agents {\n\t\tsb.WriteString(fmt.Sprintf(tpl, agent.Name(ctx), agent.Description(ctx)))\n\t}\n\n\treturn fmt.Sprintf(instruction, sb.String(), TransferToAgentToolName)\n}\n"
  },
  {
    "path": "adk/interface.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/gob\"\n\t\"fmt\"\n\t\"io\"\n\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/internal/core\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// ComponentOfAgent is the component type identifier for ADK agents in callbacks.\n// Use this to filter callback events to only agent-related events.\nconst ComponentOfAgent components.Component = \"Agent\"\n\ntype Message = *schema.Message\ntype MessageStream = *schema.StreamReader[Message]\n\ntype MessageVariant struct {\n\tIsStreaming bool\n\n\tMessage       Message\n\tMessageStream MessageStream\n\t// message role: Assistant or Tool\n\tRole schema.RoleType\n\t// only used when Role is Tool\n\tToolName string\n}\n\n// EventFromMessage wraps a message or stream into an AgentEvent with role metadata.\nfunc EventFromMessage(msg Message, msgStream MessageStream,\n\trole schema.RoleType, toolName string) *AgentEvent {\n\treturn &AgentEvent{\n\t\tOutput: &AgentOutput{\n\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\tIsStreaming:   msgStream != nil,\n\t\t\t\tMessage:       msg,\n\t\t\t\tMessageStream: msgStream,\n\t\t\t\tRole:          role,\n\t\t\t\tToolName:      toolName,\n\t\t\t},\n\t\t},\n\t}\n}\n\ntype messageVariantSerialization struct {\n\tIsStreaming   bool\n\tMessage       Message\n\tMessageStream Message\n\tRole          schema.RoleType\n\tToolName      string\n}\n\nfunc (mv *MessageVariant) GobEncode() ([]byte, error) {\n\ts := &messageVariantSerialization{\n\t\tIsStreaming: mv.IsStreaming,\n\t\tMessage:     mv.Message,\n\t\tRole:        mv.Role,\n\t\tToolName:    mv.ToolName,\n\t}\n\tif mv.IsStreaming {\n\t\tvar messages []Message\n\t\tfor {\n\t\t\tframe, err := mv.MessageStream.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"error receiving message stream: %w\", err)\n\t\t\t}\n\t\t\tmessages = append(messages, frame)\n\t\t}\n\t\tm, err := schema.ConcatMessages(messages)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to encode message: cannot concat message stream: %w\", err)\n\t\t}\n\t\ts.MessageStream = m\n\t}\n\tbuf := &bytes.Buffer{}\n\terr := gob.NewEncoder(buf).Encode(s)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to gob encode message variant: %w\", err)\n\t}\n\treturn buf.Bytes(), nil\n}\n\nfunc (mv *MessageVariant) GobDecode(b []byte) error {\n\ts := &messageVariantSerialization{}\n\terr := gob.NewDecoder(bytes.NewReader(b)).Decode(s)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to decoding message variant: %w\", err)\n\t}\n\tmv.IsStreaming = s.IsStreaming\n\tmv.Message = s.Message\n\tmv.Role = s.Role\n\tmv.ToolName = s.ToolName\n\tif s.MessageStream != nil {\n\t\tmv.MessageStream = schema.StreamReaderFromArray([]*schema.Message{s.MessageStream})\n\t}\n\treturn nil\n}\n\nfunc (mv *MessageVariant) GetMessage() (Message, error) {\n\tvar message Message\n\tif mv.IsStreaming {\n\t\tvar err error\n\t\tmessage, err = schema.ConcatMessageStream(mv.MessageStream)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t} else {\n\t\tmessage = mv.Message\n\t}\n\n\treturn message, nil\n}\n\ntype TransferToAgentAction struct {\n\tDestAgentName string\n}\n\ntype AgentOutput struct {\n\tMessageOutput *MessageVariant\n\n\tCustomizedOutput any\n}\n\n// NewTransferToAgentAction creates an action to transfer to the specified agent.\nfunc NewTransferToAgentAction(destAgentName string) *AgentAction {\n\treturn &AgentAction{TransferToAgent: &TransferToAgentAction{DestAgentName: destAgentName}}\n}\n\n// NewExitAction creates an action that signals the agent to exit.\nfunc NewExitAction() *AgentAction {\n\treturn &AgentAction{Exit: true}\n}\n\n// AgentAction represents actions that an agent can emit during execution.\n//\n// Action Scoping in Agent Tools:\n// When an agent is wrapped as an agent tool (via NewAgentTool), actions emitted by the inner agent\n// are scoped to the tool boundary:\n//   - Interrupted: Propagated via CompositeInterrupt to allow proper interrupt/resume across boundaries\n//   - Exit, TransferToAgent, BreakLoop: Ignored outside the agent tool; these actions only affect\n//     the inner agent's execution and do not propagate to the parent agent\n//\n// This scoping ensures that nested agents cannot unexpectedly terminate or transfer control\n// of their parent agent's execution flow.\ntype AgentAction struct {\n\tExit bool\n\n\tInterrupted *InterruptInfo\n\n\tTransferToAgent *TransferToAgentAction\n\n\tBreakLoop *BreakLoopAction\n\n\tCustomizedAction any\n\n\tinternalInterrupted *core.InterruptSignal\n}\n\n// RunStep CheckpointSchema: persisted via serialization.RunCtx (gob).\ntype RunStep struct {\n\tagentName string\n}\n\nfunc init() {\n\tschema.RegisterName[[]RunStep](\"eino_run_step_list\")\n}\n\nfunc (r *RunStep) String() string {\n\treturn r.agentName\n}\n\nfunc (r *RunStep) Equals(r1 RunStep) bool {\n\treturn r.agentName == r1.agentName\n}\n\nfunc (r *RunStep) GobEncode() ([]byte, error) {\n\ts := &runStepSerialization{AgentName: r.agentName}\n\tbuf := &bytes.Buffer{}\n\terr := gob.NewEncoder(buf).Encode(s)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to gob encode RunStep: %w\", err)\n\t}\n\treturn buf.Bytes(), nil\n}\n\nfunc (r *RunStep) GobDecode(b []byte) error {\n\ts := &runStepSerialization{}\n\terr := gob.NewDecoder(bytes.NewReader(b)).Decode(s)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to gob decode RunStep: %w\", err)\n\t}\n\tr.agentName = s.AgentName\n\treturn nil\n}\n\ntype runStepSerialization struct {\n\tAgentName string\n}\n\n// AgentEvent CheckpointSchema: persisted via serialization.RunCtx (gob).\ntype AgentEvent struct {\n\tAgentName string\n\n\t// RunPath represents the execution path from root agent to the current event source.\n\t// This field is managed entirely by the eino framework and cannot be set by end-users\n\t// because RunStep's fields are unexported. The framework sets RunPath exactly once:\n\t// - flowAgent sets it when the event has no RunPath (len == 0)\n\t// - agentTool prepends parent RunPath when forwarding events from nested agents\n\tRunPath []RunStep\n\n\tOutput *AgentOutput\n\n\tAction *AgentAction\n\n\tErr error\n}\n\ntype AgentInput struct {\n\tMessages        []Message\n\tEnableStreaming bool\n}\n\n//go:generate  mockgen -destination ../internal/mock/adk/Agent_mock.go --package adk -source interface.go\ntype Agent interface {\n\tName(ctx context.Context) string\n\tDescription(ctx context.Context) string\n\n\t// Run runs the agent.\n\t// The returned AgentEvent within the AsyncIterator must be safe to modify.\n\t// If the returned AgentEvent within the AsyncIterator contains MessageStream,\n\t// the MessageStream MUST be exclusive and safe to be received directly.\n\t// NOTE: it's recommended to use SetAutomaticClose() on the MessageStream of AgentEvents emitted by AsyncIterator,\n\t// so that even the events are not processed, the MessageStream can still be closed.\n\tRun(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]\n}\n\ntype OnSubAgents interface {\n\tOnSetSubAgents(ctx context.Context, subAgents []Agent) error\n\tOnSetAsSubAgent(ctx context.Context, parent Agent) error\n\n\tOnDisallowTransferToParent(ctx context.Context) error\n}\n\ntype ResumableAgent interface {\n\tAgent\n\n\tResume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]\n}\n"
  },
  {
    "path": "adk/internal/config.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package internal provides adk internal utils.\npackage internal\n\nimport (\n\t\"fmt\"\n\t\"sync/atomic\"\n)\n\n// Language represents the language setting for the ADK built-in prompts.\ntype Language uint8\n\nconst (\n\t// LanguageEnglish represents English language.\n\tLanguageEnglish Language = iota\n\t// LanguageChinese represents Chinese language.\n\tLanguageChinese\n)\n\nvar language atomic.Value\n\n// SetLanguage sets the language for the ADK built-in prompts.\n// The default language is English if not explicitly set.\nfunc SetLanguage(lang Language) error {\n\tif lang != LanguageEnglish &&\n\t\tlang != LanguageChinese {\n\t\treturn fmt.Errorf(\"invalid language: %v\", lang)\n\t}\n\tlanguage.Store(lang)\n\treturn nil\n}\n\n// GetLanguage returns the current language setting for the ADK built-in prompts.\n// Returns LanguageEnglish if no language has been set.\nfunc getLanguage() Language {\n\tif l, ok := language.Load().(Language); ok {\n\t\treturn l\n\t}\n\treturn LanguageEnglish\n}\n\n// I18nPrompts holds prompt strings for different languages.\ntype I18nPrompts struct {\n\tEnglish string\n\tChinese string\n}\n\n// SelectPrompt returns the appropriate prompt string based on the current language setting.\n// Returns an error if the current language is not supported.\nfunc SelectPrompt(prompts I18nPrompts) string {\n\tlang := getLanguage()\n\tswitch lang {\n\tcase LanguageEnglish:\n\t\treturn prompts.English\n\tcase LanguageChinese:\n\t\treturn prompts.Chinese\n\tdefault:\n\t\t// unreachable\n\t\tpanic(fmt.Sprintf(\"invalid language: %v\", lang))\n\t}\n}\n"
  },
  {
    "path": "adk/interrupt.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/gob\"\n\t\"errors\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/internal/core\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// ResumeInfo holds all the information necessary to resume an interrupted agent execution.\n// It is created by the framework and passed to an agent's Resume method.\ntype ResumeInfo struct {\n\t// EnableStreaming indicates whether the original execution was in streaming mode.\n\tEnableStreaming bool\n\n\t// Deprecated: use InterruptContexts from the embedded InterruptInfo for user-facing details,\n\t// and GetInterruptState for internal state retrieval.\n\t*InterruptInfo\n\n\tWasInterrupted bool\n\tInterruptState any\n\tIsResumeTarget bool\n\tResumeData     any\n}\n\n// InterruptInfo contains all the information about an interruption event.\n// It is created by the framework when an agent returns an interrupt action.\ntype InterruptInfo struct {\n\tData any\n\n\t// InterruptContexts provides a structured, user-facing view of the interrupt chain.\n\t// Each context represents a step in the agent hierarchy that was interrupted.\n\tInterruptContexts []*InterruptCtx\n}\n\n// Interrupt creates a basic interrupt action.\n// This is used when an agent needs to pause its execution to request external input or intervention,\n// but does not need to save any internal state to be restored upon resumption.\n// The `info` parameter is user-facing data that describes the reason for the interrupt.\nfunc Interrupt(ctx context.Context, info any) *AgentEvent {\n\tvar rp []RunStep\n\trCtx := getRunCtx(ctx)\n\tif rCtx != nil {\n\t\trp = rCtx.RunPath\n\t}\n\n\tis, err := core.Interrupt(ctx, info, nil, nil,\n\t\tcore.WithLayerPayload(rp))\n\tif err != nil {\n\t\treturn &AgentEvent{Err: err}\n\t}\n\n\tcontexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)\n\n\treturn &AgentEvent{\n\t\tAction: &AgentAction{\n\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\tInterruptContexts: contexts,\n\t\t\t},\n\t\t\tinternalInterrupted: is,\n\t\t},\n\t}\n}\n\n// StatefulInterrupt creates an interrupt action that also saves the agent's internal state.\n// This is used when an agent has internal state that must be restored for it to continue correctly.\n// The `info` parameter is user-facing data describing the interrupt.\n// The `state` parameter is the agent's internal state object, which will be serialized and stored.\nfunc StatefulInterrupt(ctx context.Context, info any, state any) *AgentEvent {\n\tvar rp []RunStep\n\trCtx := getRunCtx(ctx)\n\tif rCtx != nil {\n\t\trp = rCtx.RunPath\n\t}\n\n\tis, err := core.Interrupt(ctx, info, state, nil,\n\t\tcore.WithLayerPayload(rp))\n\tif err != nil {\n\t\treturn &AgentEvent{Err: err}\n\t}\n\n\tcontexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)\n\n\treturn &AgentEvent{\n\t\tAction: &AgentAction{\n\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\tInterruptContexts: contexts,\n\t\t\t},\n\t\t\tinternalInterrupted: is,\n\t\t},\n\t}\n}\n\n// CompositeInterrupt creates an interrupt action for a workflow agent.\n// It combines the interrupts from one or more of its sub-agents into a single, cohesive interrupt.\n// This is used by workflow agents (like Sequential, Parallel, or Loop) to propagate interrupts from their children.\n// The `info` parameter is user-facing data describing the workflow's own reason for interrupting.\n// The `state` parameter is the workflow agent's own state (e.g., the index of the sub-agent that was interrupted).\n// The `subInterruptSignals` is a variadic list of the InterruptSignal objects from the interrupted sub-agents.\nfunc CompositeInterrupt(ctx context.Context, info any, state any,\n\tsubInterruptSignals ...*InterruptSignal) *AgentEvent {\n\tvar rp []RunStep\n\trCtx := getRunCtx(ctx)\n\tif rCtx != nil {\n\t\trp = rCtx.RunPath\n\t}\n\n\tis, err := core.Interrupt(ctx, info, state, subInterruptSignals,\n\t\tcore.WithLayerPayload(rp))\n\tif err != nil {\n\t\treturn &AgentEvent{Err: err}\n\t}\n\n\tcontexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)\n\n\treturn &AgentEvent{\n\t\tAction: &AgentAction{\n\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\tInterruptContexts: contexts,\n\t\t\t},\n\t\t\tinternalInterrupted: is,\n\t\t},\n\t}\n}\n\n// Address represents the unique, hierarchical address of a component within an execution.\n// It is a slice of AddressSegments, where each segment represents one level of nesting.\n// This is a type alias for core.Address. See the core package for more details.\ntype Address = core.Address\ntype AddressSegment = core.AddressSegment\ntype AddressSegmentType = core.AddressSegmentType\n\nconst (\n\tAddressSegmentAgent AddressSegmentType = \"agent\"\n\tAddressSegmentTool  AddressSegmentType = \"tool\"\n)\n\nvar allowedAddressSegmentTypes = []AddressSegmentType{AddressSegmentAgent, AddressSegmentTool}\n\n// AppendAddressSegment adds an address segment for the current execution context.\nfunc AppendAddressSegment(ctx context.Context, segType AddressSegmentType, segID string) context.Context {\n\treturn core.AppendAddressSegment(ctx, segType, segID, \"\")\n}\n\n// InterruptCtx provides a structured, user-facing view of a single point of interruption.\n// It contains the ID and Address of the interrupted component, as well as user-defined info.\n// This is a type alias for core.InterruptCtx. See the core package for more details.\ntype InterruptCtx = core.InterruptCtx\ntype InterruptSignal = core.InterruptSignal\n\n// FromInterruptContexts converts user-facing interrupt contexts to an interrupt signal.\nfunc FromInterruptContexts(contexts []*InterruptCtx) *InterruptSignal {\n\treturn core.FromInterruptContexts(contexts)\n}\n\n// WithCheckPointID sets the checkpoint ID used for interruption persistence.\nfunc WithCheckPointID(id string) AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(t *options) {\n\t\tt.checkPointID = &id\n\t})\n}\n\nfunc init() {\n\tschema.RegisterName[*serialization](\"_eino_adk_serialization\")\n\tschema.RegisterName[*WorkflowInterruptInfo](\"_eino_adk_workflow_interrupt_info\")\n}\n\n// serialization CheckpointSchema: root checkpoint payload (gob).\n// Any type tagged with `CheckpointSchema:` is persisted and must remain backward compatible.\ntype serialization struct {\n\tRunCtx *runContext\n\t// deprecated: still keep it here for backward compatibility\n\tInfo                *InterruptInfo\n\tEnableStreaming     bool\n\tInterruptID2Address map[string]Address\n\tInterruptID2State   map[string]core.InterruptState\n}\n\nfunc (r *Runner) loadCheckPoint(ctx context.Context, checkpointID string) (\n\tcontext.Context, *runContext, *ResumeInfo, error) {\n\tdata, existed, err := r.store.Get(ctx, checkpointID)\n\tif err != nil {\n\t\treturn nil, nil, nil, fmt.Errorf(\"failed to get checkpoint from store: %w\", err)\n\t}\n\tif !existed {\n\t\treturn nil, nil, nil, fmt.Errorf(\"checkpoint[%s] not exist\", checkpointID)\n\t}\n\n\tdata = preprocessADKCheckpoint(data)\n\n\ts := &serialization{}\n\terr = gob.NewDecoder(bytes.NewReader(data)).Decode(s)\n\tif err != nil {\n\t\treturn nil, nil, nil, fmt.Errorf(\"failed to decode checkpoint: %w\", err)\n\t}\n\tctx = core.PopulateInterruptState(ctx, s.InterruptID2Address, s.InterruptID2State)\n\n\treturn ctx, s.RunCtx, &ResumeInfo{\n\t\tEnableStreaming: s.EnableStreaming,\n\t\tInterruptInfo:   s.Info,\n\t}, nil\n}\n\n// preprocessADKCheckpoint fixes a gob incompatibility when resuming old ChatModelAgent/DeepAgents checkpoints.\n//\n// Background\n//   - ADK checkpoints are gob-encoded.\n//   - Some values inside checkpoints are stored as `any`, so gob includes a concrete type name\n//     string in the wire format and uses that name to pick the local Go type to decode into.\n//\n// Problem (v0.8.0-v0.8.3 checkpoints)\n//   - In v0.8.0-v0.8.3, *State was registered under the name \"_eino_adk_react_state\" AND\n//     implemented GobEncode/GobDecode, so the wire format for that name is \"GobEncoder payload\"\n//     (opaque bytes).\n//   - In v0.7.*, the same name \"_eino_adk_react_state\" was used but encoded as a normal struct\n//     (no GobEncode). Gob treats these two wire formats as incompatible.\n//   - Gob only allows one local Go type per name. Today we register \"_eino_adk_react_state\" to\n//     a v0.7-compatible struct decoder (stateV07). If we try to decode a v0.8.0-v0.8.3\n//     checkpoint under that same name, gob fails with a \"want struct; got non-struct\" mismatch.\n//\n// Solution\n//   - We keep \"_eino_adk_react_state\" mapped to the v0.7 decoder.\n//   - For v0.8.0-v0.8.3 checkpoints only, we rewrite the on-wire name to a same-length alias\n//     \"_eino_adk_state_v080_\", which is registered to a GobDecoder-compatible type (stateV080).\n//   - The alias is the same length as the original, so we can safely replace the length-prefixed\n//     bytes without re-encoding the whole stream.\nfunc preprocessADKCheckpoint(data []byte) []byte {\n\tconst (\n\t\tlenPrefixedReactStateName         = \"\\x15\" + stateGobNameV07\n\t\tlenPrefixedCompatName             = \"\\x15\" + stateGobNameV080\n\t\tlenPrefixedStateSerializationName = \"\\x12stateSerialization\"\n\t)\n\n\t// the following line checks whether the checkpoint is persisted through v0.8.0-v0.8.3\n\tif !bytes.Contains(data, []byte(lenPrefixedReactStateName)) || !bytes.Contains(data, []byte(lenPrefixedStateSerializationName)) {\n\t\treturn data\n\t}\n\treturn bytes.ReplaceAll(data,\n\t\t[]byte(lenPrefixedReactStateName),\n\t\t[]byte(lenPrefixedCompatName))\n}\n\nfunc (r *Runner) saveCheckPoint(\n\tctx context.Context,\n\tkey string,\n\tinfo *InterruptInfo,\n\tis *core.InterruptSignal,\n) error {\n\trunCtx := getRunCtx(ctx)\n\n\tid2Addr, id2State := core.SignalToPersistenceMaps(is)\n\n\tbuf := &bytes.Buffer{}\n\terr := gob.NewEncoder(buf).Encode(&serialization{\n\t\tRunCtx:              runCtx,\n\t\tInfo:                info,\n\t\tInterruptID2Address: id2Addr,\n\t\tInterruptID2State:   id2State,\n\t\tEnableStreaming:     r.enableStreaming,\n\t})\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to encode checkpoint: %w\", err)\n\t}\n\treturn r.store.Set(ctx, key, buf.Bytes())\n}\n\nconst bridgeCheckpointID = \"adk_react_mock_key\"\n\nfunc newBridgeStore() *bridgeStore {\n\treturn &bridgeStore{}\n}\n\nfunc newResumeBridgeStore(data []byte) *bridgeStore {\n\treturn &bridgeStore{\n\t\tData:  data,\n\t\tValid: true,\n\t}\n}\n\ntype bridgeStore struct {\n\tData  []byte\n\tValid bool\n}\n\nfunc (m *bridgeStore) Get(_ context.Context, _ string) ([]byte, bool, error) {\n\tif m.Valid {\n\t\treturn m.Data, true, nil\n\t}\n\treturn nil, false, nil\n}\n\nfunc (m *bridgeStore) Set(_ context.Context, _ string, checkPoint []byte) error {\n\tm.Data = checkPoint\n\tm.Valid = true\n\treturn nil\n}\n\nfunc getNextResumeAgent(ctx context.Context, info *ResumeInfo) (string, error) {\n\tnextAgents, err := core.GetNextResumptionPoints(ctx)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to get next agent leading to interruption: %w\", err)\n\t}\n\n\tif len(nextAgents) == 0 {\n\t\treturn \"\", errors.New(\"no child agents leading to interrupted agent were found\")\n\t}\n\n\tif len(nextAgents) > 1 {\n\t\treturn \"\", errors.New(\"agent has multiple child agents leading to interruption, \" +\n\t\t\t\"but concurrent transfer is not supported\")\n\t}\n\n\t// get the single next agent to delegate to.\n\tvar nextAgentID string\n\tfor id := range nextAgents {\n\t\tnextAgentID = id\n\t\tbreak\n\t}\n\n\treturn nextAgentID, nil\n}\n\nfunc getNextResumeAgents(ctx context.Context, info *ResumeInfo) (map[string]bool, error) {\n\tnextAgents, err := core.GetNextResumptionPoints(ctx)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to get next agents leading to interruption: %w\", err)\n\t}\n\n\tif len(nextAgents) == 0 {\n\t\treturn nil, errors.New(\"no child agents leading to interrupted agent were found\")\n\t}\n\n\treturn nextAgents, nil\n}\n\nfunc buildResumeInfo(ctx context.Context, nextAgentID string, info *ResumeInfo) (\n\tcontext.Context, *ResumeInfo) {\n\tctx = AppendAddressSegment(ctx, AddressSegmentAgent, nextAgentID)\n\tnextResumeInfo := &ResumeInfo{\n\t\tEnableStreaming: info.EnableStreaming,\n\t\tInterruptInfo:   info.InterruptInfo,\n\t}\n\n\twasInterrupted, hasState, state := core.GetInterruptState[any](ctx)\n\tnextResumeInfo.WasInterrupted = wasInterrupted\n\tif hasState {\n\t\tnextResumeInfo.InterruptState = state\n\t}\n\n\tif wasInterrupted {\n\t\tisResumeTarget, hasData, data := core.GetResumeContext[any](ctx)\n\t\tnextResumeInfo.IsResumeTarget = isResumeTarget\n\t\tif hasData {\n\t\t\tnextResumeInfo.ResumeData = data\n\t\t}\n\t}\n\n\tctx = updateRunPathOnly(ctx, nextAgentID)\n\n\treturn ctx, nextResumeInfo\n}\n"
  },
  {
    "path": "adk/interrupt_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype interruptTestToolsHandler struct {\n\t*BaseChatModelAgentMiddleware\n\ttools []tool.BaseTool\n}\n\nfunc TestPreprocessADKCheckpoint(t *testing.T) {\n\tt.Run(\"no-op when missing markers\", func(t *testing.T) {\n\t\tin := []byte(\"random\")\n\t\tout := preprocessADKCheckpoint(append([]byte(nil), in...))\n\t\tassert.Equal(t, in, out)\n\t})\n\n\tt.Run(\"rewrite legacy name for v0.8.0-v0.8.3\", func(t *testing.T) {\n\t\tconst (\n\t\t\tlenPrefixedReactStateName         = \"\\x15\" + stateGobNameV07\n\t\t\tlenPrefixedCompatName             = \"\\x15\" + stateGobNameV080\n\t\t\tlenPrefixedStateSerializationName = \"\\x12stateSerialization\"\n\t\t)\n\n\t\tin := []byte(lenPrefixedReactStateName + \"xxx\" + lenPrefixedStateSerializationName + \"yyy\")\n\t\tout := preprocessADKCheckpoint(append([]byte(nil), in...))\n\t\tassert.True(t, bytes.Contains(out, []byte(lenPrefixedCompatName)))\n\t\tassert.False(t, bytes.Contains(out, []byte(lenPrefixedReactStateName)))\n\t})\n}\n\nfunc (h *interruptTestToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {\n\trunCtx.Tools = append(runCtx.Tools, h.tools...)\n\treturn ctx, runCtx, nil\n}\n\nfunc TestSaveAgentEventWrapper(t *testing.T) {\n\tsr, sw := schema.Pipe[Message](1)\n\tsw.Send(schema.UserMessage(\"test\"), nil)\n\tsw.Close()\n\tsr = sr.Copy(2)[1]\n\n\tw := &agentEventWrapper{\n\t\tAgentEvent: &AgentEvent{\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessageStream: sr,\n\t\t\t\t},\n\t\t\t},\n\t\t\tRunPath: []RunStep{\n\t\t\t\t{\n\t\t\t\t\t\"a1\",\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\t\"a2\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\tmu:                  sync.Mutex{},\n\t\tconcatenatedMessage: nil,\n\t}\n\n\t_, err := getMessageFromWrappedEvent(w)\n\tassert.NoError(t, err)\n\n\tbuf, err := w.GobEncode()\n\tassert.NoError(t, err)\n\tassert.NoError(t, err)\n\n\tw1 := &agentEventWrapper{}\n\terr = w1.GobDecode(buf)\n\tassert.NoError(t, err)\n}\n\nfunc TestInterruptFunctionsPopulateInterruptContextsImmediately(t *testing.T) {\n\tctx := context.Background()\n\tctx, _ = initRunCtx(ctx, \"TestAgent\", &AgentInput{Messages: []Message{}})\n\tctx = AppendAddressSegment(ctx, AddressSegmentAgent, \"TestAgent\")\n\n\tt.Run(\"Interrupt populates InterruptContexts\", func(t *testing.T) {\n\t\tevent := Interrupt(ctx, \"test info\")\n\t\tassert.NotNil(t, event.Action)\n\t\tassert.NotNil(t, event.Action.Interrupted)\n\t\tassert.NotNil(t, event.Action.Interrupted.InterruptContexts)\n\t\tassert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))\n\t\tassert.Equal(t, \"test info\", event.Action.Interrupted.InterruptContexts[0].Info)\n\t\tassert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)\n\t\tassert.Equal(t, Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"TestAgent\"},\n\t\t}, event.Action.Interrupted.InterruptContexts[0].Address)\n\t})\n\n\tt.Run(\"StatefulInterrupt populates InterruptContexts\", func(t *testing.T) {\n\t\tevent := StatefulInterrupt(ctx, \"stateful info\", \"my state\")\n\t\tassert.NotNil(t, event.Action)\n\t\tassert.NotNil(t, event.Action.Interrupted)\n\t\tassert.NotNil(t, event.Action.Interrupted.InterruptContexts)\n\t\tassert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))\n\t\tassert.Equal(t, \"stateful info\", event.Action.Interrupted.InterruptContexts[0].Info)\n\t\tassert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)\n\t})\n\n\tt.Run(\"CompositeInterrupt populates InterruptContexts with filtered parent chain\", func(t *testing.T) {\n\t\tsubCtx := AppendAddressSegment(ctx, AddressSegmentAgent, \"SubAgent\")\n\t\tsubEvent := Interrupt(subCtx, \"sub info\")\n\t\tevent := CompositeInterrupt(ctx, \"composite info\", \"composite state\", subEvent.Action.internalInterrupted)\n\t\tassert.NotNil(t, event.Action)\n\t\tassert.NotNil(t, event.Action.Interrupted)\n\t\tassert.NotNil(t, event.Action.Interrupted.InterruptContexts)\n\t\tassert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))\n\n\t\trootCause := event.Action.Interrupted.InterruptContexts[0]\n\t\tassert.Equal(t, \"sub info\", rootCause.Info)\n\t\tassert.True(t, rootCause.IsRootCause)\n\t\tassert.Equal(t, Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"TestAgent\"},\n\t\t\t{Type: AddressSegmentAgent, ID: \"SubAgent\"},\n\t\t}, rootCause.Address)\n\n\t\tassert.NotNil(t, rootCause.Parent, \"Parent should not be nil for composite interrupt\")\n\t\tassert.Equal(t, \"composite info\", rootCause.Parent.Info)\n\t\tassert.Equal(t, Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"TestAgent\"},\n\t\t}, rootCause.Parent.Address)\n\t})\n\n\tt.Run(\"Address only contains agent/tool segments\", func(t *testing.T) {\n\t\tevent := Interrupt(ctx, \"test info\")\n\t\taddr := event.Action.Interrupted.InterruptContexts[0].Address\n\t\tfor _, seg := range addr {\n\t\t\tassert.True(t, seg.Type == AddressSegmentAgent || seg.Type == AddressSegmentTool,\n\t\t\t\t\"Address should only contain agent/tool segments, got: %s\", seg.Type)\n\t\t}\n\t})\n}\n\nfunc TestSimpleInterrupt(t *testing.T) {\n\tdata := \"hello world\"\n\tagent := &myAgent{\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tIsStreaming: true,\n\t\t\t\t\t\tMessage:     nil,\n\t\t\t\t\t\tMessageStream: schema.StreamReaderFromArray([]Message{\n\t\t\t\t\t\t\tschema.UserMessage(\"hello \"),\n\t\t\t\t\t\t\tschema.UserMessage(\"world\"),\n\t\t\t\t\t\t}),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tintEvent := Interrupt(ctx, data)\n\t\t\tintEvent.Action.Interrupted.Data = data\n\t\t\tgenerator.Send(intEvent)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.True(t, info.WasInterrupted)\n\t\t\tassert.Nil(t, info.InterruptState)\n\t\t\tassert.True(t, info.EnableStreaming)\n\t\t\tassert.Equal(t, data, info.Data)\n\n\t\t\tassert.True(t, info.IsResumeTarget)\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\tstore := newMyStore()\n\tctx := context.Background()\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           agent,\n\t\tEnableStreaming: true,\n\t\tCheckPointStore: store,\n\t})\n\titer := runner.Query(ctx, \"hello world\", WithCheckPointID(\"1\"))\n\t_, ok := iter.Next()\n\tassert.True(t, ok)\n\tinterruptEvent, ok := iter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, data, interruptEvent.Action.Interrupted.Data)\n\tassert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID)\n\tassert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause)\n\tassert.Equal(t, data, interruptEvent.Action.Interrupted.InterruptContexts[0].Info)\n\tassert.Equal(t, Address{{Type: AddressSegmentAgent, ID: \"myAgent\"}},\n\t\tinterruptEvent.Action.Interrupted.InterruptContexts[0].Address)\n\t_, ok = iter.Next()\n\tassert.False(t, ok)\n\n\titer, err := runner.ResumeWithParams(ctx, \"1\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\tinterruptEvent.Action.Interrupted.InterruptContexts[0].ID: nil,\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\t_, ok = iter.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestMultiAgentInterrupt(t *testing.T) {\n\tctx := context.Background()\n\tsa1 := &myAgent{\n\t\tname: \"sa1\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"sa1\",\n\t\t\t\tAction: &AgentAction{\n\t\t\t\t\tTransferToAgent: &TransferToAgentAction{\n\t\t\t\t\t\tDestAgentName: \"sa2\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\tsa2 := &myAgent{\n\t\tname: \"sa2\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tintEvent := StatefulInterrupt(ctx, \"hello world\", \"temp state\")\n\t\t\tintEvent.Action.Interrupted.Data = \"hello world\"\n\t\t\tgenerator.Send(intEvent)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.NotNil(t, info)\n\t\t\tassert.Equal(t, info.Data, \"hello world\")\n\n\t\t\tassert.True(t, info.WasInterrupted)\n\t\t\tassert.NotNil(t, info.InterruptState)\n\t\t\tassert.Equal(t, \"temp state\", info.InterruptState)\n\n\t\t\tassert.True(t, info.IsResumeTarget)\n\t\t\tassert.NotNil(t, info.ResumeData)\n\t\t\tassert.Equal(t, \"resume data\", info.ResumeData)\n\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"sa2\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{Message: schema.UserMessage(info.ResumeData.(string))},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\ta, err := SetSubAgents(ctx, sa1, []Agent{sa2})\n\tassert.NoError(t, err)\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           a,\n\t\tEnableStreaming: false,\n\t\tCheckPointStore: newMyStore(),\n\t})\n\titer := runner.Query(ctx, \"\", WithCheckPointID(\"1\"))\n\tevent, ok := iter.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event.Action.TransferToAgent)\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event.Action.Interrupted)\n\tassert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))\n\tassert.Equal(t, \"hello world\", event.Action.Interrupted.InterruptContexts[0].Info)\n\tassert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)\n\tassert.Equal(t, Address{\n\t\t{Type: AddressSegmentAgent, ID: \"sa1\"},\n\t\t{Type: AddressSegmentAgent, ID: \"sa2\"},\n\t}, event.Action.Interrupted.InterruptContexts[0].Address)\n\tassert.NotEmpty(t, event.Action.Interrupted.InterruptContexts[0].ID)\n\n\tinterruptID := event.Action.Interrupted.InterruptContexts[0].ID\n\t_, ok = iter.Next()\n\tassert.False(t, ok)\n\n\titer, err = runner.ResumeWithParams(ctx, \"1\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\tinterruptID: \"resume data\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, event.Output.MessageOutput.Message.Content, \"resume data\")\n\t_, ok = iter.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestWorkflowInterrupt(t *testing.T) {\n\tctx := context.Background()\n\tsa1 := &myAgent{\n\t\tname: \"sa1\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\t\t\tintEvent := Interrupt(ctx, \"sa1 interrupt data\")\n\t\t\tintEvent.Action.Interrupted.Data = \"sa1 interrupt data\"\n\t\t\tgenerator.Send(intEvent)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.Equal(t, info.InterruptInfo.Data, \"sa1 interrupt data\")\n\t\t\tassert.True(t, info.WasInterrupted)\n\t\t\tassert.Nil(t, info.InterruptState)\n\t\t\tassert.True(t, info.IsResumeTarget)\n\t\t\tassert.Equal(t, \"resume sa1\", info.ResumeData)\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t} // interrupt once\n\tsa2 := &myAgent{\n\t\tname: \"sa2\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\t\t\tintEvent := StatefulInterrupt(ctx, \"sa2 interrupt data\", \"sa2 interrupt\")\n\t\t\tintEvent.Action.Interrupted.Data = \"sa2 interrupt data\"\n\t\t\tgenerator.Send(intEvent)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.Equal(t, info.InterruptInfo.Data, \"sa2 interrupt data\")\n\t\t\tassert.True(t, info.WasInterrupted)\n\t\t\tassert.NotNil(t, info.InterruptState)\n\t\t\tassert.Equal(t, \"sa2 interrupt\", info.InterruptState)\n\n\t\t\tassert.True(t, info.IsResumeTarget)\n\t\t\tassert.NotNil(t, info.ResumeData)\n\t\t\tassert.Equal(t, \"resume sa2\", info.ResumeData)\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t} // interrupt once\n\tsa3 := &myAgent{\n\t\tname: \"sa3\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"sa3\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa3 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t} // won't interrupt\n\tsa4 := &myAgent{\n\t\tname: \"sa4\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"sa4\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa4 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t} // won't interrupt\n\n\tfirstInterruptEvent := &AgentEvent{\n\t\tAgentName: \"sa1\",\n\t\tRunPath:   []RunStep{{\"sequential\"}, {\"sa1\"}},\n\t\tAction: &AgentAction{\n\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\tData: &WorkflowInterruptInfo{\n\t\t\t\t\tOrigInput: &AgentInput{\n\t\t\t\t\t\tMessages: []Message{schema.UserMessage(\"hello world\")},\n\t\t\t\t\t},\n\t\t\t\t\tSequentialInterruptIndex: 0,\n\t\t\t\t\tSequentialInterruptInfo: &InterruptInfo{\n\t\t\t\t\t\tData: \"sa1 interrupt data\",\n\t\t\t\t\t},\n\t\t\t\t\tLoopIterations: 0,\n\t\t\t\t},\n\t\t\t\tInterruptContexts: []*InterruptCtx{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:   \"agent:sequential;agent:sa1\",\n\t\t\t\t\t\tInfo: \"sa1 interrupt data\",\n\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID:   \"sequential\",\n\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID:   \"sa1\",\n\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tIsRootCause: true,\n\t\t\t\t\t\tParent: &InterruptCtx{\n\t\t\t\t\t\t\tID:   \"agent:sequential\",\n\t\t\t\t\t\t\tInfo: \"Sequential workflow interrupted\",\n\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"sequential\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\t_ = firstInterruptEvent\n\tsecondInterruptEvent := &AgentEvent{\n\t\tAgentName: \"sa2\",\n\t\tRunPath:   []RunStep{{\"sequential\"}, {\"sa1\"}, {\"sa2\"}},\n\t\tAction: &AgentAction{\n\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\tData: &WorkflowInterruptInfo{\n\t\t\t\t\tOrigInput: &AgentInput{\n\t\t\t\t\t\tMessages: []Message{schema.UserMessage(\"hello world\")},\n\t\t\t\t\t},\n\t\t\t\t\tSequentialInterruptIndex: 1,\n\t\t\t\t\tSequentialInterruptInfo: &InterruptInfo{\n\t\t\t\t\t\tData: \"sa2 interrupt data\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\tInterruptContexts: []*InterruptCtx{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:   \"agent:sequential;agent:sa1;agent:sa2\",\n\t\t\t\t\t\tInfo: \"sa2 interrupt data\",\n\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID:   \"sequential\",\n\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID:   \"sa2\",\n\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tIsRootCause: true,\n\t\t\t\t\t\tParent: &InterruptCtx{\n\t\t\t\t\t\t\tID:   \"agent:sequential\",\n\t\t\t\t\t\t\tInfo: \"Sequential workflow interrupted\",\n\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"sequential\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\t_ = secondInterruptEvent\n\tmessageEvents := []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"sa3\",\n\t\t\tRunPath:   []RunStep{{\"sequential\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}},\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tMessage: schema.UserMessage(\"sa3 completed\"),\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tAgentName: \"sa4\",\n\t\t\tRunPath:   []RunStep{{\"sequential\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}, {\"sa4\"}},\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tMessage: schema.UserMessage(\"sa4 completed\"),\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\t_ = messageEvents\n\n\tt.Run(\"test sequential workflow agent\", func(t *testing.T) {\n\n\t\t// sequential\n\t\ta, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\t\tName:        \"sequential\",\n\t\t\tDescription: \"sequential agent\",\n\t\t\tSubAgents:   []Agent{sa1, sa2, sa3, sa4},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\trunner := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           a,\n\t\t\tCheckPointStore: newMyStore(),\n\t\t})\n\t\tvar events []*AgentEvent\n\t\titer := runner.Query(ctx, \"hello world\", WithCheckPointID(\"sequential-1\"))\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tassert.Equal(t, 1, len(events))\n\t\tassert.Equal(t, firstInterruptEvent.AgentName, events[0].AgentName)\n\t\tassert.Equal(t, firstInterruptEvent.RunPath, events[0].RunPath)\n\t\tassert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(firstInterruptEvent.Action.Interrupted.InterruptContexts[0]))\n\t\tinterruptID1 := events[0].Action.Interrupted.InterruptContexts[0].ID\n\t\tevents = []*AgentEvent{}\n\n\t\t// Resume after sa1 interrupt\n\t\titer, err = runner.ResumeWithParams(ctx, \"sequential-1\", &ResumeParams{\n\t\t\tTargets: map[string]any{\n\t\t\t\tinterruptID1: \"resume sa1\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tassert.Equal(t, 1, len(events))\n\t\tassert.Equal(t, secondInterruptEvent.AgentName, events[0].AgentName)\n\t\tassert.Equal(t, secondInterruptEvent.RunPath, events[0].RunPath)\n\t\tassert.True(t, events[0].Action.Interrupted.InterruptContexts[0].\n\t\t\tEqualsWithoutID(secondInterruptEvent.Action.Interrupted.InterruptContexts[0]))\n\t\tinterruptID2 := events[0].Action.Interrupted.InterruptContexts[0].ID\n\t\tevents = []*AgentEvent{}\n\n\t\t// Resume after sa2 interrupt\n\t\titer, err = runner.ResumeWithParams(ctx, \"sequential-1\", &ResumeParams{\n\t\t\tTargets: map[string]any{\n\t\t\t\tinterruptID2: \"resume sa2\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tassert.Equal(t, 2, len(events))\n\t\tassert.Equal(t, messageEvents, events)\n\t})\n\n\tt.Run(\"test loop workflow agent\", func(t *testing.T) {\n\t\t// loop\n\t\ta, err := NewLoopAgent(ctx, &LoopAgentConfig{\n\t\t\tName:          \"loop\",\n\t\t\tSubAgents:     []Agent{sa1, sa2, sa3, sa4},\n\t\t\tMaxIterations: 2,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\trunner := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           a,\n\t\t\tCheckPointStore: newMyStore(),\n\t\t})\n\t\tvar events []*AgentEvent\n\t\titer := runner.Query(ctx, \"hello world\", WithCheckPointID(\"loop-1\"))\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tloopFirstInterruptEvent := &AgentEvent{\n\t\t\tAgentName: \"sa1\",\n\t\t\tRunPath:   []RunStep{{\"loop\"}, {\"sa1\"}},\n\t\t\tAction: &AgentAction{\n\t\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\t\tData: &WorkflowInterruptInfo{\n\t\t\t\t\t\tOrigInput: &AgentInput{\n\t\t\t\t\t\t\tMessages: []Message{schema.UserMessage(\"hello world\")},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tSequentialInterruptIndex: 0,\n\t\t\t\t\t\tSequentialInterruptInfo: &InterruptInfo{\n\t\t\t\t\t\t\tData: \"sa1 interrupt data\",\n\t\t\t\t\t\t},\n\t\t\t\t\t\tLoopIterations: 0,\n\t\t\t\t\t},\n\t\t\t\t\tInterruptContexts: []*InterruptCtx{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID:   \"agent:loop;agent:sa1\",\n\t\t\t\t\t\t\tInfo: \"sa1 interrupt data\",\n\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"loop\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"sa1\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\tIsRootCause: true,\n\t\t\t\t\t\t\tParent: &InterruptCtx{\n\t\t\t\t\t\t\t\tID:   \"agent:loop\",\n\t\t\t\t\t\t\t\tInfo: \"Loop workflow interrupted\",\n\t\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\t\tID:   \"loop\",\n\t\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tassert.Equal(t, 1, len(events))\n\t\tassert.Equal(t, loopFirstInterruptEvent.AgentName, events[0].AgentName)\n\t\tassert.Equal(t, loopFirstInterruptEvent.RunPath, events[0].RunPath)\n\t\tassert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopFirstInterruptEvent.Action.Interrupted.InterruptContexts[0]))\n\t\tloopInterruptID1 := events[0].Action.Interrupted.InterruptContexts[0].ID\n\t\tevents = []*AgentEvent{}\n\n\t\t// Resume after sa1 interrupt\n\t\titer, err = runner.ResumeWithParams(ctx, \"loop-1\", &ResumeParams{\n\t\t\tTargets: map[string]any{\n\t\t\t\tloopInterruptID1: \"resume sa1\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tloopSecondInterruptEvent := &AgentEvent{\n\t\t\tAgentName: \"sa2\",\n\t\t\tRunPath:   []RunStep{{\"loop\"}, {\"sa1\"}, {\"sa2\"}},\n\t\t\tAction: &AgentAction{\n\t\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\t\tData: &WorkflowInterruptInfo{\n\t\t\t\t\t\tOrigInput: &AgentInput{\n\t\t\t\t\t\t\tMessages: []Message{schema.UserMessage(\"hello world\")},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tSequentialInterruptIndex: 1,\n\t\t\t\t\t\tSequentialInterruptInfo: &InterruptInfo{\n\t\t\t\t\t\t\tData: \"sa2 interrupt data\",\n\t\t\t\t\t\t},\n\t\t\t\t\t\tLoopIterations: 0,\n\t\t\t\t\t},\n\t\t\t\t\tInterruptContexts: []*InterruptCtx{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID:   \"agent:loop;agent:sa1;agent:sa2\",\n\t\t\t\t\t\t\tInfo: \"sa2 interrupt data\",\n\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"loop\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"sa2\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\tIsRootCause: true,\n\t\t\t\t\t\t\tParent: &InterruptCtx{\n\t\t\t\t\t\t\t\tID:   \"agent:loop\",\n\t\t\t\t\t\t\t\tInfo: \"Loop workflow interrupted\",\n\t\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\t\tID:   \"loop\",\n\t\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tassert.Equal(t, 1, len(events))\n\t\tassert.Equal(t, loopSecondInterruptEvent.AgentName, events[0].AgentName)\n\t\tassert.Equal(t, loopSecondInterruptEvent.RunPath, events[0].RunPath)\n\t\tassert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopSecondInterruptEvent.Action.Interrupted.InterruptContexts[0]))\n\t\tloopInterruptID2 := events[0].Action.Interrupted.InterruptContexts[0].ID\n\t\tevents = []*AgentEvent{}\n\n\t\t// Resume after sa2 interrupt\n\t\titer, err = runner.ResumeWithParams(ctx, \"loop-1\", &ResumeParams{\n\t\t\tTargets: map[string]any{\n\t\t\t\tloopInterruptID2: \"resume sa2\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\n\t\tloopThirdInterruptEvent := &AgentEvent{\n\t\t\tAgentName: \"sa1\",\n\t\t\tRunPath:   []RunStep{{\"loop\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}, {\"sa4\"}, {\"sa1\"}},\n\t\t\tAction: &AgentAction{\n\t\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\t\tData: &WorkflowInterruptInfo{\n\t\t\t\t\t\tOrigInput: &AgentInput{\n\t\t\t\t\t\t\tMessages: []Message{schema.UserMessage(\"hello world\")},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tSequentialInterruptIndex: 0,\n\t\t\t\t\t\tSequentialInterruptInfo: &InterruptInfo{\n\t\t\t\t\t\t\tData: \"sa1 interrupt data\",\n\t\t\t\t\t\t},\n\t\t\t\t\t\tLoopIterations: 1,\n\t\t\t\t\t},\n\t\t\t\t\tInterruptContexts: []*InterruptCtx{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID:   \"agent:loop;agent:sa1;agent:sa2;agent:sa3;agent:sa4;agent:sa1\",\n\t\t\t\t\t\t\tInfo: \"sa1 interrupt data\",\n\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"loop\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"sa1\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\tIsRootCause: true,\n\t\t\t\t\t\t\tParent: &InterruptCtx{\n\t\t\t\t\t\t\t\tID:   \"agent:loop\",\n\t\t\t\t\t\t\t\tInfo: \"Loop workflow interrupted\",\n\t\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\t\tID:   \"loop\",\n\t\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tloopFourthInterruptEvent := &AgentEvent{\n\t\t\tAgentName: \"sa2\",\n\t\t\tRunPath:   []RunStep{{\"loop\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}, {\"sa4\"}, {\"sa1\"}, {\"sa2\"}},\n\t\t\tAction: &AgentAction{\n\t\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\t\tData: &WorkflowInterruptInfo{\n\t\t\t\t\t\tOrigInput: &AgentInput{\n\t\t\t\t\t\t\tMessages: []Message{schema.UserMessage(\"hello world\")},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tSequentialInterruptIndex: 1,\n\t\t\t\t\t\tSequentialInterruptInfo: &InterruptInfo{\n\t\t\t\t\t\t\tData: \"sa2 interrupt data\",\n\t\t\t\t\t\t},\n\t\t\t\t\t\tLoopIterations: 1,\n\t\t\t\t\t},\n\t\t\t\t\tInterruptContexts: []*InterruptCtx{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID:   \"agent:loop;agent:sa1;agent:sa2;agent:sa3;agent:sa4;agent:sa1;agent:sa2\",\n\t\t\t\t\t\t\tInfo: \"sa2 interrupt data\",\n\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"loop\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID:   \"sa2\",\n\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\tIsRootCause: true,\n\t\t\t\t\t\t\tParent: &InterruptCtx{\n\t\t\t\t\t\t\t\tID:   \"agent:loop\",\n\t\t\t\t\t\t\t\tInfo: \"Loop workflow interrupted\",\n\t\t\t\t\t\t\t\tAddress: Address{\n\t\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\t\tID:   \"loop\",\n\t\t\t\t\t\t\t\t\t\tType: AddressSegmentAgent,\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tloopMessageEvents := []*AgentEvent{\n\t\t\t{\n\t\t\t\tAgentName: \"sa3\",\n\t\t\t\tRunPath:   []RunStep{{\"loop\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}},\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa3 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tAgentName: \"sa4\",\n\t\t\t\tRunPath:   []RunStep{{\"loop\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}, {\"sa4\"}},\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa4 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\tloopThirdInterruptEvent,\n\t\t}\n\t\tassert.Equal(t, 3, len(events))\n\t\t// Check the first two message events\n\t\tassert.Equal(t, loopMessageEvents[0].AgentName, events[0].AgentName)\n\t\tassert.Equal(t, loopMessageEvents[0].RunPath, events[0].RunPath)\n\t\tassert.Equal(t, loopMessageEvents[0].Output.MessageOutput.Message.Content, events[0].Output.MessageOutput.Message.Content)\n\n\t\tassert.Equal(t, loopMessageEvents[1].AgentName, events[1].AgentName)\n\t\tassert.Equal(t, loopMessageEvents[1].RunPath, events[1].RunPath)\n\t\tassert.Equal(t, loopMessageEvents[1].Output.MessageOutput.Message.Content, events[1].Output.MessageOutput.Message.Content)\n\n\t\t// Check the third interrupt event using EqualsWithoutID\n\t\tassert.Equal(t, loopMessageEvents[2].AgentName, events[2].AgentName)\n\t\tassert.Equal(t, loopMessageEvents[2].RunPath, events[2].RunPath)\n\t\tassert.True(t, events[2].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopMessageEvents[2].Action.Interrupted.InterruptContexts[0]))\n\t\tloopInterruptID3 := events[2].Action.Interrupted.InterruptContexts[0].ID\n\t\tevents = []*AgentEvent{}\n\n\t\t// Resume after third interrupt\n\t\titer, err = runner.ResumeWithParams(ctx, \"loop-1\", &ResumeParams{\n\t\t\tTargets: map[string]any{\n\t\t\t\tloopInterruptID3: \"resume sa1\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\t\tassert.Equal(t, 1, len(events))\n\t\tassert.Equal(t, loopFourthInterruptEvent.AgentName, events[0].AgentName)\n\t\tassert.Equal(t, loopFourthInterruptEvent.RunPath, events[0].RunPath)\n\t\tassert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopFourthInterruptEvent.Action.Interrupted.InterruptContexts[0]))\n\t\tloopInterruptID4 := events[0].Action.Interrupted.InterruptContexts[0].ID\n\t\tevents = []*AgentEvent{}\n\n\t\t// Resume after fourth interrupt\n\t\titer, err = runner.ResumeWithParams(ctx, \"loop-1\", &ResumeParams{\n\t\t\tTargets: map[string]any{\n\t\t\t\tloopInterruptID4: \"resume sa2\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\t\tloopFinalMessageEvents := []*AgentEvent{\n\t\t\t{\n\t\t\t\tAgentName: \"sa3\",\n\t\t\t\tRunPath:   []RunStep{{\"loop\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}, {\"sa4\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}},\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa3 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tAgentName: \"sa4\",\n\t\t\t\tRunPath:   []RunStep{{\"loop\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}, {\"sa4\"}, {\"sa1\"}, {\"sa2\"}, {\"sa3\"}, {\"sa4\"}},\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa4 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tassert.Equal(t, 2, len(events))\n\t\tassert.Equal(t, loopFinalMessageEvents, events)\n\t})\n\n\tt.Run(\"test parallel workflow agent\", func(t *testing.T) {\n\t\t// parallel\n\t\ta, err := NewParallelAgent(ctx, &ParallelAgentConfig{\n\t\t\tName:      \"parallel agent\",\n\t\t\tSubAgents: []Agent{sa1, sa2, sa3, sa4},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\trunner := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           a,\n\t\t\tCheckPointStore: newMyStore(),\n\t\t})\n\t\titer := runner.Query(ctx, \"hello world\", WithCheckPointID(\"1\"))\n\t\tvar (\n\t\t\tevents         []*AgentEvent\n\t\t\tinterruptEvent *AgentEvent\n\t\t)\n\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\t\tinterruptEvent = event\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\t\tassert.Equal(t, 2, len(events))\n\n\t\t// Debug: Print actual events to see what we're getting\n\t\tfor i, event := range events {\n\t\t\tt.Logf(\"Event %d: AgentName=%s, RunPath=%v, Output=%v\", i, event.AgentName, event.RunPath, event.Output)\n\t\t}\n\n\t\t// Define parallel message events separately\n\t\tparallelMessageEvents := []*AgentEvent{\n\t\t\t{\n\t\t\t\tAgentName: \"sa4\",\n\t\t\t\tRunPath:   []RunStep{{\"parallel agent\"}, {\"sa4\"}},\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa4 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tAgentName: \"sa3\",\n\t\t\t\tRunPath:   []RunStep{{\"parallel agent\"}, {\"sa3\"}},\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa3 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tassert.Contains(t, events, parallelMessageEvents[0])\n\t\tassert.Contains(t, events, parallelMessageEvents[1])\n\n\t\tassert.NotNil(t, interruptEvent)\n\t\tassert.Equal(t, \"parallel agent\", interruptEvent.AgentName)\n\t\tassert.Equal(t, []RunStep{{\"parallel agent\"}}, interruptEvent.RunPath)\n\t\tassert.NotNil(t, interruptEvent.Action.Interrupted)\n\t\twii, ok := interruptEvent.Action.Interrupted.Data.(*WorkflowInterruptInfo)\n\t\tassert.True(t, ok)\n\t\tassert.Equal(t, 2, len(wii.ParallelInterruptInfo))\n\n\t\tvar sa1Found, sa2Found bool\n\t\tfor _, info := range wii.ParallelInterruptInfo {\n\t\t\tswitch info.Data {\n\t\t\tcase \"sa1 interrupt data\":\n\t\t\t\tsa1Found = true\n\t\t\tcase \"sa2 interrupt data\":\n\t\t\t\tsa2Found = true\n\t\t\t}\n\t\t}\n\t\tassert.True(t, sa1Found)\n\t\tassert.True(t, sa2Found)\n\n\t\tvar sa1InfoFound, sa2InfoFound bool\n\t\tfor _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\t\tif ctx.Info == \"sa1 interrupt data\" {\n\t\t\t\tsa1InfoFound = true\n\t\t\t} else if ctx.Info == \"sa2 interrupt data\" {\n\t\t\t\tsa2InfoFound = true\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts))\n\t\tassert.True(t, sa1InfoFound)\n\t\tassert.True(t, sa2InfoFound)\n\n\t\tvar parallelInterruptID1, parallelInterruptID2 string\n\t\tfor _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\t\tif ctx.Info == \"sa1 interrupt data\" {\n\t\t\t\tparallelInterruptID1 = ctx.ID\n\t\t\t} else if ctx.Info == \"sa2 interrupt data\" {\n\t\t\t\tparallelInterruptID2 = ctx.ID\n\t\t\t}\n\t\t}\n\t\tassert.NotEmpty(t, parallelInterruptID1)\n\t\tassert.NotEmpty(t, parallelInterruptID2)\n\n\t\titer, err = runner.ResumeWithParams(ctx, \"1\", &ResumeParams{\n\t\t\tTargets: map[string]any{\n\t\t\t\tparallelInterruptID1: \"resume sa1\",\n\t\t\t\tparallelInterruptID2: \"resume sa2\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\t_, ok = iter.Next()\n\t\tassert.False(t, ok)\n\t})\n}\n\nfunc TestChatModelInterrupt(t *testing.T) {\n\tctx := context.Background()\n\ta, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"name\",\n\t\tDescription: \"description\",\n\t\tInstruction: \"instruction\",\n\t\tModel: &myModel{\n\t\t\tvalidator: func(i int, messages []*schema.Message) bool {\n\t\t\t\tif i > 0 && (len(messages) != 4 || messages[2].Content != \"new user message\") {\n\t\t\t\t\treturn false\n\t\t\t\t}\n\t\t\t\treturn true\n\t\t\t},\n\t\t\tmessages: []*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      \"tool1\",\n\t\t\t\t\t\t\tArguments: \"arguments\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t\tschema.AssistantMessage(\"completed\", nil),\n\t\t\t},\n\t\t},\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{&myTool1{}},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           a,\n\t\tCheckPointStore: newMyStore(),\n\t})\n\titer := runner.Query(ctx, \"hello world\", WithCheckPointID(\"1\"))\n\tevent, ok := iter.Next()\n\tassert.True(t, ok)\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.NoError(t, event.Err)\n\tassert.NotNil(t, event.Action.Interrupted)\n\tassert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))\n\tassert.Equal(t, Address{\n\t\t{Type: AddressSegmentAgent, ID: \"name\"},\n\t\t{Type: AddressSegmentTool, ID: \"tool1\", SubID: \"1\"},\n\t}, event.Action.Interrupted.InterruptContexts[0].Address)\n\n\tvar (\n\t\tchatModelAgentID string\n\t\ttoolID           string\n\t)\n\n\tintCtx := event.Action.Interrupted.InterruptContexts[0]\n\tfor intCtx != nil {\n\t\tif intCtx.Address[len(intCtx.Address)-1].Type == AddressSegmentTool {\n\t\t\ttoolID = intCtx.ID\n\t\t} else if intCtx.Address[len(intCtx.Address)-1].Type == AddressSegmentAgent {\n\t\t\tchatModelAgentID = intCtx.ID\n\t\t}\n\t\tintCtx = intCtx.Parent\n\t}\n\n\tevent, ok = iter.Next()\n\tassert.False(t, ok)\n\n\titer, err = runner.ResumeWithParams(ctx, \"1\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\tchatModelAgentID: &ChatModelAgentResumeData{\n\t\t\t\tHistoryModifier: func(ctx context.Context, history []Message) []Message {\n\t\t\t\t\thistory[2].Content = \"new user message\"\n\t\t\t\t\treturn history\n\t\t\t\t},\n\t\t\t},\n\t\t\ttoolID: \"tool resume result\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.NoError(t, event.Err)\n\tassert.Equal(t, event.Output.MessageOutput.Message.Content, \"tool resume result\")\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.NoError(t, event.Err)\n\tassert.Equal(t, event.Output.MessageOutput.Message.Content, \"completed\")\n}\n\nfunc TestChatModelAgentToolInterrupt(t *testing.T) {\n\tsa := &myAgent{\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tintAct := Interrupt(ctx, \"hello world\")\n\t\t\tintAct.Action.Interrupted.Data = \"hello world\"\n\t\t\tgenerator.Send(intAct)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.NotNil(t, info)\n\t\t\tassert.False(t, info.EnableStreaming)\n\n\t\t\tif !info.IsResumeTarget {\n\t\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\t\tintAct := Interrupt(ctx, \"interrupt again\")\n\t\t\t\tintAct.Action.Interrupted.Data = \"interrupt again\"\n\t\t\t\tgenerator.Send(intAct)\n\t\t\t\tgenerator.Close()\n\t\t\t\treturn iter\n\t\t\t}\n\n\t\t\tassert.NotNil(t, info.ResumeData)\n\t\t\tassert.Equal(t, \"resume sa\", info.ResumeData)\n\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.UserMessage(fmt.Sprintf(\"my agent completed with data %s\", info.ResumeData))}}})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\tctx := context.Background()\n\ta, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"name\",\n\t\tDescription: \"description\",\n\t\tInstruction: \"instruction\",\n\t\tModel: &myModel{\n\t\t\tmessages: []*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      \"myAgent\",\n\t\t\t\t\t\t\tArguments: \"{\\\"request\\\":\\\"123\\\"}\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t\tschema.AssistantMessage(\"completed\", nil),\n\t\t\t},\n\t\t},\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{NewAgentTool(ctx, sa)},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           a,\n\t\tCheckPointStore: newMyStore(),\n\t})\n\n\titer := runner.Query(ctx, \"hello world\", WithCheckPointID(\"1\"))\n\tevent, ok := iter.Next()\n\tassert.True(t, ok)\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.NoError(t, event.Err)\n\tassert.NotNil(t, event.Action.Interrupted)\n\tevent, ok = iter.Next()\n\tassert.False(t, ok)\n\n\titer, err = runner.Resume(ctx, \"1\")\n\tassert.NoError(t, err)\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.NoError(t, event.Err)\n\tassert.NotNil(t, event.Action.Interrupted)\n\tassert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))\n\tfor _, ctx := range event.Action.Interrupted.InterruptContexts {\n\t\tif ctx.IsRootCause {\n\t\t\tassert.Equal(t, Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"name\"},\n\t\t\t\t{Type: AddressSegmentTool, ID: \"myAgent\", SubID: \"1\"},\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"myAgent\"},\n\t\t\t}, ctx.Address)\n\t\t\tassert.Equal(t, \"interrupt again\", ctx.Info)\n\t\t}\n\t}\n\n\tvar toolInterruptID string\n\tfor _, ctx := range event.Action.Interrupted.InterruptContexts {\n\t\tif ctx.IsRootCause {\n\t\t\ttoolInterruptID = ctx.ID\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.NotEmpty(t, toolInterruptID)\n\n\tevent, ok = iter.Next()\n\tassert.False(t, ok)\n\n\titer, err = runner.ResumeWithParams(ctx, \"1\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\ttoolInterruptID: \"resume sa\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.NoError(t, event.Err)\n\tassert.Equal(t, event.Output.MessageOutput.Message.Content, \"my agent completed with data resume sa\")\n\tevent, ok = iter.Next()\n\tassert.True(t, ok)\n\tassert.NoError(t, event.Err)\n\tassert.Equal(t, event.Output.MessageOutput.Message.Content, \"completed\")\n\t_, ok = iter.Next()\n\tassert.False(t, ok)\n}\n\nfunc newMyStore() *myStore {\n\treturn &myStore{\n\t\tm: map[string][]byte{},\n\t}\n}\n\ntype myStore struct {\n\tm map[string][]byte\n}\n\nfunc (m *myStore) Set(_ context.Context, key string, value []byte) error {\n\tm.m[key] = value\n\treturn nil\n}\n\nfunc (m *myStore) Get(_ context.Context, key string) ([]byte, bool, error) {\n\tv, ok := m.m[key]\n\treturn v, ok, nil\n}\n\ntype myAgentOptions struct {\n\tinterrupt bool\n\n\tvalue string\n}\n\nfunc withValue(value string) AgentRunOption {\n\treturn WrapImplSpecificOptFn(func(t *myAgentOptions) {\n\t\tt.value = value\n\t})\n}\n\ntype myAgent struct {\n\tname     string\n\trunFn    func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]\n\tresumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]\n}\n\nfunc (m *myAgent) Name(_ context.Context) string {\n\tif len(m.name) > 0 {\n\t\treturn m.name\n\t}\n\treturn \"myAgent\"\n}\n\nfunc (m *myAgent) Description(_ context.Context) string {\n\treturn \"myAgent description\"\n}\n\nfunc (m *myAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\treturn m.runFn(ctx, input, options...)\n}\n\nfunc (m *myAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\treturn m.resumeFn(ctx, info, opts...)\n}\n\ntype myModel struct {\n\ttimes     int\n\tmessages  []*schema.Message\n\tvalidator func(int, []*schema.Message) bool\n}\n\nfunc (m *myModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) {\n\tif m.validator != nil && !m.validator(m.times, input) {\n\t\treturn nil, errors.New(\"invalid input\")\n\t}\n\tif m.times >= len(m.messages) {\n\t\treturn nil, errors.New(\"exceeded max number of messages\")\n\t}\n\tt := m.times\n\tm.times++\n\treturn m.messages[t], nil\n}\n\nfunc (m *myModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tpanic(\"implement me\")\n}\n\nfunc (m *myModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\treturn m, nil\n}\n\ntype myTool1 struct{}\n\nfunc (m *myTool1) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"tool1\",\n\t\tDesc: \"desc\",\n\t}, nil\n}\n\nfunc (m *myTool1) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {\n\tif wasInterrupted, _, _ := tool.GetInterruptState[any](ctx); !wasInterrupted {\n\t\treturn \"\", tool.Interrupt(ctx, nil)\n\t}\n\n\tif isResumeFlow, hasResumeData, data := tool.GetResumeContext[string](ctx); !isResumeFlow {\n\t\treturn \"\", tool.Interrupt(ctx, nil)\n\t} else if hasResumeData {\n\t\treturn data, nil\n\t}\n\n\treturn \"result\", nil\n}\n\nfunc TestCyclicalAgentInterrupt(t *testing.T) {\n\tctx := context.Background()\n\n\tvar agentA, agentB, agentC Agent\n\n\t// agentC interrupts\n\tagentC = &myAgent{\n\t\tname: \"C\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tintAct := Interrupt(ctx, \"interrupt from C\")\n\t\t\tgenerator.Send(intAct)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.True(t, info.IsResumeTarget)\n\t\t\tassert.NotNil(t, info.ResumeData)\n\t\t\tassert.Equal(t, \"resume C\", info.ResumeData)\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"C\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{Message: schema.UserMessage(\"C completed\")},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\t// agentB transfers back to its parent A\n\tagentB = &myAgent{\n\t\tname: \"B\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"B\",\n\t\t\t\tAction: &AgentAction{\n\t\t\t\t\tTransferToAgent: &TransferToAgentAction{\n\t\t\t\t\t\tDestAgentName: \"A\", // Transfer back to parent\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\t// agentA is the parent, orchestrating the A->B->A->C flow\n\tagentA = &myAgent{\n\t\tname: \"A\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\t\t\t// If the last agent was B, we are in the A->B->A path, so transfer to C.\n\t\t\t// Otherwise, it's the first run, transfer to B.\n\t\t\tdest := \"B\"\n\t\t\tif len(runCtx.RunPath) > 1 && runCtx.RunPath[len(runCtx.RunPath)-2].agentName == \"B\" {\n\t\t\t\tdest = \"C\"\n\t\t\t}\n\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"A\",\n\t\t\t\tAction: &AgentAction{\n\t\t\t\t\tTransferToAgent: &TransferToAgentAction{\n\t\t\t\t\t\tDestAgentName: dest,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\t// Set up the hierarchy: A is parent of B and C.\n\tagentA, err := SetSubAgents(ctx, agentA, []Agent{agentB, agentC})\n\tassert.NoError(t, err)\n\n\t// Run the test\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           agentA,\n\t\tCheckPointStore: newMyStore(),\n\t})\n\titer := runner.Query(ctx, \"start\", WithCheckPointID(\"cyclical-1\"))\n\n\tvar events []*AgentEvent\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\t// We expect 3 transfer events (A->B, B->A, A->C) and 1 interrupt event from C.\n\tassert.Equal(t, 4, len(events))\n\n\tinterruptEvent := events[3]\n\tassert.NotNil(t, interruptEvent.Action.Interrupted)\n\tassert.Equal(t, \"C\", interruptEvent.AgentName)\n\n\t// Check the interrupt context\n\tassert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts))\n\tinterruptCtx := interruptEvent.Action.Interrupted.InterruptContexts[0]\n\tassert.True(t, interruptCtx.IsRootCause)\n\tassert.Equal(t, \"interrupt from C\", interruptCtx.Info)\n\n\texpectedAddr := Address{\n\t\t{Type: AddressSegmentAgent, ID: \"A\"},\n\t\t{Type: AddressSegmentAgent, ID: \"B\"},\n\t\t{Type: AddressSegmentAgent, ID: \"A\"},\n\t\t{Type: AddressSegmentAgent, ID: \"C\"},\n\t}\n\tassert.Equal(t, expectedAddr, interruptCtx.Address)\n\tassert.NotEmpty(t, interruptCtx.ID)\n\n\t// Resume the execution\n\titer, err = runner.ResumeWithParams(ctx, \"cyclical-1\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\tinterruptCtx.ID: \"resume C\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tevents = []*AgentEvent{}\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\t// We expect one output event from C\n\tassert.Equal(t, 1, len(events))\n\tassert.Equal(t, \"C completed\", events[0].Output.MessageOutput.Message.Content)\n}\n\n// myStatefulTool is a tool that can interrupt and has internal state to track invocations.\n\ntype myStatefulTool struct {\n\tname string\n\tt    *testing.T\n}\n\nfunc (m *myStatefulTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: m.name,\n\t\tDesc: \"desc\",\n\t}, nil\n}\n\ntype myStatefulToolState struct {\n\tInterruptCount int\n}\n\nfunc init() {\n\tschema.Register[myStatefulToolState]()\n}\n\nfunc (m *myStatefulTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {\n\twasInterrupted, hasState, state := tool.GetInterruptState[myStatefulToolState](ctx)\n\tif !wasInterrupted {\n\t\treturn \"\", tool.StatefulInterrupt(ctx, fmt.Sprintf(\"interrupt from %s\", m.name), myStatefulToolState{InterruptCount: 1})\n\t}\n\n\tisResumeFlow, hasResumeData, data := tool.GetResumeContext[string](ctx)\n\tif !isResumeFlow || !hasResumeData {\n\t\tassert.True(m.t, hasState, \"tool %s should have interrupt state on resume\", m.name)\n\t\treturn \"\", tool.StatefulInterrupt(ctx, fmt.Sprintf(\"interrupt from %s\", m.name), myStatefulToolState{InterruptCount: state.InterruptCount + 1})\n\t}\n\n\treturn data, nil\n}\n\nfunc TestChatModelParallelToolInterruptAndResume(t *testing.T) {\n\tctx := context.Background()\n\n\ttoolA := &myStatefulTool{name: \"toolA\", t: t}\n\ttoolB := &myStatefulTool{name: \"toolB\", t: t}\n\n\tchatModel, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"ParallelToolAgent\",\n\t\tDescription: \"An agent that uses parallel tools\",\n\t\tModel: &myModel{\n\t\t\tmessages: []*schema.Message{\n\t\t\t\t// 1. First model response: call toolA and toolB in parallel\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{ID: \"1\", Function: schema.FunctionCall{Name: \"toolA\", Arguments: \"{}\"}},\n\t\t\t\t\t{ID: \"2\", Function: schema.FunctionCall{Name: \"toolB\", Arguments: \"{}\"}},\n\t\t\t\t}),\n\t\t\t\t// 2. Second model response (after tools are resumed): call them again to check state\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{ID: \"3\", Function: schema.FunctionCall{Name: \"toolA\", Arguments: \"{}\"}},\n\t\t\t\t\t{ID: \"4\", Function: schema.FunctionCall{Name: \"toolB\", Arguments: \"{}\"}},\n\t\t\t\t}),\n\t\t\t\t// 3. Final completion\n\t\t\t\tschema.AssistantMessage(\"all done\", nil),\n\t\t\t},\n\t\t},\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{toolA, toolB},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           chatModel,\n\t\tCheckPointStore: newMyStore(),\n\t})\n\n\t// 1. Initial query -> parallel interrupt from toolA and toolB\n\titer := runner.Query(ctx, \"start\", WithCheckPointID(\"parallel-tool-test-1\"))\n\tnormalEvents, interruptEvent := consumeUntilInterrupt(iter)\n\n\tassert.Equal(t, 1, len(normalEvents))\n\tassert.NotNil(t, interruptEvent)\n\tassert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts),\n\t\t\"should have 2 interrupts\")\n\n\tvar toolAInterruptID, toolBInterruptID string\n\tfor _, info := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif info.Info == \"interrupt from toolA\" {\n\t\t\ttoolAInterruptID = info.ID\n\t\t\tassert.True(t, info.IsRootCause)\n\t\t} else if info.Info == \"interrupt from toolB\" {\n\t\t\ttoolBInterruptID = info.ID\n\t\t\tassert.True(t, info.IsRootCause)\n\t\t}\n\t}\n\tassert.NotEmpty(t, toolAInterruptID)\n\tassert.NotEmpty(t, toolBInterruptID)\n\n\t// 2. Resume, targeting only toolA. toolB should re-interrupt.\n\titer, err = runner.ResumeWithParams(ctx, \"parallel-tool-test-1\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\ttoolAInterruptID: \"toolA resumed\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\t_, interruptEvent = consumeUntilInterrupt(iter)\n\n\tassert.NotNil(t, interruptEvent, \"expected a re-interrupt from toolB\")\n\tassert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts),\n\t\t\"should have 1 remaining interrupts\")\n\n\tvar rootCause *InterruptCtx\n\tfor _, info := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif info.IsRootCause {\n\t\t\trootCause = info\n\t\t\tbreak\n\t\t}\n\t}\n\n\tif rootCause == nil {\n\t\tt.Fatal(\"expected a root cause interrupt from toolB\")\n\t}\n\tassert.Equal(t, \"interrupt from toolB\", rootCause.Info)\n\ttoolBReInterruptID := rootCause.ID\n\n\t// 3. Resume the re-interrupted toolB. The agent should then call the tools again.\n\titer, err = runner.ResumeWithParams(ctx, \"parallel-tool-test-1\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\ttoolBReInterruptID: \"toolB resumed\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\t// 4. Consume all final events. The internal assertions in the tools will check the wasInterrupted flag.\n\t// We expect to see the results of the second tool calls, and then the final agent completion.\n\tfinalEvents, interruptEvent := consumeUntilInterrupt(iter)\n\tassert.Equal(t, 2, len(finalEvents))\n\tassert.NotNil(t, interruptEvent)\n}\n\n// TestNestedChatModelAgentWithAgentTool verifies that the shouldFire method correctly prevents\n// duplicate event firing in nested ChatModelAgent scenarios (ChatModelAgent -> AgentTool -> ChatModelAgent).\n// This ensures that only the inner agent's cbHandler fires, not the outer agent's.\nfunc TestNestedChatModelAgentWithAgentTool(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create an interruptible tool for the inner agent\n\tinnerTool := &myStatefulTool{name: \"innerTool\", t: t}\n\n\t// Create the inner ChatModelAgent that will be wrapped by AgentTool\n\tinnerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"InnerAgent\",\n\t\tDescription: \"Inner agent with interruptible tool\",\n\t\tModel: &myModel{\n\t\t\tmessages: []*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{ID: \"1\", Function: schema.FunctionCall{Name: \"innerTool\", Arguments: \"{}\"}},\n\t\t\t\t}),\n\t\t\t\tschema.AssistantMessage(\"inner agent completed\", nil),\n\t\t\t},\n\t\t},\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{innerTool},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\t// Wrap the inner agent in an AgentTool\n\tagentTool := NewAgentTool(ctx, innerAgent)\n\n\t// Create the outer ChatModelAgent that uses the AgentTool\n\touterAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"OuterAgent\",\n\t\tDescription: \"Outer agent with AgentTool containing inner agent\",\n\t\tModel: &myModel{\n\t\t\tmessages: []*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{ID: \"1\", Function: schema.FunctionCall{Name: \"InnerAgent\", Arguments: \"{}\"}},\n\t\t\t\t}),\n\t\t\t\tschema.AssistantMessage(\"outer agent completed\", nil),\n\t\t\t},\n\t\t},\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{agentTool},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           outerAgent,\n\t\tCheckPointStore: newMyStore(),\n\t})\n\n\t// Run the query - this should trigger the nested agent structure\n\titer := runner.Query(ctx, \"start\", WithCheckPointID(\"nested-agent-test-1\"))\n\n\t// Collect all events to verify no duplicates\n\tvar allEvents []*AgentEvent\n\tvar interruptEvent *AgentEvent\n\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tassert.Nil(t, interruptEvent)\n\t\t\tinterruptEvent = event\n\t\t}\n\n\t\tallEvents = append(allEvents, event)\n\t}\n\n\tif interruptEvent == nil {\n\t\tt.Fatal(\"expected an interrupt event\")\n\t}\n\n\t// Verify we got exactly one interrupt event (not duplicated)\n\tassert.NotNil(t, interruptEvent, \"should have an interrupt event\")\n\tassert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts),\n\t\t\"should have exactly one interrupt context\")\n\n\t// Verify the interrupt comes from the inner tool, not duplicated\n\tinterruptCtx := interruptEvent.Action.Interrupted.InterruptContexts[0]\n\tassert.True(t, interruptCtx.IsRootCause, \"interrupt should be root cause\")\n\tassert.Equal(t, \"interrupt from innerTool\", interruptCtx.Info)\n\n\t// Verify the address path shows the correct nested structure\n\texpectedAddress := Address{\n\t\t{Type: AddressSegmentAgent, ID: \"OuterAgent\"},\n\t\t{Type: AddressSegmentTool, ID: \"InnerAgent\", SubID: \"1\"},\n\t\t{Type: AddressSegmentAgent, ID: \"InnerAgent\"},\n\t\t{Type: AddressSegmentTool, ID: \"innerTool\", SubID: \"1\"},\n\t}\n\tassert.Equal(t, expectedAddress, interruptCtx.Address,\n\t\t\"interrupt address should show correct nested structure\")\n\n\t// Verify no duplicate events by checking agent names in events\n\tvar agentNames []string\n\tfor _, event := range allEvents {\n\t\tif event.AgentName != \"\" {\n\t\t\tagentNames = append(agentNames, event.AgentName)\n\t\t}\n\t}\n\n\t// Should only have events from the outer agent (the inner agent's events should be handled\n\t// by the AgentTool and not duplicated by the outer agent's cbHandler)\n\tfor _, name := range agentNames {\n\t\tassert.Equal(t, \"OuterAgent\", name,\n\t\t\t\"all events should come from OuterAgent, not duplicated from InnerAgent\")\n\t}\n\n\t// Now resume the interrupt\n\tinterruptID := interruptCtx.ID\n\titer, err = runner.ResumeWithParams(ctx, \"nested-agent-test-1\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\tinterruptID: \"resume inner tool\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\t// Collect final events after resume\n\tvar finalEvents []*AgentEvent\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tfinalEvents = append(finalEvents, event)\n\t}\n\n\t// Verify completion events\n\tassert.Greater(t, len(finalEvents), 0, \"should have completion events after resume\")\n\n\t// Check that we get the expected completion messages\n\tvar foundInnerCompletion, foundOuterCompletion bool\n\tfor _, event := range finalEvents {\n\t\tif event.Output != nil && event.Output.MessageOutput != nil {\n\t\t\tif event.Output.MessageOutput.Message != nil {\n\t\t\t\tcontent := event.Output.MessageOutput.Message.Content\n\t\t\t\tif content == \"inner agent completed\" {\n\t\t\t\t\tfoundInnerCompletion = true\n\t\t\t\t} else if content == \"outer agent completed\" {\n\t\t\t\t\tfoundOuterCompletion = true\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tassert.True(t, foundInnerCompletion, \"should have inner agent completion\")\n\tassert.True(t, foundOuterCompletion, \"should have outer agent completion\")\n}\n\n// consumeUntilInterrupt consumes events from the iterator until an interrupt is found or it's exhausted.\nfunc consumeUntilInterrupt(iter *AsyncIterator[*AgentEvent]) (normalEvents []*AgentEvent, interruptEvent *AgentEvent) {\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tinterruptEvent = event\n\t\t\tcontinue\n\t\t}\n\t\tnormalEvents = append(normalEvents, event)\n\t}\n\treturn\n}\n\ntype returnDirectlyTool struct {\n\tname string\n}\n\nfunc (t *returnDirectlyTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: t.name,\n\t\tDesc: \"A tool that returns directly\",\n\t}, nil\n}\n\nfunc (t *returnDirectlyTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn \"return directly result\", nil\n}\n\ntype interruptingTool struct {\n\tname string\n}\n\nfunc (i *interruptingTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: i.name,\n\t\tDesc: \"A tool that interrupts\",\n\t}, nil\n}\n\nfunc (i *interruptingTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {\n\tif wasInterrupted, _, _ := compose.GetInterruptState[any](ctx); !wasInterrupted {\n\t\treturn \"\", compose.Interrupt(ctx, \"interrupt data\")\n\t}\n\n\tif isResumeFlow, hasResumeData, data := compose.GetResumeContext[string](ctx); isResumeFlow && hasResumeData {\n\t\treturn data, nil\n\t}\n\n\treturn \"resumed without data\", nil\n}\n\ntype twoToolCallModel struct {\n\treturnDirectlyToolName string\n\tinterruptingToolName   string\n\tcallCount              int\n\treceivedTools          []*schema.ToolInfo\n\tmu                     sync.Mutex\n}\n\nfunc (m *twoToolCallModel) Generate(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tm.mu.Lock()\n\tm.callCount++\n\tcallNum := m.callCount\n\toptions := model.GetCommonOptions(&model.Options{}, opts...)\n\tif options.Tools != nil {\n\t\tm.receivedTools = options.Tools\n\t}\n\tm.mu.Unlock()\n\n\tif callNum == 1 {\n\t\treturn &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"\",\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID:   \"call_return_directly\",\n\t\t\t\t\tType: \"function\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      m.returnDirectlyToolName,\n\t\t\t\t\t\tArguments: \"{}\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tID:   \"call_interrupting\",\n\t\t\t\t\tType: \"function\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      m.interruptingToolName,\n\t\t\t\t\t\tArguments: \"{}\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, nil\n\t}\n\treturn schema.AssistantMessage(\"final response\", nil), nil\n}\n\nfunc (m *twoToolCallModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tpanic(\"not implemented\")\n}\n\nfunc (m *twoToolCallModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\treturn m, nil\n}\n\nfunc (m *twoToolCallModel) GetReceivedTools() []*schema.ToolInfo {\n\tm.mu.Lock()\n\tdefer m.mu.Unlock()\n\treturn m.receivedTools\n}\n\ntype dynamicTool struct {\n\tname string\n}\n\nfunc (t *dynamicTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: t.name,\n\t\tDesc: \"A dynamically added tool\",\n\t}, nil\n}\n\nfunc (t *dynamicTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn \"dynamic tool result\", nil\n}\n\nfunc TestReturnDirectlyEventSentAfterResume(t *testing.T) {\n\tctx := context.Background()\n\n\treturnDirectlyToolName := \"return_directly_tool\"\n\tinterruptingToolName := \"interrupting_tool\"\n\tdynamicToolName := \"dynamic_tool\"\n\n\tmdl := &twoToolCallModel{\n\t\treturnDirectlyToolName: returnDirectlyToolName,\n\t\tinterruptingToolName:   interruptingToolName,\n\t}\n\n\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\tName:        \"TestAgent\",\n\t\tDescription: \"Test agent for return directly + interrupt\",\n\t\tModel:       mdl,\n\t\tToolsConfig: ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{\n\t\t\t\t\t&returnDirectlyTool{name: returnDirectlyToolName},\n\t\t\t\t\t&interruptingTool{name: interruptingToolName},\n\t\t\t\t},\n\t\t\t},\n\t\t\tReturnDirectly: map[string]bool{\n\t\t\t\treturnDirectlyToolName: true,\n\t\t\t},\n\t\t},\n\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t&interruptTestToolsHandler{tools: []tool.BaseTool{&dynamicTool{name: dynamicToolName}}},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tstore := newMyStore()\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           agent,\n\t\tEnableStreaming: false,\n\t\tCheckPointStore: store,\n\t})\n\n\titer := runner.Query(ctx, \"test input\", WithCheckPointID(\"test_checkpoint\"))\n\n\tvar interruptEvent *AgentEvent\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tinterruptEvent = event\n\t\t}\n\t}\n\n\tassert.NotNil(t, interruptEvent, \"Should have an interrupt event\")\n\tassert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)\n\n\treceivedToolsBeforeResume := mdl.GetReceivedTools()\n\tvar hasDynamicToolBeforeResume bool\n\tfor _, ti := range receivedToolsBeforeResume {\n\t\tif ti.Name == dynamicToolName {\n\t\t\thasDynamicToolBeforeResume = true\n\t\t}\n\t}\n\tassert.True(t, hasDynamicToolBeforeResume, \"Dynamic tool should be in tool list before interrupt\")\n\n\tinterruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID\n\tresumeIter, err := runner.ResumeWithParams(ctx, \"test_checkpoint\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\tinterruptID: \"resume data\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tvar resumeEvents []*AgentEvent\n\tfor {\n\t\tevent, ok := resumeIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tresumeEvents = append(resumeEvents, event)\n\t}\n\n\tvar hasReturnDirectlyEvent bool\n\tfor _, e := range resumeEvents {\n\t\tif e.Output != nil && e.Output.MessageOutput != nil {\n\t\t\tif e.Output.MessageOutput.Role == schema.Tool && e.Output.MessageOutput.ToolName == returnDirectlyToolName {\n\t\t\t\thasReturnDirectlyEvent = true\n\t\t\t}\n\t\t}\n\t}\n\tassert.True(t, hasReturnDirectlyEvent, \"ReturnDirectlyEvent should be sent after resume\")\n\n\treceivedToolsAfterResume := mdl.GetReceivedTools()\n\tvar hasDynamicToolAfterResume bool\n\tfor _, ti := range receivedToolsAfterResume {\n\t\tif ti.Name == dynamicToolName {\n\t\t\thasDynamicToolAfterResume = true\n\t\t}\n\t}\n\tassert.True(t, hasDynamicToolAfterResume, \"Dynamic tool should be in tool list after resume (bc.toolUpdated path)\")\n}\n"
  },
  {
    "path": "adk/middlewares/dynamictool/toolsearch/toolsearch.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package toolsearch provides tool search middleware.\npackage toolsearch\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"regexp\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Config is the configuration for the tool search middleware.\ntype Config struct {\n\t// DynamicTools is a list of tools that can be dynamically searched and loaded by the agent.\n\tDynamicTools []tool.BaseTool\n}\n\n// New constructs and returns the tool search middleware.\n//\n// The tool search middleware enables dynamic tool selection for agents with large tool libraries.\n// Instead of passing all tools to the model at once (which can overwhelm context limits),\n// this middleware:\n//\n//  1. Adds a \"tool_search\" meta-tool that accepts a regex pattern to search tool names\n//  2. Initially hides all dynamic tools from the model's tool list\n//  3. When the model calls tool_search, matching tools become available for subsequent calls\n//\n// Example usage:\n//\n//\tmiddleware, _ := toolsearch.New(ctx, &toolsearch.Config{\n//\t    DynamicTools: []tool.BaseTool{weatherTool, stockTool, currencyTool, ...},\n//\t})\n//\tagent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n//\t    // ...\n//\t    Handlers: []adk.ChatModelAgentMiddleware{middleware},\n//\t})\nfunc New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) {\n\tif config == nil {\n\t\treturn nil, fmt.Errorf(\"config is required\")\n\t}\n\tif len(config.DynamicTools) == 0 {\n\t\treturn nil, fmt.Errorf(\"tools is required\")\n\t}\n\n\treturn &middleware{\n\t\tdynamicTools: config.DynamicTools,\n\t}, nil\n}\n\ntype middleware struct {\n\tadk.BaseChatModelAgentMiddleware\n\tdynamicTools []tool.BaseTool\n}\n\nfunc (m *middleware) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) {\n\tif runCtx == nil {\n\t\treturn ctx, runCtx, nil\n\t}\n\n\tnRunCtx := *runCtx\n\ttoolNames, err := getToolNames(ctx, m.dynamicTools)\n\tif err != nil {\n\t\treturn ctx, nil, fmt.Errorf(\"failed to get tool names: %w\", err)\n\t}\n\tnRunCtx.Tools = append(nRunCtx.Tools, newToolSearchTool(toolNames))\n\tnRunCtx.Tools = append(nRunCtx.Tools, m.dynamicTools...)\n\treturn ctx, &nRunCtx, nil\n}\n\nfunc (m *middleware) WrapModel(_ context.Context, cm model.BaseChatModel, mc *adk.ModelContext) (model.BaseChatModel, error) {\n\treturn &wrapper{allTools: mc.Tools, cm: cm, dynamicTools: m.dynamicTools}, nil\n}\n\ntype wrapper struct {\n\tallTools     []*schema.ToolInfo\n\tdynamicTools []tool.BaseTool\n\n\tcm model.BaseChatModel\n}\n\nfunc (w *wrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\ttools, err := removeTools(ctx, w.allTools, w.dynamicTools, input)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to load dynamic tools: %w\", err)\n\t}\n\treturn w.cm.Generate(ctx, input, append(opts, model.WithTools(tools))...)\n}\n\nfunc (w *wrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\ttools, err := removeTools(ctx, w.allTools, w.dynamicTools, input)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to load dynamic tools: %w\", err)\n\t}\n\treturn w.cm.Stream(ctx, input, append(opts, model.WithTools(tools))...)\n}\n\nfunc newToolSearchTool(toolNames []string) *toolSearchTool {\n\treturn &toolSearchTool{toolNames: toolNames}\n}\n\ntype toolSearchTool struct {\n\ttoolNames []string\n}\n\nconst (\n\ttoolSearchToolName = \"tool_search\"\n)\n\nfunc (t *toolSearchTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"tool_search\",\n\t\tDesc: \"Search for tools using a regex pattern that matches tool names. Returns a list of matching tool names. Use this when you need a tool but don't have it available yet.\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"regex_pattern\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"A regex pattern to match tool names against.\",\n\t\t\t\tRequired: true,\n\t\t\t},\n\t\t}),\n\t}, nil\n}\n\ntype toolSearchArgs struct {\n\tRegexPattern string `json:\"regex_pattern\"`\n}\n\ntype toolSearchResult struct {\n\tSelectedTools []string `json:\"selectedTools\"`\n}\n\nfunc (t *toolSearchTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tvar args toolSearchArgs\n\tif err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to unmarshal tool search arguments: %w\", err)\n\t}\n\n\tif args.RegexPattern == \"\" {\n\t\treturn \"\", fmt.Errorf(\"regex_pattern is required\")\n\t}\n\n\tre, err := regexp.Compile(args.RegexPattern)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"invalid regex pattern: %w\", err)\n\t}\n\n\tvar matchedTools []string\n\tfor _, name := range t.toolNames {\n\t\tif re.MatchString(name) {\n\t\t\tmatchedTools = append(matchedTools, name)\n\t\t}\n\t}\n\n\tresult := toolSearchResult{\n\t\tSelectedTools: matchedTools,\n\t}\n\n\toutput, err := json.Marshal(result)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to marshal result: %w\", err)\n\t}\n\n\treturn string(output), nil\n}\n\nfunc getToolNames(ctx context.Context, tools []tool.BaseTool) ([]string, error) {\n\tret := make([]string, 0, len(tools))\n\tfor _, t := range tools {\n\t\tinfo, err := t.Info(ctx)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tret = append(ret, info.Name)\n\t}\n\treturn ret, nil\n}\n\nfunc extractSelectedTools(ctx context.Context, messages []*schema.Message) ([]string, error) {\n\tvar selectedTools []string\n\tfor _, message := range messages {\n\t\tif message.Role != schema.Tool || message.ToolName != toolSearchToolName {\n\t\t\tcontinue\n\t\t}\n\n\t\tresult := &toolSearchResult{}\n\t\terr := json.Unmarshal([]byte(message.Content), result)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to unmarshal tool search tool result: %w\", err)\n\t\t}\n\t\tselectedTools = append(selectedTools, result.SelectedTools...)\n\t}\n\treturn selectedTools, nil\n}\n\nfunc invertSelect[T comparable](all []T, selected []T) map[T]struct{} {\n\tselectedSet := make(map[T]struct{}, len(selected))\n\tfor _, s := range selected {\n\t\tselectedSet[s] = struct{}{}\n\t}\n\n\tresult := make(map[T]struct{})\n\tfor _, item := range all {\n\t\tif _, ok := selectedSet[item]; !ok {\n\t\t\tresult[item] = struct{}{}\n\t\t}\n\t}\n\treturn result\n}\n\nfunc removeTools(ctx context.Context, all []*schema.ToolInfo, dynamicTools []tool.BaseTool, messages []*schema.Message) ([]*schema.ToolInfo, error) {\n\tselectedToolNames, err := extractSelectedTools(ctx, messages)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdynamicToolNames, err := getToolNames(ctx, dynamicTools)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tremoveMap := invertSelect(dynamicToolNames, selectedToolNames)\n\tret := make([]*schema.ToolInfo, 0, len(all)-len(dynamicTools))\n\tfor _, info := range all {\n\t\tif _, ok := removeMap[info.Name]; ok {\n\t\t\tcontinue\n\t\t}\n\t\tret = append(ret, info)\n\t}\n\treturn ret, nil\n}\n"
  },
  {
    "path": "adk/middlewares/dynamictool/toolsearch/toolsearch_test.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage toolsearch\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype mockTool struct {\n\tname string\n\tdesc string\n}\n\nfunc (m *mockTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: m.name,\n\t\tDesc: m.desc,\n\t}, nil\n}\n\nfunc newMockTool(name, desc string) *mockTool {\n\treturn &mockTool{name: name, desc: desc}\n}\n\nfunc TestNew(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"nil config returns error\", func(t *testing.T) {\n\t\tm, err := New(ctx, nil)\n\t\tassert.Nil(t, m)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"config is required\")\n\t})\n\n\tt.Run(\"empty tools returns error\", func(t *testing.T) {\n\t\tm, err := New(ctx, &Config{DynamicTools: []tool.BaseTool{}})\n\t\tassert.Nil(t, m)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"tools is required\")\n\t})\n\n\tt.Run(\"valid config returns middleware\", func(t *testing.T) {\n\t\ttools := []tool.BaseTool{\n\t\t\tnewMockTool(\"tool1\", \"desc1\"),\n\t\t\tnewMockTool(\"tool2\", \"desc2\"),\n\t\t}\n\t\tm, err := New(ctx, &Config{DynamicTools: tools})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, m)\n\t})\n}\n\nfunc TestMiddleware_BeforeAgent(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"nil runCtx returns nil\", func(t *testing.T) {\n\t\ttools := []tool.BaseTool{newMockTool(\"tool1\", \"desc1\")}\n\t\tm, err := New(ctx, &Config{DynamicTools: tools})\n\t\trequire.NoError(t, err)\n\n\t\tnewCtx, newRunCtx, err := m.BeforeAgent(ctx, nil)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, ctx, newCtx)\n\t\tassert.Nil(t, newRunCtx)\n\t})\n\n\tt.Run(\"adds tool_search and dynamic tools\", func(t *testing.T) {\n\t\ttools := []tool.BaseTool{\n\t\t\tnewMockTool(\"tool1\", \"desc1\"),\n\t\t\tnewMockTool(\"tool2\", \"desc2\"),\n\t\t}\n\t\tm, err := New(ctx, &Config{DynamicTools: tools})\n\t\trequire.NoError(t, err)\n\n\t\tmiddleware := m.(*middleware)\n\t\trunCtx := &adk.ChatModelAgentContext{\n\t\t\tTools: []tool.BaseTool{},\n\t\t}\n\n\t\t_, newRunCtx, err := middleware.BeforeAgent(ctx, runCtx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, newRunCtx)\n\t\tassert.Len(t, newRunCtx.Tools, 3)\n\t})\n}\n\nfunc TestToolSearchTool_Info(t *testing.T) {\n\tctx := context.Background()\n\ttoolNames := []string{\"tool1\", \"tool2\", \"tool3\"}\n\ttst := newToolSearchTool(toolNames)\n\n\tinfo, err := tst.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"tool_search\", info.Name)\n\tassert.Contains(t, info.Desc, \"regex pattern\")\n\tassert.NotNil(t, info.ParamsOneOf)\n}\n\nfunc TestToolSearchTool_InvokableRun(t *testing.T) {\n\tctx := context.Background()\n\ttoolNames := []string{\"get_weather\", \"get_time\", \"search_web\", \"calculate_sum\"}\n\ttst := newToolSearchTool(toolNames)\n\n\tt.Run(\"empty regex pattern returns error\", func(t *testing.T) {\n\t\targs := `{\"regex_pattern\": \"\"}`\n\t\tresult, err := tst.InvokableRun(ctx, args)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"regex_pattern is required\")\n\t\tassert.Empty(t, result)\n\t})\n\n\tt.Run(\"invalid json returns error\", func(t *testing.T) {\n\t\targs := `{invalid json}`\n\t\tresult, err := tst.InvokableRun(ctx, args)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"failed to unmarshal\")\n\t\tassert.Empty(t, result)\n\t})\n\n\tt.Run(\"invalid regex returns error\", func(t *testing.T) {\n\t\targs := `{\"regex_pattern\": \"[invalid\"}`\n\t\tresult, err := tst.InvokableRun(ctx, args)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"invalid regex pattern\")\n\t\tassert.Empty(t, result)\n\t})\n\n\tt.Run(\"matches tools with prefix pattern\", func(t *testing.T) {\n\t\targs := `{\"regex_pattern\": \"^get_\"}`\n\t\tresult, err := tst.InvokableRun(ctx, args)\n\t\tassert.NoError(t, err)\n\n\t\tvar res toolSearchResult\n\t\terr = json.Unmarshal([]byte(result), &res)\n\t\tassert.NoError(t, err)\n\t\tassert.ElementsMatch(t, []string{\"get_weather\", \"get_time\"}, res.SelectedTools)\n\t})\n\n\tt.Run(\"matches tools with suffix pattern\", func(t *testing.T) {\n\t\targs := `{\"regex_pattern\": \"_sum$\"}`\n\t\tresult, err := tst.InvokableRun(ctx, args)\n\t\tassert.NoError(t, err)\n\n\t\tvar res toolSearchResult\n\t\terr = json.Unmarshal([]byte(result), &res)\n\t\tassert.NoError(t, err)\n\t\tassert.ElementsMatch(t, []string{\"calculate_sum\"}, res.SelectedTools)\n\t})\n\n\tt.Run(\"matches all tools with wildcard\", func(t *testing.T) {\n\t\targs := `{\"regex_pattern\": \".*\"}`\n\t\tresult, err := tst.InvokableRun(ctx, args)\n\t\tassert.NoError(t, err)\n\n\t\tvar res toolSearchResult\n\t\terr = json.Unmarshal([]byte(result), &res)\n\t\tassert.NoError(t, err)\n\t\tassert.ElementsMatch(t, toolNames, res.SelectedTools)\n\t})\n\n\tt.Run(\"no matches returns empty list\", func(t *testing.T) {\n\t\targs := `{\"regex_pattern\": \"^nonexistent_\"}`\n\t\tresult, err := tst.InvokableRun(ctx, args)\n\t\tassert.NoError(t, err)\n\n\t\tvar res toolSearchResult\n\t\terr = json.Unmarshal([]byte(result), &res)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, res.SelectedTools)\n\t})\n}\n\nfunc TestGetToolNames(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"returns tool names\", func(t *testing.T) {\n\t\ttools := []tool.BaseTool{\n\t\t\tnewMockTool(\"tool1\", \"desc1\"),\n\t\t\tnewMockTool(\"tool2\", \"desc2\"),\n\t\t\tnewMockTool(\"tool3\", \"desc3\"),\n\t\t}\n\t\tnames, err := getToolNames(ctx, tools)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []string{\"tool1\", \"tool2\", \"tool3\"}, names)\n\t})\n\n\tt.Run(\"empty tools returns empty slice\", func(t *testing.T) {\n\t\tnames, err := getToolNames(ctx, []tool.BaseTool{})\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, names)\n\t})\n}\n\nfunc TestExtractSelectedTools(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"extracts selected tools from messages\", func(t *testing.T) {\n\t\tresult := toolSearchResult{SelectedTools: []string{\"tool1\", \"tool2\"}}\n\t\tresultJSON, _ := json.Marshal(result)\n\n\t\tmessages := []*schema.Message{\n\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t{Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultJSON)},\n\t\t}\n\n\t\tselected, err := extractSelectedTools(ctx, messages)\n\t\tassert.NoError(t, err)\n\t\tassert.ElementsMatch(t, []string{\"tool1\", \"tool2\"}, selected)\n\t})\n\n\tt.Run(\"handles multiple tool_search results\", func(t *testing.T) {\n\t\tresult1 := toolSearchResult{SelectedTools: []string{\"tool1\"}}\n\t\tresult1JSON, _ := json.Marshal(result1)\n\t\tresult2 := toolSearchResult{SelectedTools: []string{\"tool2\", \"tool3\"}}\n\t\tresult2JSON, _ := json.Marshal(result2)\n\n\t\tmessages := []*schema.Message{\n\t\t\t{Role: schema.Tool, ToolName: toolSearchToolName, Content: string(result1JSON)},\n\t\t\tschema.UserMessage(\"continue\"),\n\t\t\t{Role: schema.Tool, ToolName: toolSearchToolName, Content: string(result2JSON)},\n\t\t}\n\n\t\tselected, err := extractSelectedTools(ctx, messages)\n\t\tassert.NoError(t, err)\n\t\tassert.ElementsMatch(t, []string{\"tool1\", \"tool2\", \"tool3\"}, selected)\n\t})\n\n\tt.Run(\"ignores non-tool_search messages\", func(t *testing.T) {\n\t\tmessages := []*schema.Message{\n\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t{Role: schema.Tool, ToolName: \"other_tool\", Content: \"some content\"},\n\t\t\t{Role: schema.Assistant, Content: \"response\"},\n\t\t}\n\n\t\tselected, err := extractSelectedTools(ctx, messages)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, selected)\n\t})\n\n\tt.Run(\"returns error for invalid json\", func(t *testing.T) {\n\t\tmessages := []*schema.Message{\n\t\t\t{Role: schema.Tool, ToolName: toolSearchToolName, Content: \"invalid json\"},\n\t\t}\n\n\t\tselected, err := extractSelectedTools(ctx, messages)\n\t\tassert.Error(t, err)\n\t\tassert.Nil(t, selected)\n\t})\n}\n\nfunc TestInvertSelect(t *testing.T) {\n\tt.Run(\"returns items not in selected\", func(t *testing.T) {\n\t\tall := []string{\"a\", \"b\", \"c\", \"d\"}\n\t\tselected := []string{\"b\", \"d\"}\n\n\t\tresult := invertSelect(all, selected)\n\t\tassert.Len(t, result, 2)\n\t\t_, hasA := result[\"a\"]\n\t\t_, hasC := result[\"c\"]\n\t\tassert.True(t, hasA)\n\t\tassert.True(t, hasC)\n\t})\n\n\tt.Run(\"empty selected returns all\", func(t *testing.T) {\n\t\tall := []string{\"a\", \"b\", \"c\"}\n\t\tselected := []string{}\n\n\t\tresult := invertSelect(all, selected)\n\t\tassert.Len(t, result, 3)\n\t})\n\n\tt.Run(\"all selected returns empty\", func(t *testing.T) {\n\t\tall := []string{\"a\", \"b\"}\n\t\tselected := []string{\"a\", \"b\"}\n\n\t\tresult := invertSelect(all, selected)\n\t\tassert.Empty(t, result)\n\t})\n\n\tt.Run(\"works with integers\", func(t *testing.T) {\n\t\tall := []int{1, 2, 3, 4, 5}\n\t\tselected := []int{2, 4}\n\n\t\tresult := invertSelect(all, selected)\n\t\tassert.Len(t, result, 3)\n\t\t_, has1 := result[1]\n\t\t_, has3 := result[3]\n\t\t_, has5 := result[5]\n\t\tassert.True(t, has1)\n\t\tassert.True(t, has3)\n\t\tassert.True(t, has5)\n\t})\n}\n\nfunc TestRemoveTools(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"removes unselected dynamic tools\", func(t *testing.T) {\n\t\tallTools := []*schema.ToolInfo{\n\t\t\t{Name: \"static_tool\"},\n\t\t\t{Name: \"dynamic_tool1\"},\n\t\t\t{Name: \"dynamic_tool2\"},\n\t\t\t{Name: \"dynamic_tool3\"},\n\t\t}\n\n\t\tdynamicTools := []tool.BaseTool{\n\t\t\tnewMockTool(\"dynamic_tool1\", \"\"),\n\t\t\tnewMockTool(\"dynamic_tool2\", \"\"),\n\t\t\tnewMockTool(\"dynamic_tool3\", \"\"),\n\t\t}\n\n\t\tresult := toolSearchResult{SelectedTools: []string{\"dynamic_tool1\"}}\n\t\tresultJSON, _ := json.Marshal(result)\n\t\tmessages := []*schema.Message{\n\t\t\t{Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultJSON)},\n\t\t}\n\n\t\ttools, err := removeTools(ctx, allTools, dynamicTools, messages)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 2)\n\n\t\ttoolNames := make([]string, len(tools))\n\t\tfor i, t := range tools {\n\t\t\ttoolNames[i] = t.Name\n\t\t}\n\t\tassert.ElementsMatch(t, []string{\"static_tool\", \"dynamic_tool1\"}, toolNames)\n\t})\n\n\tt.Run(\"remove all dynamic tools when no tool_search result\", func(t *testing.T) {\n\t\tallTools := []*schema.ToolInfo{\n\t\t\t{Name: \"static_tool\"},\n\t\t\t{Name: \"dynamic_tool1\"},\n\t\t}\n\n\t\tdynamicTools := []tool.BaseTool{\n\t\t\tnewMockTool(\"dynamic_tool1\", \"\"),\n\t\t}\n\n\t\tmessages := []*schema.Message{\n\t\t\tschema.UserMessage(\"hello\"),\n\t\t}\n\n\t\ttools, err := removeTools(ctx, allTools, dynamicTools, messages)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 1)\n\t\tassert.Equal(t, \"static_tool\", tools[0].Name)\n\t})\n\n\tt.Run(\"handles empty dynamic tools\", func(t *testing.T) {\n\t\tallTools := []*schema.ToolInfo{\n\t\t\t{Name: \"static_tool1\"},\n\t\t\t{Name: \"static_tool2\"},\n\t\t}\n\n\t\tdynamicTools := []tool.BaseTool{}\n\t\tmessages := []*schema.Message{}\n\n\t\ttools, err := removeTools(ctx, allTools, dynamicTools, messages)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 2)\n\t})\n}\n\ntype mockChatModel struct {\n\tgenerateFunc func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error)\n\tstreamFunc   func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error)\n}\n\nfunc (m *mockChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tif m.generateFunc != nil {\n\t\treturn m.generateFunc(ctx, input, opts...)\n\t}\n\treturn &schema.Message{Role: schema.Assistant, Content: \"response\"}, nil\n}\n\nfunc (m *mockChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tif m.streamFunc != nil {\n\t\treturn m.streamFunc(ctx, input, opts...)\n\t}\n\treturn nil, nil\n}\n\nfunc TestWrapper_Generate(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"filters tools based on tool_search result\", func(t *testing.T) {\n\t\tallTools := []*schema.ToolInfo{\n\t\t\t{Name: \"static_tool\"},\n\t\t\t{Name: \"dynamic_tool1\"},\n\t\t\t{Name: \"dynamic_tool2\"},\n\t\t}\n\n\t\tdynamicTools := []tool.BaseTool{\n\t\t\tnewMockTool(\"dynamic_tool1\", \"\"),\n\t\t\tnewMockTool(\"dynamic_tool2\", \"\"),\n\t\t}\n\n\t\tresult := toolSearchResult{SelectedTools: []string{\"dynamic_tool1\"}}\n\t\tresultJSON, _ := json.Marshal(result)\n\n\t\tmessages := []*schema.Message{\n\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t{Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultJSON)},\n\t\t}\n\n\t\tw := &wrapper{\n\t\t\tallTools:     allTools,\n\t\t\tdynamicTools: dynamicTools,\n\t\t\tcm: &mockChatModel{\n\t\t\t\tgenerateFunc: func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t\t\toptions := model.GetCommonOptions(nil, opts...)\n\t\t\t\t\tassert.Len(t, options.Tools, 2)\n\t\t\t\t\tassert.Equal(t, \"static_tool\", options.Tools[0].Name)\n\t\t\t\t\tassert.Equal(t, \"dynamic_tool1\", options.Tools[1].Name)\n\t\t\t\t\treturn nil, nil\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := w.Generate(ctx, messages)\n\t\tassert.NoError(t, err)\n\t})\n}\n\nfunc TestWrapper_Stream(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"filters tools based on tool_search result\", func(t *testing.T) {\n\t\tallTools := []*schema.ToolInfo{\n\t\t\t{Name: \"static_tool\"},\n\t\t\t{Name: \"dynamic_tool1\"},\n\t\t\t{Name: \"dynamic_tool2\"},\n\t\t}\n\n\t\tdynamicTools := []tool.BaseTool{\n\t\t\tnewMockTool(\"dynamic_tool1\", \"\"),\n\t\t\tnewMockTool(\"dynamic_tool2\", \"\"),\n\t\t}\n\n\t\tresult := toolSearchResult{SelectedTools: []string{\"dynamic_tool1\"}}\n\t\tresultJSON, _ := json.Marshal(result)\n\n\t\tmessages := []*schema.Message{\n\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t{Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultJSON)},\n\t\t}\n\n\t\tw := &wrapper{\n\t\t\tallTools:     allTools,\n\t\t\tdynamicTools: dynamicTools,\n\t\t\tcm: &mockChatModel{\n\t\t\t\tstreamFunc: func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\t\t\t\t\toptions := model.GetCommonOptions(nil, opts...)\n\t\t\t\t\tassert.Len(t, options.Tools, 2)\n\t\t\t\t\tassert.Equal(t, \"static_tool\", options.Tools[0].Name)\n\t\t\t\t\tassert.Equal(t, \"dynamic_tool1\", options.Tools[1].Name)\n\t\t\t\t\treturn nil, nil\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tstream, err := w.Stream(ctx, messages)\n\t\tassert.NoError(t, err)\n\t\tassert.Nil(t, stream)\n\t})\n}\n"
  },
  {
    "path": "adk/middlewares/filesystem/backend.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package filesystem provides middlewares.\npackage filesystem\n\nimport (\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n)\n\ntype FileInfo = filesystem.FileInfo\ntype GrepMatch = filesystem.GrepMatch\ntype LsInfoRequest = filesystem.LsInfoRequest\ntype ReadRequest = filesystem.ReadRequest\ntype GrepRequest = filesystem.GrepRequest\ntype GlobInfoRequest = filesystem.GlobInfoRequest\ntype WriteRequest = filesystem.WriteRequest\ntype EditRequest = filesystem.EditRequest\ntype FileContent = filesystem.FileContent\n"
  },
  {
    "path": "adk/middlewares/filesystem/filesystem.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage filesystem\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"path/filepath\"\n\t\"runtime/debug\"\n\t\"sort\"\n\t\"strconv\"\n\t\"strings\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/components/tool/utils\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nconst (\n\tToolNameLs        = \"ls\"\n\tToolNameReadFile  = \"read_file\"\n\tToolNameWriteFile = \"write_file\"\n\tToolNameEditFile  = \"edit_file\"\n\tToolNameGlob      = \"glob\"\n\tToolNameGrep      = \"grep\"\n\tToolNameExecute   = \"execute\"\n\n\tnoFilesFound   = \"No files found\"\n\tnoMatchesFound = \"No matches found\"\n)\n\n// ToolConfig configures a filesystem tool\ntype ToolConfig struct {\n\t// Name overrides the tool name used in tool registration\n\t// optional, default tool name will be used if not set (empty string)\n\tName string\n\n\t// Desc overrides the tool description used in tool registration\n\t// optional, default tool description will be used if not set (nil pointer)\n\tDesc *string\n\n\t// CustomTool provides a custom implementation for this tool.\n\t// If set, this custom tool will be used instead of the default implementation associated with Backend.\n\t// If not set, the default tool implementation associated with Backend will be created automatically.\n\t// optional\n\tCustomTool tool.BaseTool\n\n\t// Disable disables this tool\n\t// If true, the tool will not be registered\n\t// optional, false by default\n\tDisable bool\n}\n\n// Config is the configuration for the filesystem middleware\ntype Config struct {\n\t// Backend provides filesystem operations used by tools and offloading.\n\t// If set, filesystem tools (read_file, write_file, edit_file, glob, grep) will be registered.\n\t// At least one of Backend, Shell, or StreamingShell must be set.\n\tBackend filesystem.Backend\n\n\t// Shell provides shell command execution capability.\n\t// If set, an execute tool will be registered to support shell command execution.\n\t// At least one of Backend, Shell, or StreamingShell must be set.\n\t// Mutually exclusive with StreamingShell.\n\tShell filesystem.Shell\n\t// StreamingShell provides streaming shell command execution capability.\n\t// If set, a streaming execute tool will be registered to support streaming shell command execution.\n\t// At least one of Backend, Shell, or StreamingShell must be set.\n\t// Mutually exclusive with Shell.\n\tStreamingShell filesystem.StreamingShell\n\n\t// LsToolConfig configures the ls tool\n\t// optional\n\tLsToolConfig *ToolConfig\n\t// ReadFileToolConfig configures the read_file tool\n\t// optional\n\tReadFileToolConfig *ToolConfig\n\t// WriteFileToolConfig configures the write_file tool\n\t// optional\n\tWriteFileToolConfig *ToolConfig\n\t// EditFileToolConfig configures the edit_file tool\n\t// optional\n\tEditFileToolConfig *ToolConfig\n\t// GlobToolConfig configures the glob tool\n\t// optional\n\tGlobToolConfig *ToolConfig\n\t// GrepToolConfig configures the grep tool\n\t// optional\n\tGrepToolConfig *ToolConfig\n\n\t// WithoutLargeToolResultOffloading disables automatic offloading of large tool result to Backend\n\t// optional, false(enabled) by default\n\tWithoutLargeToolResultOffloading bool\n\t// LargeToolResultOffloadingTokenLimit sets the token threshold to trigger offloading\n\t// optional, 20000 by default\n\tLargeToolResultOffloadingTokenLimit int\n\t// LargeToolResultOffloadingPathGen generates the write path for offloaded results based on context and ToolInput\n\t// optional, \"/large_tool_result/{ToolCallID}\" by default\n\tLargeToolResultOffloadingPathGen func(ctx context.Context, input *compose.ToolInput) (string, error)\n\n\t// CustomSystemPrompt overrides the default ToolsSystemPrompt appended to agent instruction\n\t// optional, ToolsSystemPrompt by default\n\tCustomSystemPrompt *string\n\n\t// CustomLsToolDesc overrides the ls tool description used in tool registration\n\t// optional, ListFilesToolDesc by default\n\t// Deprecated: Use LsToolConfig.Desc instead\n\tCustomLsToolDesc *string\n\t// CustomReadFileToolDesc overrides the read_file tool description\n\t// optional, ReadFileToolDesc by default\n\t// Deprecated: Use ReadFileToolConfig.Desc instead\n\tCustomReadFileToolDesc *string\n\t// CustomGrepToolDesc overrides the grep tool description\n\t// optional, GrepToolDesc by default\n\t// Deprecated: Use GrepToolConfig.Desc instead\n\tCustomGrepToolDesc *string\n\t// CustomGlobToolDesc overrides the glob tool description\n\t// optional, GlobToolDesc by default\n\t// Deprecated: Use GlobToolConfig.Desc instead\n\tCustomGlobToolDesc *string\n\t// CustomWriteFileToolDesc overrides the write_file tool description\n\t// optional, WriteFileToolDesc by default\n\t// Deprecated: Use WriteFileToolConfig.Desc instead\n\tCustomWriteFileToolDesc *string\n\t// CustomEditToolDesc overrides the edit_file tool description\n\t// optional, EditFileToolDesc by default\n\t// Deprecated: Use EditFileToolConfig.Desc instead\n\tCustomEditToolDesc *string\n}\n\nfunc (c *Config) Validate() error {\n\tif c == nil {\n\t\treturn errors.New(\"config should not be nil\")\n\t}\n\tif c.Backend == nil {\n\t\treturn errors.New(\"backend should not be nil\")\n\t}\n\tif c.StreamingShell != nil && c.Shell != nil {\n\t\treturn errors.New(\"shell and streaming shell should not be both set\")\n\t}\n\treturn nil\n}\n\n// NewMiddleware constructs and returns the filesystem middleware.\n//\n// Deprecated: Use New instead. New returns\n// a ChatModelAgentMiddleware which provides better context propagation through wrapper methods\n// and is the recommended approach for new code. See ChatModelAgentMiddleware documentation\n// for details on the benefits over AgentMiddleware.\nfunc NewMiddleware(ctx context.Context, config *Config) (adk.AgentMiddleware, error) {\n\terr := config.Validate()\n\tif err != nil {\n\t\treturn adk.AgentMiddleware{}, err\n\t}\n\tts, err := getFilesystemTools(ctx, &MiddlewareConfig{\n\t\tBackend:                 config.Backend,\n\t\tShell:                   config.Shell,\n\t\tStreamingShell:          config.StreamingShell,\n\t\tLsToolConfig:            config.LsToolConfig,\n\t\tReadFileToolConfig:      config.ReadFileToolConfig,\n\t\tWriteFileToolConfig:     config.WriteFileToolConfig,\n\t\tEditFileToolConfig:      config.EditFileToolConfig,\n\t\tGlobToolConfig:          config.GlobToolConfig,\n\t\tGrepToolConfig:          config.GrepToolConfig,\n\t\tCustomSystemPrompt:      config.CustomSystemPrompt,\n\t\tCustomLsToolDesc:        config.CustomLsToolDesc,\n\t\tCustomReadFileToolDesc:  config.CustomReadFileToolDesc,\n\t\tCustomGrepToolDesc:      config.CustomGrepToolDesc,\n\t\tCustomGlobToolDesc:      config.CustomGlobToolDesc,\n\t\tCustomWriteFileToolDesc: config.CustomWriteFileToolDesc,\n\t\tCustomEditToolDesc:      config.CustomEditToolDesc,\n\t})\n\tif err != nil {\n\t\treturn adk.AgentMiddleware{}, err\n\t}\n\n\tvar systemPrompt string\n\tif config.CustomSystemPrompt != nil {\n\t\tsystemPrompt = *config.CustomSystemPrompt\n\t}\n\n\tm := adk.AgentMiddleware{\n\t\tAdditionalInstruction: systemPrompt,\n\t\tAdditionalTools:       ts,\n\t}\n\n\tif !config.WithoutLargeToolResultOffloading {\n\t\tm.WrapToolCall = newToolResultOffloading(ctx, &toolResultOffloadingConfig{\n\t\t\tBackend:       config.Backend,\n\t\t\tTokenLimit:    config.LargeToolResultOffloadingTokenLimit,\n\t\t\tPathGenerator: config.LargeToolResultOffloadingPathGen,\n\t\t})\n\t}\n\n\treturn m, nil\n}\n\n// MiddlewareConfig is the configuration for the filesystem middleware\ntype MiddlewareConfig struct {\n\t// Backend provides filesystem operations used by tools and offloading.\n\t// required\n\tBackend filesystem.Backend\n\n\t// Shell provides shell command execution capability.\n\t// If set, an execute tool will be registered to support shell command execution.\n\t// optional, mutually exclusive with StreamingShell\n\tShell filesystem.Shell\n\t// StreamingShell provides streaming shell command execution capability.\n\t// If set, a streaming execute tool will be registered for real-time output.\n\t// optional, mutually exclusive with Shell\n\tStreamingShell filesystem.StreamingShell\n\n\t// LsToolConfig configures the ls tool\n\t// optional\n\tLsToolConfig *ToolConfig\n\t// ReadFileToolConfig configures the read_file tool\n\t// optional\n\tReadFileToolConfig *ToolConfig\n\t// WriteFileToolConfig configures the write_file tool\n\t// optional\n\tWriteFileToolConfig *ToolConfig\n\t// EditFileToolConfig configures the edit_file tool\n\t// optional\n\tEditFileToolConfig *ToolConfig\n\t// GlobToolConfig configures the glob tool\n\t// optional\n\tGlobToolConfig *ToolConfig\n\t// GrepToolConfig configures the grep tool\n\t// optional\n\tGrepToolConfig *ToolConfig\n\n\t// CustomSystemPrompt overrides the default ToolsSystemPrompt appended to agent instruction\n\t// optional, ToolsSystemPrompt by default\n\tCustomSystemPrompt *string\n\n\t// CustomLsToolDesc overrides the ls tool description used in tool registration\n\t// optional, ListFilesToolDesc by default\n\t// Deprecated: Use LsToolConfig.Desc instead\n\tCustomLsToolDesc *string\n\t// CustomReadFileToolDesc overrides the read_file tool description\n\t// optional, ReadFileToolDesc by default\n\t// Deprecated: Use ReadFileToolConfig.Desc instead\n\tCustomReadFileToolDesc *string\n\t// CustomGrepToolDesc overrides the grep tool description\n\t// optional, GrepToolDesc by default\n\t// Deprecated: Use GrepToolConfig.Desc instead\n\tCustomGrepToolDesc *string\n\t// CustomGlobToolDesc overrides the glob tool description\n\t// optional, GlobToolDesc by default\n\t// Deprecated: Use GlobToolConfig.Desc instead\n\tCustomGlobToolDesc *string\n\t// CustomWriteFileToolDesc overrides the write_file tool description\n\t// optional, WriteFileToolDesc by default\n\t// Deprecated: Use WriteFileToolConfig.Desc instead\n\tCustomWriteFileToolDesc *string\n\t// CustomEditToolDesc overrides the edit_file tool description\n\t// optional, EditFileToolDesc by default\n\t// Deprecated: Use EditFileToolConfig.Desc instead\n\tCustomEditToolDesc *string\n}\n\nfunc (c *MiddlewareConfig) Validate() error {\n\tif c == nil {\n\t\treturn errors.New(\"config should not be nil\")\n\t}\n\tif c.Backend == nil {\n\t\treturn errors.New(\"backend should not be nil\")\n\t}\n\tif c.StreamingShell != nil && c.Shell != nil {\n\t\treturn errors.New(\"shell and streaming shell should not be both set\")\n\t}\n\treturn nil\n}\n\n// mergeToolConfigWithDesc merges ToolConfig with legacy Desc field\n// Priority: ToolConfig.Desc > legacy Desc\n// Returns an empty ToolConfig if both are nil (to allow backend default implementation)\nfunc (c *MiddlewareConfig) mergeToolConfigWithDesc(\n\ttoolConfig *ToolConfig,\n\tlegacyDesc *string,\n) *ToolConfig {\n\tif toolConfig == nil && legacyDesc == nil {\n\t\treturn &ToolConfig{}\n\t}\n\n\tif toolConfig == nil {\n\t\treturn &ToolConfig{\n\t\t\tDesc: legacyDesc,\n\t\t}\n\t}\n\n\tif toolConfig.Desc == nil && legacyDesc != nil {\n\t\tmerged := *toolConfig\n\t\tmerged.Desc = legacyDesc\n\t\treturn &merged\n\t}\n\n\treturn toolConfig\n}\n\n// New constructs and returns the filesystem middleware as a ChatModelAgentMiddleware.\n//\n// This is the recommended constructor for new code. It returns a ChatModelAgentMiddleware which provides:\n//   - Better context propagation through WrapInvokableToolCall and WrapStreamableToolCall methods\n//   - BeforeAgent hook for modifying agent instruction and tools at runtime\n//   - More flexible extension points compared to the struct-based AgentMiddleware\n//\n// The middleware provides filesystem tools (ls, read_file, write_file, edit_file, glob, grep)\n// and optionally an execute tool if the Backend implements ShellBackend or StreamingShellBackend.\n//\n// Example usage:\n//\n//\tmiddleware, err := filesystem.New(ctx, &filesystem.Config{\n//\t    Backend: myBackend,\n//\t})\n//\tagent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n//\t    // ...\n//\t    Handlers: []adk.ChatModelAgentMiddleware{middleware},\n//\t})\nfunc New(ctx context.Context, config *MiddlewareConfig) (adk.ChatModelAgentMiddleware, error) {\n\terr := config.Validate()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tts, err := getFilesystemTools(ctx, config)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tvar systemPrompt string\n\tif config.CustomSystemPrompt != nil {\n\t\tsystemPrompt = *config.CustomSystemPrompt\n\t}\n\n\tm := &filesystemMiddleware{\n\t\tadditionalInstruction: systemPrompt,\n\t\tadditionalTools:       ts,\n\t}\n\n\treturn m, nil\n}\n\ntype filesystemMiddleware struct {\n\tadk.BaseChatModelAgentMiddleware\n\tadditionalInstruction string\n\tadditionalTools       []tool.BaseTool\n}\n\nfunc (m *filesystemMiddleware) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) {\n\tif runCtx == nil {\n\t\treturn ctx, runCtx, nil\n\t}\n\n\tnRunCtx := *runCtx\n\tif m.additionalInstruction != \"\" {\n\t\tnRunCtx.Instruction = nRunCtx.Instruction + \"\\n\" + m.additionalInstruction\n\t}\n\tnRunCtx.Tools = append(nRunCtx.Tools, m.additionalTools...)\n\treturn ctx, &nRunCtx, nil\n}\n\n// toolSpec defines a specification for creating a filesystem tool.\n// It unifies the tool creation process by encapsulating the tool configuration,\n// legacy descriptor, and the creation function.\ntype toolSpec struct {\n\tconfig     *ToolConfig\n\tlegacyDesc *string\n\tcreateFunc func(name, desc string) (tool.BaseTool, error)\n}\n\nfunc getFilesystemTools(_ context.Context, middlewareConfig *MiddlewareConfig) ([]tool.BaseTool, error) {\n\tvar tools []tool.BaseTool\n\n\ttoolSpecs := []toolSpec{\n\t\t{\n\t\t\tconfig:     middlewareConfig.LsToolConfig,\n\t\t\tlegacyDesc: middlewareConfig.CustomLsToolDesc,\n\t\t\tcreateFunc: func(name, desc string) (tool.BaseTool, error) {\n\t\t\t\tif middlewareConfig.Backend != nil {\n\t\t\t\t\treturn newLsTool(middlewareConfig.Backend, name, desc)\n\t\t\t\t}\n\t\t\t\treturn nil, nil\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tconfig:     middlewareConfig.ReadFileToolConfig,\n\t\t\tlegacyDesc: middlewareConfig.CustomReadFileToolDesc,\n\t\t\tcreateFunc: func(name, desc string) (tool.BaseTool, error) {\n\t\t\t\tif middlewareConfig.Backend != nil {\n\t\t\t\t\treturn newReadFileTool(middlewareConfig.Backend, name, desc)\n\t\t\t\t}\n\t\t\t\treturn nil, nil\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tconfig:     middlewareConfig.WriteFileToolConfig,\n\t\t\tlegacyDesc: middlewareConfig.CustomWriteFileToolDesc,\n\t\t\tcreateFunc: func(name, desc string) (tool.BaseTool, error) {\n\t\t\t\tif middlewareConfig.Backend != nil {\n\t\t\t\t\treturn newWriteFileTool(middlewareConfig.Backend, name, desc)\n\t\t\t\t}\n\t\t\t\treturn nil, nil\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tconfig:     middlewareConfig.EditFileToolConfig,\n\t\t\tlegacyDesc: middlewareConfig.CustomEditToolDesc,\n\t\t\tcreateFunc: func(name, desc string) (tool.BaseTool, error) {\n\t\t\t\tif middlewareConfig.Backend != nil {\n\t\t\t\t\treturn newEditFileTool(middlewareConfig.Backend, name, desc)\n\t\t\t\t}\n\t\t\t\treturn nil, nil\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tconfig:     middlewareConfig.GlobToolConfig,\n\t\t\tlegacyDesc: middlewareConfig.CustomGlobToolDesc,\n\t\t\tcreateFunc: func(name, desc string) (tool.BaseTool, error) {\n\t\t\t\tif middlewareConfig.Backend != nil {\n\t\t\t\t\treturn newGlobTool(middlewareConfig.Backend, name, desc)\n\t\t\t\t}\n\t\t\t\treturn nil, nil\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tconfig:     middlewareConfig.GrepToolConfig,\n\t\t\tlegacyDesc: middlewareConfig.CustomGrepToolDesc,\n\t\t\tcreateFunc: func(name, desc string) (tool.BaseTool, error) {\n\t\t\t\tif middlewareConfig.Backend != nil {\n\t\t\t\t\treturn newGrepTool(middlewareConfig.Backend, name, desc)\n\t\t\t\t}\n\t\t\t\treturn nil, nil\n\t\t\t},\n\t\t},\n\t}\n\n\tfor _, spec := range toolSpecs {\n\t\tt, err := createToolFromSpec(middlewareConfig, spec)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tif t != nil {\n\t\t\ttools = append(tools, t)\n\t\t}\n\t}\n\n\t// Create execute tool if Shell or StreamingShell is available\n\tif middlewareConfig.StreamingShell != nil {\n\t\texecuteDesc, err := selectToolDesc(\"\", ExecuteToolDesc, ExecuteToolDescChinese)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\texecuteTool, err := newStreamingExecuteTool(middlewareConfig.StreamingShell, ToolNameExecute, executeDesc)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\ttools = append(tools, executeTool)\n\t} else if middlewareConfig.Shell != nil {\n\t\texecuteDesc, err := selectToolDesc(\"\", ExecuteToolDesc, ExecuteToolDescChinese)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\texecuteTool, err := newExecuteTool(middlewareConfig.Shell, ToolNameExecute, executeDesc)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\ttools = append(tools, executeTool)\n\t}\n\n\treturn tools, nil\n}\n\n// createToolFromSpec creates a tool instance based on the provided toolSpec.\n// It handles configuration merging (ToolConfig + legacy Desc), checks if the tool\n// is disabled, and prioritizes CustomTool over the default implementation.\nfunc createToolFromSpec(middlewareConfig *MiddlewareConfig, spec toolSpec) (tool.BaseTool, error) {\n\tmergedConfig := middlewareConfig.mergeToolConfigWithDesc(spec.config, spec.legacyDesc)\n\n\tif mergedConfig.Disable {\n\t\treturn nil, nil\n\t}\n\n\treturn getOrCreateTool(mergedConfig.CustomTool, func() (tool.BaseTool, error) {\n\t\tdesc := \"\"\n\t\tif mergedConfig.Desc != nil {\n\t\t\tdesc = *mergedConfig.Desc\n\t\t}\n\t\treturn spec.createFunc(mergedConfig.Name, desc)\n\t})\n}\n\nfunc getOrCreateTool(customTool tool.BaseTool, createFunc func() (tool.BaseTool, error)) (tool.BaseTool, error) {\n\tif customTool != nil {\n\t\treturn customTool, nil\n\t}\n\treturn createFunc()\n}\n\ntype lsArgs struct {\n\tPath string `json:\"path\"`\n}\n\nfunc newLsTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) {\n\ttoolName := selectToolName(name, ToolNameLs)\n\td, err := selectToolDesc(desc, ListFilesToolDesc, ListFilesToolDescChinese)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn utils.InferTool(toolName, d, func(ctx context.Context, input lsArgs) (string, error) {\n\t\tinfos, err := fs.LsInfo(ctx, &filesystem.LsInfoRequest{Path: input.Path})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tif len(infos) == 0 {\n\t\t\treturn noFilesFound, nil\n\t\t}\n\t\tpaths := make([]string, 0, len(infos))\n\t\tfor _, fi := range infos {\n\t\t\tpaths = append(paths, fi.Path)\n\t\t}\n\t\treturn strings.Join(paths, \"\\n\"), nil\n\t})\n}\n\ntype readFileArgs struct {\n\t// FilePath is the path to the file to read.\n\tFilePath string `json:\"file_path\" jsonschema:\"description=The path to the file to read\"`\n\n\t// Offset is the line number to start reading from.\n\tOffset int `json:\"offset\" jsonschema:\"description=The line number to start reading from. Only provide if the file is too large to read at once\"`\n\n\t// Limit is the number of lines to read.\n\tLimit int `json:\"limit\" jsonschema:\"description=The number of lines to read. Only provide if the file is too large to read at once.\"`\n}\n\nfunc newReadFileTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) {\n\ttoolName := selectToolName(name, ToolNameReadFile)\n\td, err := selectToolDesc(desc, ReadFileToolDesc, ReadFileToolDescChinese)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn utils.InferTool(toolName, d, func(ctx context.Context, input readFileArgs) (string, error) {\n\t\tif input.Offset <= 0 {\n\t\t\tinput.Offset = 1\n\t\t}\n\n\t\tfileCt, err := fs.Read(ctx, &filesystem.ReadRequest{\n\t\t\tFilePath: input.FilePath,\n\t\t\tOffset:   input.Offset,\n\t\t\tLimit:    input.Limit,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\tstartLine := input.Offset\n\t\tlines := strings.Split(fileCt.Content, \"\\n\")\n\t\tvar b strings.Builder\n\t\tfor i, line := range lines {\n\t\t\tif i < len(lines)-1 {\n\t\t\t\tfmt.Fprintf(&b, \"%6d\\t%s\\n\", startLine+i, line)\n\t\t\t} else {\n\t\t\t\tfmt.Fprintf(&b, \"%6d\\t%s\", startLine+i, line)\n\t\t\t}\n\n\t\t}\n\t\treturn b.String(), nil\n\t})\n}\n\ntype writeFileArgs struct {\n\t// FilePath is the path to the file to write.\n\tFilePath string `json:\"file_path\" jsonschema:\"description=The path to the file to write\"`\n\n\t// Content is the content to write to the file.\n\tContent string `json:\"content\" jsonschema:\"description=The content to write to the file\"`\n}\n\nfunc newWriteFileTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) {\n\ttoolName := selectToolName(name, ToolNameWriteFile)\n\td, err := selectToolDesc(desc, WriteFileToolDesc, WriteFileToolDescChinese)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn utils.InferTool(toolName, d, func(ctx context.Context, input writeFileArgs) (string, error) {\n\t\terr := fs.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: input.FilePath,\n\t\t\tContent:  input.Content,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn fmt.Sprintf(\"Updated file %s\", input.FilePath), nil\n\t})\n}\n\ntype editFileArgs struct {\n\t// FilePath is the path to the file to modify.\n\tFilePath string `json:\"file_path\" jsonschema:\"description=The path to the file to modify\"`\n\n\t// OldString is the text to replace.\n\tOldString string `json:\"old_string\" jsonschema:\"description=The text to replace\"`\n\n\t// NewString is the text to replace it with.\n\tNewString string `json:\"new_string\" jsonschema:\"description=The text to replace it with (must be different from old_string)\"`\n\n\t// ReplaceAll indicates whether to replace all occurrences of old_string.\n\tReplaceAll bool `json:\"replace_all\" jsonschema:\"description=Replace all occurrences of old_string (default false),default=false\"`\n}\n\nfunc newEditFileTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) {\n\ttoolName := selectToolName(name, ToolNameEditFile)\n\td, err := selectToolDesc(desc, EditFileToolDesc, EditFileToolDescChinese)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn utils.InferTool(toolName, d, func(ctx context.Context, input editFileArgs) (string, error) {\n\t\terr := fs.Edit(ctx, &filesystem.EditRequest{\n\t\t\tFilePath:   input.FilePath,\n\t\t\tOldString:  input.OldString,\n\t\t\tNewString:  input.NewString,\n\t\t\tReplaceAll: input.ReplaceAll,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn fmt.Sprintf(\"Successfully replaced the string in '%s'\", input.FilePath), nil\n\t})\n}\n\ntype globArgs struct {\n\t// Pattern is the glob pattern to match files against.\n\tPattern string `json:\"pattern\" jsonschema:\"description=The glob pattern to match files against\"`\n\n\t// Path is the directory to search in.\n\tPath string `json:\"path\" jsonschema:\"description=The directory to search in. If not specified\\\\, the current working directory will be used. IMPORTANT: Omit this field to use the default directory. DO NOT enter 'undefined' or 'null' - simply omit it for the default behavior. Must be a valid directory path if provided.\"`\n}\n\nfunc newGlobTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) {\n\ttoolName := selectToolName(name, ToolNameGlob)\n\td, err := selectToolDesc(desc, GlobToolDesc, GlobToolDescChinese)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn utils.InferTool(toolName, d, func(ctx context.Context, input globArgs) (string, error) {\n\t\tinfos, err := fs.GlobInfo(ctx, &filesystem.GlobInfoRequest{\n\t\t\tPattern: input.Pattern,\n\t\t\tPath:    input.Path,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tif len(infos) == 0 {\n\t\t\treturn noFilesFound, nil\n\t\t}\n\t\tpaths := make([]string, 0, len(infos))\n\t\tfor _, fi := range infos {\n\t\t\tpaths = append(paths, fi.Path)\n\t\t}\n\t\treturn strings.Join(paths, \"\\n\"), nil\n\t})\n}\n\ntype grepArgs struct {\n\t// Pattern is the regular expression pattern to search for in file contents.\n\tPattern string `json:\"pattern\" jsonschema:\"description=The regular expression pattern to search for in file contents\"`\n\n\t// Path is the file or directory to search in. Defaults to current working directory.\n\tPath *string `json:\"path,omitempty\" jsonschema:\"description=File or directory to search in (rg PATH). Defaults to current working directory.\"`\n\n\t// Glob is the glob pattern to filter files (e.g. \"*.js\", \"*.{ts,tsx}\").\n\tGlob *string `json:\"glob,omitempty\" jsonschema:\"description=Glob pattern to filter files (e.g. '*.js'\\\\, '*.{ts\\\\,tsx}') - maps to rg --glob\"`\n\n\t// OutputMode specifies the output format.\n\t// \"content\" shows matching lines (supports context, line numbers, head_limit).\n\t// \"files_with_matches\" shows file paths (supports head_limit).\n\t// \"count\" shows match counts (supports head_limit).\n\t// Defaults to \"files_with_matches\".\n\tOutputMode string `json:\"output_mode,omitempty\" jsonschema:\"description=Output mode: 'content' shows matching lines (supports -A/-B/-C context\\\\, -n line numbers\\\\, head_limit)\\\\, 'files_with_matches' shows file paths (supports head_limit)\\\\, 'count' shows match counts (supports head_limit). Defaults to 'files_with_matches'.,enum=content,enum=files_with_matches,enum=count\"`\n\n\t// Context is the number of lines to show before and after each match.\n\t// Only applicable when output_mode is \"content\".\n\tContext *int `json:\"-C,omitempty\" jsonschema:\"description=Number of lines to show before and after each match (rg -C). Requires output_mode: 'content'\\\\, ignored otherwise.\"`\n\n\t// BeforeLines is the number of lines to show before each match.\n\t// Only applicable when output_mode is \"content\".\n\tBeforeLines *int `json:\"-B,omitempty\" jsonschema:\"description=Number of lines to show before each match (rg -B). Requires output_mode: 'content'\\\\, ignored otherwise.\"`\n\n\t// AfterLines is the number of lines to show after each match.\n\t// Only applicable when output_mode is \"content\".\n\tAfterLines *int `json:\"-A,omitempty\" jsonschema:\"description=Number of lines to show after each match (rg -A). Requires output_mode: 'content'\\\\, ignored otherwise.\"`\n\n\t// ShowLineNumbers enables showing line numbers in output.\n\t// Only applicable when output_mode is \"content\". Defaults to true.\n\tShowLineNumbers *bool `json:\"-n,omitempty\" jsonschema:\"description=Show line numbers in output (rg -n). Requires output_mode: 'content'\\\\, ignored otherwise. Defaults to true.\"`\n\n\t// CaseInsensitive enables case insensitive search.\n\tCaseInsensitive *bool `json:\"-i,omitempty\" jsonschema:\"description=Case insensitive search (rg -i)\"`\n\n\t// FileType is the file type to search (e.g., js, py, rust, go, java).\n\t// More efficient than Glob for standard file types.\n\tFileType *string `json:\"type,omitempty\" jsonschema:\"description=File type to search (rg --type). Common types: js\\\\, py\\\\, rust\\\\, go\\\\, java\\\\, etc. More efficient than include for standard file types.\"`\n\n\t// HeadLimit limits output to first N lines/entries.\n\t// Works across all output modes. Defaults to 0 (unlimited).\n\tHeadLimit *int `json:\"head_limit,omitempty\" jsonschema:\"description=Limit output to first N lines/entries\\\\, equivalent to '| head -N'. Works across all output modes: content (limits output lines)\\\\, files_with_matches (limits file paths)\\\\, count (limits count entries). Defaults to 0 (unlimited).\"`\n\n\t// Offset skips first N lines/entries before applying HeadLimit.\n\t// Works across all output modes. Defaults to 0.\n\tOffset *int `json:\"offset,omitempty\" jsonschema:\"description=Skip first N lines/entries before applying head_limit\\\\, equivalent to '| tail -n +N | head -N'. Works across all output modes. Defaults to 0.\"`\n\n\t// Multiline enables multiline mode where patterns can span lines.\n\t//   - true: Allows patterns to match across lines, \".\" matches newlines\n\t//   - false: Default, matches only within single lines\n\tMultiline *bool `json:\"multiline,omitempty\" jsonschema:\"description=Enable multiline mode where . matches newlines and patterns can span lines (rg -U --multiline-dotall). Default: false.\"`\n}\n\nfunc newGrepTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) {\n\ttoolName := selectToolName(name, ToolNameGrep)\n\td, err := selectToolDesc(desc, GrepToolDesc, GrepToolDescChinese)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn utils.InferTool(toolName, d, func(ctx context.Context, input grepArgs) (string, error) {\n\t\t// Extract string parameters\n\t\tpath := valueOrDefault(input.Path, \"\")\n\t\tglob := valueOrDefault(input.Glob, \"\")\n\t\tfileType := valueOrDefault(input.FileType, \"\")\n\t\tvar beforeLines, afterLines int\n\n\t\tif input.Context != nil {\n\t\t\tbeforeLines = valueOrDefault(input.Context, 0)\n\t\t\tafterLines = valueOrDefault(input.Context, 0)\n\t\t} else {\n\t\t\t// Extract context parameters\n\t\t\tbeforeLines = valueOrDefault(input.BeforeLines, 0)\n\t\t\tafterLines = valueOrDefault(input.AfterLines, 0)\n\t\t}\n\n\t\t// Extract boolean flags\n\t\tcaseInsensitive := valueOrDefault(input.CaseInsensitive, false)\n\t\tenableMultiline := valueOrDefault(input.Multiline, false)\n\n\t\t// Extract pagination parameters\n\t\theadLimit := valueOrDefault(input.HeadLimit, 0)\n\t\toffset := valueOrDefault(input.Offset, 0)\n\n\t\tmatches, err := fs.GrepRaw(ctx, &filesystem.GrepRequest{\n\t\t\tPattern:         input.Pattern,\n\t\t\tPath:            path,\n\t\t\tGlob:            glob,\n\t\t\tFileType:        fileType,\n\t\t\tCaseInsensitive: caseInsensitive,\n\t\t\tAfterLines:      afterLines,\n\t\t\tBeforeLines:     beforeLines,\n\t\t\tEnableMultiline: enableMultiline,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\tsort.SliceStable(matches, func(i, j int) bool {\n\t\t\treturn filepath.Base(matches[i].Path) < filepath.Base(matches[j].Path)\n\t\t})\n\n\t\tswitch input.OutputMode {\n\t\tcase \"content\":\n\t\t\tmatches = applyPagination(matches, offset, headLimit)\n\t\t\treturn formatContentMatches(matches, valueOrDefault(input.ShowLineNumbers, true)), nil\n\n\t\tcase \"count\":\n\t\t\treturn formatCountMatches(matches, offset, headLimit), nil\n\n\t\tcase \"files_with_matches\":\n\t\t\treturn formatFileMatches(matches, offset, headLimit), nil\n\n\t\tdefault:\n\t\t\treturn formatFileMatches(matches, offset, headLimit), nil\n\t\t}\n\t})\n}\n\ntype executeArgs struct {\n\tCommand string `json:\"command\"`\n}\n\nfunc newExecuteTool(sb filesystem.Shell, name string, desc string) (tool.BaseTool, error) {\n\ttoolName := selectToolName(name, ToolNameExecute)\n\td, err := selectToolDesc(desc, ExecuteToolDesc, ExecuteToolDescChinese)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn utils.InferTool(toolName, d, func(ctx context.Context, input executeArgs) (string, error) {\n\t\tresult, err := sb.Execute(ctx, &filesystem.ExecuteRequest{\n\t\t\tCommand: input.Command,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn convExecuteResponse(result), nil\n\t})\n}\n\nfunc newStreamingExecuteTool(sb filesystem.StreamingShell, name string, desc string) (tool.BaseTool, error) {\n\ttoolName := selectToolName(name, ToolNameExecute)\n\td, err := selectToolDesc(desc, ExecuteToolDesc, ExecuteToolDescChinese)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn utils.InferStreamTool(toolName, d, func(ctx context.Context, input executeArgs) (*schema.StreamReader[string], error) {\n\t\tresult, err := sb.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{\n\t\t\tCommand: input.Command,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tsr, sw := schema.Pipe[string](10)\n\t\tgo func() {\n\t\t\tdefer func() {\n\t\t\t\te := recover()\n\t\t\t\tif e != nil {\n\t\t\t\t\tsw.Send(\"\", fmt.Errorf(\"panic: %v,\\n stack: %s\", e, string(debug.Stack())))\n\t\t\t\t}\n\t\t\t\tsw.Close()\n\t\t\t}()\n\n\t\t\tvar hasSentContent bool\n\t\t\tvar exitCode *int\n\n\t\t\tfor {\n\t\t\t\tchunk, recvErr := result.Recv()\n\t\t\t\tif recvErr == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tif recvErr != nil {\n\t\t\t\t\tsw.Send(\"\", recvErr)\n\t\t\t\t\treturn\n\t\t\t\t}\n\n\t\t\t\tif chunk == nil {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\t\t\t\tif chunk.ExitCode != nil {\n\t\t\t\t\texitCode = chunk.ExitCode\n\t\t\t\t}\n\n\t\t\t\tparts := make([]string, 0, 2)\n\t\t\t\tif chunk.Output != \"\" {\n\t\t\t\t\tparts = append(parts, chunk.Output)\n\t\t\t\t}\n\t\t\t\tif chunk.Truncated {\n\t\t\t\t\tparts = append(parts, \"[Output was truncated due to size limits]\")\n\t\t\t\t}\n\t\t\t\tif len(parts) > 0 {\n\t\t\t\t\tsw.Send(strings.Join(parts, \"\\n\"), nil)\n\t\t\t\t\thasSentContent = true\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tif exitCode != nil && *exitCode != 0 {\n\t\t\t\tsw.Send(fmt.Sprintf(\"\\n[Command failed with exit code %d]\", *exitCode), nil)\n\t\t\t} else if !hasSentContent {\n\t\t\t\tsw.Send(\"[Command executed successfully with no output]\", nil)\n\t\t\t}\n\t\t}()\n\n\t\treturn sr, nil\n\t})\n}\n\nfunc convExecuteResponse(response *filesystem.ExecuteResponse) string {\n\tif response == nil {\n\t\treturn \"\"\n\t}\n\tparts := []string{response.Output}\n\tif response.ExitCode != nil && *response.ExitCode != 0 {\n\t\tparts = append(parts, fmt.Sprintf(\"[Command failed with exit code %d]\", *response.ExitCode))\n\t}\n\tif response.Truncated {\n\t\tparts = append(parts, \"[Output was truncated due to size limits]\")\n\t}\n\n\tresult := strings.Join(parts, \"\\n\")\n\tif result == \"\" && (response.ExitCode == nil || *response.ExitCode == 0) {\n\t\treturn \"[Command executed successfully with no output]\"\n\t}\n\treturn result\n}\n\n// valueOrDefault returns the value pointed to by ptr, or defaultValue if ptr is nil.\nfunc valueOrDefault[T any](ptr *T, defaultValue T) T {\n\tif ptr != nil {\n\t\treturn *ptr\n\t}\n\treturn defaultValue\n}\n\nfunc applyPagination[T any](items []T, offset, headLimit int) []T {\n\tif offset < 0 {\n\t\toffset = 0\n\t}\n\tif offset >= len(items) {\n\t\treturn []T{}\n\t}\n\titems = items[offset:]\n\n\tif headLimit > 0 && headLimit < len(items) {\n\t\titems = items[:headLimit]\n\t}\n\treturn items\n}\n\nfunc formatFileMatches(matches []filesystem.GrepMatch, offset, headLimit int) string {\n\tif len(matches) == 0 {\n\t\treturn noFilesFound\n\t}\n\tseen := make(map[string]bool)\n\tvar uniquePaths []string\n\tfor _, match := range matches {\n\t\tif !seen[match.Path] {\n\t\t\tseen[match.Path] = true\n\t\t\tuniquePaths = append(uniquePaths, match.Path)\n\t\t}\n\t}\n\ttotalFiles := len(uniquePaths)\n\tuniquePaths = applyPagination(uniquePaths, offset, headLimit)\n\n\tfileWord := \"files\"\n\tif totalFiles == 1 {\n\t\tfileWord = \"file\"\n\t}\n\treturn fmt.Sprintf(\"Found %d %s\\n%s\", totalFiles, fileWord, strings.Join(uniquePaths, \"\\n\"))\n}\n\nfunc formatContentMatches(matches []filesystem.GrepMatch, showLineNum bool) string {\n\tif len(matches) == 0 {\n\t\treturn noMatchesFound\n\t}\n\tvar b strings.Builder\n\tfor _, match := range matches {\n\t\tb.WriteString(match.Path)\n\t\tif showLineNum {\n\t\t\tb.WriteString(\":\")\n\t\t\tb.WriteString(strconv.Itoa(match.Line))\n\t\t}\n\t\tb.WriteString(\":\")\n\t\tb.WriteString(match.Content)\n\t\tb.WriteString(\"\\n\")\n\t}\n\treturn strings.TrimSuffix(b.String(), \"\\n\")\n}\n\nfunc formatCountMatches(matches []filesystem.GrepMatch, offset, headLimit int) string {\n\tcountMap := make(map[string]int)\n\tfor _, match := range matches {\n\t\tcountMap[match.Path]++\n\t}\n\n\tvar paths []string\n\tfor path := range countMap {\n\t\tpaths = append(paths, path)\n\t}\n\tsort.Strings(paths)\n\n\ttotalOccurrences := len(matches)\n\ttotalFiles := len(paths)\n\n\toccurrenceWord := \"occurrences\"\n\tif totalOccurrences == 1 {\n\t\toccurrenceWord = \"occurrence\"\n\t}\n\tfileWord := \"files\"\n\tif totalFiles == 1 {\n\t\tfileWord = \"file\"\n\t}\n\n\tif totalOccurrences == 0 {\n\t\treturn fmt.Sprintf(\"%s\\n\\nFound %d total %s across %d %s.\", noMatchesFound, totalOccurrences, occurrenceWord, totalFiles, fileWord)\n\t}\n\n\tpaths = applyPagination(paths, offset, headLimit)\n\n\tvar b strings.Builder\n\tfor _, path := range paths {\n\t\tb.WriteString(path)\n\t\tb.WriteString(\":\")\n\t\tb.WriteString(strconv.Itoa(countMap[path]))\n\t\tb.WriteString(\"\\n\")\n\t}\n\tresult := strings.TrimSuffix(b.String(), \"\\n\")\n\treturn fmt.Sprintf(\"%s\\n\\nFound %d total %s across %d %s.\", result, totalOccurrences, occurrenceWord, totalFiles, fileWord)\n}\n\n// selectToolDesc returns the custom description if provided, otherwise selects the appropriate\n// i18n description based on the current language setting.\nfunc selectToolDesc(customDesc string, defaultEnglish, defaultChinese string) (string, error) {\n\tif customDesc != \"\" {\n\t\treturn customDesc, nil\n\t}\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: defaultEnglish,\n\t\tChinese: defaultChinese,\n\t}), nil\n}\n\n// selectToolName returns the custom tool name if provided, otherwise returns the default name.\nfunc selectToolName(customName string, defaultName string) string {\n\tif customName != \"\" {\n\t\treturn customName\n\t}\n\treturn defaultName\n}\n"
  },
  {
    "path": "adk/middlewares/filesystem/filesystem_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage filesystem\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// setupTestBackend creates a test backend with some initial files\nfunc setupTestBackend() *filesystem.InMemoryBackend {\n\tbackend := filesystem.NewInMemoryBackend()\n\tctx := context.Background()\n\n\t// Create test files\n\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\tFilePath: \"/file1.txt\",\n\t\tContent:  \"line1\\nline2\\nline3\\nline4\\nline5\",\n\t})\n\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\tFilePath: \"/file2.go\",\n\t\tContent:  \"package main\\n\\nfunc main() {\\n\\tprintln(\\\"hello\\\")\\n}\",\n\t})\n\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\tFilePath: \"/dir1/file3.txt\",\n\t\tContent:  \"hello world\\nfoo bar\\nhello again\",\n\t})\n\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\tFilePath: \"/dir1/file4.py\",\n\t\tContent:  \"print('hello')\\nprint('world')\",\n\t})\n\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\tFilePath: \"/dir2/file5.go\",\n\t\tContent:  \"package test\\n\\nfunc test() {}\",\n\t})\n\n\treturn backend\n}\n\n// invokeTool is a helper to invoke a tool with JSON input\nfunc invokeTool(_ *testing.T, bt tool.BaseTool, input string) (string, error) {\n\tctx := context.Background()\n\tresult, err := bt.(tool.InvokableTool).InvokableRun(ctx, input)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn result, nil\n}\n\nfunc TestLsTool(t *testing.T) {\n\tbackend := setupTestBackend()\n\tlsTool, err := newLsTool(backend, \"\", \"\")\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to create ls tool: %v\", err)\n\t}\n\n\ttests := []struct {\n\t\tname     string\n\t\tinput    string\n\t\texpected []string // expected paths in output\n\t}{\n\t\t{\n\t\t\tname:     \"list root\",\n\t\t\tinput:    `{\"path\": \"/\"}`,\n\t\t\texpected: []string{\"file1.txt\", \"file2.go\", \"dir1\", \"dir2\"},\n\t\t},\n\t\t{\n\t\t\tname:     \"list empty path (defaults to root)\",\n\t\t\tinput:    `{\"path\": \"\"}`,\n\t\t\texpected: []string{\"file1.txt\", \"file2.go\", \"dir1\", \"dir2\"},\n\t\t},\n\t\t{\n\t\t\tname:     \"list dir1\",\n\t\t\tinput:    `{\"path\": \"/dir1\"}`,\n\t\t\texpected: []string{\"file3.txt\", \"file4.py\"},\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := invokeTool(t, lsTool, tt.input)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"ls tool failed: %v\", err)\n\t\t\t}\n\n\t\t\tfor _, expectedPath := range tt.expected {\n\t\t\t\tif !strings.Contains(result, expectedPath) {\n\t\t\t\t\tt.Errorf(\"Expected output to contain %q, got: %s\", expectedPath, result)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestReadFileTool(t *testing.T) {\n\tbackend := setupTestBackend()\n\treadTool, err := newReadFileTool(backend, \"\", \"\")\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to create read_file tool: %v\", err)\n\t}\n\n\ttests := []struct {\n\t\tname        string\n\t\tinput       string\n\t\texpected    string\n\t\tshouldError bool\n\t}{\n\t\t{\n\t\t\tname:     \"read full file\",\n\t\t\tinput:    `{\"file_path\": \"/file1.txt\", \"offset\": 0, \"limit\": 100}`,\n\t\t\texpected: \"     1\\tline1\\n     2\\tline2\\n     3\\tline3\\n     4\\tline4\\n     5\\tline5\",\n\t\t},\n\t\t{\n\t\t\tname:     \"read with offset\",\n\t\t\tinput:    `{\"file_path\": \"/file1.txt\", \"offset\": 2, \"limit\": 2}`,\n\t\t\texpected: \"     2\\tline2\\n     3\\tline3\",\n\t\t},\n\t\t{\n\t\t\tname:     \"read with default limit\",\n\t\t\tinput:    `{\"file_path\": \"/file1.txt\", \"offset\": 0, \"limit\": 0}`,\n\t\t\texpected: \"     1\\tline1\\n     2\\tline2\\n     3\\tline3\\n     4\\tline4\\n     5\\tline5\",\n\t\t},\n\t\t{\n\t\t\tname:     \"read with negative offset (treated as 0)\",\n\t\t\tinput:    `{\"file_path\": \"/file1.txt\", \"offset\": -1, \"limit\": 2}`,\n\t\t\texpected: \"     1\\tline1\\n     2\\tline2\",\n\t\t},\n\t\t{\n\t\t\tname:        \"read non-existent file\",\n\t\t\tinput:       `{\"file_path\": \"/nonexistent.txt\", \"offset\": 0, \"limit\": 10}`,\n\t\t\tshouldError: true,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := invokeTool(t, readTool, tt.input)\n\t\t\tif tt.shouldError {\n\t\t\t\tif err == nil {\n\t\t\t\t\tt.Error(\"Expected error but got none\")\n\t\t\t\t}\n\t\t\t\treturn\n\t\t\t}\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"read_file tool failed: %v\", err)\n\t\t\t}\n\t\t\tif result != tt.expected {\n\t\t\t\tt.Errorf(\"Expected %q, got %q\", tt.expected, result)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestWriteFileTool(t *testing.T) {\n\tbackend := setupTestBackend()\n\twriteTool, err := newWriteFileTool(backend, \"\", \"\")\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to create write_file tool: %v\", err)\n\t}\n\n\ttests := []struct {\n\t\tname     string\n\t\tinput    string\n\t\texpected string\n\t\tisError  bool\n\t}{\n\t\t{\n\t\t\tname:     \"write new file\",\n\t\t\tinput:    `{\"file_path\": \"/newfile.txt\", \"content\": \"new content\"}`,\n\t\t\texpected: \"Updated file /newfile.txt\",\n\t\t},\n\t\t{\n\t\t\tname:     \"overwrite existing file\",\n\t\t\tinput:    `{\"file_path\": \"/file1.txt\", \"content\": \"overwritten\"}`,\n\t\t\tisError:  false,\n\t\t\texpected: \"Updated file /file1.txt\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := invokeTool(t, writeTool, tt.input)\n\t\t\tif tt.isError {\n\t\t\t\tif err == nil {\n\t\t\t\t\tt.Errorf(\"Expected an error, but got none\")\n\t\t\t\t}\n\t\t\t\treturn\n\t\t\t}\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"write_file tool failed: %v\", err)\n\t\t\t}\n\t\t\tif result != tt.expected {\n\t\t\t\tt.Errorf(\"Expected %q, got %q\", tt.expected, result)\n\t\t\t}\n\t\t})\n\t}\n\n\t// Verify the file was actually written\n\tctx := context.Background()\n\tcontent, err := backend.Read(ctx, &filesystem.ReadRequest{\n\t\tFilePath: \"/newfile.txt\",\n\t\tOffset:   0,\n\t\tLimit:    100,\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to read written file: %v\", err)\n\t}\n\tif content.Content != \"new content\" {\n\t\tt.Errorf(\"Expected written content to be 'new content', got %q\", content)\n\t}\n}\n\nfunc TestEditFileTool(t *testing.T) {\n\tbackend := setupTestBackend()\n\teditTool, err := newEditFileTool(backend, \"\", \"\")\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to create edit_file tool: %v\", err)\n\t}\n\n\ttests := []struct {\n\t\tname         string\n\t\tsetupFile    string\n\t\tsetupContent string\n\t\tinput        string\n\t\texpected     string\n\t\tshouldError  bool\n\t}{\n\t\t{\n\t\t\tname:         \"replace first occurrence\",\n\t\t\tsetupFile:    \"/edit1.txt\",\n\t\t\tsetupContent: \"hello world\\nhello again\\nhello world\",\n\t\t\tinput:        `{\"file_path\": \"/edit1.txt\", \"old_string\": \"hello again\", \"new_string\": \"hi\", \"replace_all\": false}`,\n\t\t\texpected:     \"hello world\\nhi\\nhello world\",\n\t\t},\n\t\t{\n\t\t\tname:         \"replace all occurrences\",\n\t\t\tsetupFile:    \"/edit2.txt\",\n\t\t\tsetupContent: \"hello world\\nhello again\\nhello world\",\n\t\t\tinput:        `{\"file_path\": \"/edit2.txt\", \"old_string\": \"hello\", \"new_string\": \"hi\", \"replace_all\": true}`,\n\t\t\texpected:     \"hi world\\nhi again\\nhi world\",\n\t\t},\n\t\t{\n\t\t\tname:         \"non-existent file\",\n\t\t\tsetupFile:    \"\",\n\t\t\tsetupContent: \"\",\n\t\t\tinput:        `{\"file_path\": \"/nonexistent.txt\", \"old_string\": \"old\", \"new_string\": \"new\", \"replace_all\": false}`,\n\t\t\tshouldError:  true,\n\t\t},\n\t\t{\n\t\t\tname:         \"empty old_string\",\n\t\t\tsetupFile:    \"/edit3.txt\",\n\t\t\tsetupContent: \"content\",\n\t\t\tinput:        `{\"file_path\": \"/edit3.txt\", \"old_string\": \"\", \"new_string\": \"new\", \"replace_all\": false}`,\n\t\t\tshouldError:  true,\n\t\t},\n\t}\n\n\tctx := context.Background()\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\t// Setup file if needed\n\t\t\tif tt.setupFile != \"\" {\n\t\t\t\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\t\t\tFilePath: tt.setupFile,\n\t\t\t\t\tContent:  tt.setupContent,\n\t\t\t\t})\n\t\t\t}\n\n\t\t\t_, err := invokeTool(t, editTool, tt.input)\n\t\t\tif tt.shouldError {\n\t\t\t\tif err == nil {\n\t\t\t\t\tt.Error(\"Expected error but got none\")\n\t\t\t\t}\n\t\t\t\treturn\n\t\t\t}\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"edit_file tool failed: %v\", err)\n\t\t\t}\n\t\t\tresult, err := backend.Read(ctx, &filesystem.ReadRequest{\n\t\t\t\tFilePath: tt.setupFile,\n\t\t\t\tOffset:   0,\n\t\t\t\tLimit:    0,\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"edit_file tool failed: %v\", err)\n\t\t\t}\n\t\t\tif result.Content != tt.expected {\n\t\t\t\tt.Errorf(\"Expected %q, got %q\", tt.expected, result.Content)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestGlobTool(t *testing.T) {\n\tbackend := setupTestBackend()\n\tglobTool, err := newGlobTool(backend, \"\", \"\")\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to create glob tool: %v\", err)\n\t}\n\n\ttests := []struct {\n\t\tname     string\n\t\tinput    string\n\t\texpected []string\n\t}{\n\t\t{\n\t\t\tname:     \"match all .txt files in root\",\n\t\t\tinput:    `{\"pattern\": \"*.txt\", \"path\": \"/\"}`,\n\t\t\texpected: []string{\"file1.txt\"},\n\t\t},\n\t\t{\n\t\t\tname:     \"match all .go files in root\",\n\t\t\tinput:    `{\"pattern\": \"*.go\", \"path\": \"/\"}`,\n\t\t\texpected: []string{\"file2.go\"},\n\t\t},\n\t\t{\n\t\t\tname:     \"match all .txt files in dir1\",\n\t\t\tinput:    `{\"pattern\": \"*.txt\", \"path\": \"/dir1\"}`,\n\t\t\texpected: []string{\"file3.txt\"},\n\t\t},\n\t\t{\n\t\t\tname:     \"match all .py files in dir1\",\n\t\t\tinput:    `{\"pattern\": \"*.py\", \"path\": \"/dir1\"}`,\n\t\t\texpected: []string{\"file4.py\"},\n\t\t},\n\n\t\t{\n\t\t\tname:     \"empty path defaults to root\",\n\t\t\tinput:    `{\"pattern\": \"*.go\", \"path\": \"\"}`,\n\t\t\texpected: []string{\"file2.go\"},\n\t\t},\n\n\t\t{\n\t\t\tname:     \"match all .txt files in dir1 in root dir\",\n\t\t\tinput:    `{\"pattern\": \"/dir1/*.txt\", \"path\": \"/\"}`,\n\t\t\texpected: []string{\"/dir1/file3.txt\"},\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := invokeTool(t, globTool, tt.input)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"glob tool failed: %v\", err)\n\t\t\t}\n\n\t\t\tfor _, expectedPath := range tt.expected {\n\t\t\t\tif !strings.Contains(result, expectedPath) {\n\t\t\t\t\tt.Errorf(\"Expected output to contain %q, got: %s\", expectedPath, result)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestGrepTool(t *testing.T) {\n\tbackend := setupTestBackend()\n\tgrepTool, err := newGrepTool(backend, \"\", \"\")\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to create grep tool: %v\", err)\n\t}\n\n\ttests := []struct {\n\t\tname     string\n\t\tinput    string\n\t\texpected string\n\t\tcontains []string\n\t}{\n\t\t{\n\t\t\tname:     \"grep with count mode\",\n\t\t\tinput:    `{\"pattern\": \"hello\", \"output_mode\": \"count\"}`,\n\t\t\texpected: \"/dir1/file3.txt:2\\n/dir1/file4.py:1\\n/file2.go:1\\n\\nFound 4 total occurrences across 3 files.\", // 2 in file3.txt, 1 in file4.py, 1 in file2.go\n\t\t},\n\t\t{\n\t\t\tname:     \"grep with content mode\",\n\t\t\tinput:    `{\"pattern\": \"hello\", \"output_mode\": \"content\"}`,\n\t\t\tcontains: []string{\"/dir1/file3.txt:1:hello world\", \"/dir1/file3.txt:3:hello again\", \"/dir1/file4.py:1:print('hello')\"},\n\t\t},\n\t\t{\n\t\t\tname:     \"grep with files_with_matches mode (default)\",\n\t\t\tinput:    `{\"pattern\": \"hello\", \"output_mode\": \"files_with_matches\"}`,\n\t\t\tcontains: []string{\"/dir1/file3.txt\", \"/dir1/file4.py\"},\n\t\t},\n\t\t{\n\t\t\tname:     \"grep with glob filter\",\n\t\t\tinput:    `{\"pattern\": \"hello\", \"glob\": \"*.txt\", \"output_mode\": \"count\"}`,\n\t\t\texpected: \"/dir1/file3.txt:2\\n\\nFound 2 total occurrences across 1 file.\", // only in file3.txt\n\t\t},\n\t\t{\n\t\t\tname:     \"grep  withpath filter\",\n\t\t\tinput:    `{\"pattern\": \"package\", \"path\": \"/dir2\", \"output_mode\": \"count\"}`,\n\t\t\texpected: \"/dir2/file5.go:1\\n\\nFound 1 total occurrence across 1 file.\", // only in dir2/file5.go\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := invokeTool(t, grepTool, tt.input)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"grep tool failed: %v\", err)\n\t\t\t}\n\n\t\t\tif tt.expected != \"\" {\n\t\t\t\tif result != tt.expected {\n\t\t\t\t\tt.Errorf(\"Expected %q, got %q\", tt.expected, result)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tfor _, expectedStr := range tt.contains {\n\t\t\t\tif !strings.Contains(result, expectedStr) {\n\t\t\t\t\tt.Errorf(\"Expected output to contain %q, got: %s\", expectedStr, result)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestExecuteTool(t *testing.T) {\n\tbackend := setupTestBackend()\n\n\ttests := []struct {\n\t\tname        string\n\t\tresp        *filesystem.ExecuteResponse\n\t\tinput       string\n\t\texpected    string\n\t\tshouldError bool\n\t}{\n\t\t{\n\t\t\tname: \"successful command execution\",\n\t\t\tresp: &filesystem.ExecuteResponse{\n\t\t\t\tOutput:   \"hello world\",\n\t\t\t\tExitCode: ptrOf(0),\n\t\t\t},\n\t\t\tinput:    `{\"command\": \"echo hello world\"}`,\n\t\t\texpected: \"hello world\",\n\t\t},\n\t\t{\n\t\t\tname: \"command with non-zero exit code\",\n\t\t\tresp: &filesystem.ExecuteResponse{\n\t\t\t\tOutput:   \"error: file not found\",\n\t\t\t\tExitCode: ptrOf(1),\n\t\t\t},\n\t\t\tinput:    `{\"command\": \"cat nonexistent.txt\"}`,\n\t\t\texpected: \"error: file not found\\n[Command failed with exit code 1]\",\n\t\t},\n\t\t{\n\t\t\tname: \"command with truncated output\",\n\t\t\tresp: &filesystem.ExecuteResponse{\n\t\t\t\tOutput:    \"partial output...\",\n\t\t\t\tExitCode:  ptrOf(0),\n\t\t\t\tTruncated: true,\n\t\t\t},\n\t\t\tinput:    `{\"command\": \"cat largefile.txt\"}`,\n\t\t\texpected: \"partial output...\\n[Output was truncated due to size limits]\",\n\t\t},\n\t\t{\n\t\t\tname: \"command with both non-zero exit code and truncated output\",\n\t\t\tresp: &filesystem.ExecuteResponse{\n\t\t\t\tOutput:    \"error output...\",\n\t\t\t\tExitCode:  ptrOf(2),\n\t\t\t\tTruncated: true,\n\t\t\t},\n\t\t\tinput:    `{\"command\": \"failing command\"}`,\n\t\t\texpected: \"error output...\\n[Command failed with exit code 2]\\n[Output was truncated due to size limits]\",\n\t\t},\n\t\t{\n\t\t\tname: \"successful command with no output\",\n\t\t\tresp: &filesystem.ExecuteResponse{\n\t\t\t\tOutput:   \"\",\n\t\t\t\tExitCode: ptrOf(0),\n\t\t\t},\n\t\t\tinput:    `{\"command\": \"mkdir /tmp/test\"}`,\n\t\t\texpected: \"[Command executed successfully with no output]\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\texecuteTool, err := newExecuteTool(&mockShellBackend{\n\t\t\t\tBackend: backend,\n\t\t\t\tresp:    tt.resp,\n\t\t\t}, \"\", \"\")\n\t\t\tassert.NoError(t, err)\n\n\t\t\tresult, err := invokeTool(t, executeTool, tt.input)\n\t\t\tif tt.shouldError {\n\t\t\t\tassert.Error(t, err)\n\t\t\t\treturn\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\tassert.Equal(t, tt.expected, result)\n\t\t})\n\t}\n}\n\nfunc ptrOf[T any](t T) *T {\n\treturn &t\n}\n\ntype mockShellBackend struct {\n\tfilesystem.Backend\n\tresp *filesystem.ExecuteResponse\n}\n\nfunc (m *mockShellBackend) Execute(ctx context.Context, req *filesystem.ExecuteRequest) (*filesystem.ExecuteResponse, error) {\n\treturn m.resp, nil\n}\n\nfunc TestGetFilesystemTools(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tt.Run(\"returns 6 tools for regular Backend\", func(t *testing.T) {\n\t\ttools, err := getFilesystemTools(ctx, &MiddlewareConfig{Backend: backend})\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 6)\n\n\t\t// Verify tool names\n\t\ttoolNames := make([]string, 0, len(tools))\n\t\tfor _, to := range tools {\n\t\t\tinfo, _ := to.Info(ctx)\n\t\t\ttoolNames = append(toolNames, info.Name)\n\t\t}\n\t\tassert.Contains(t, toolNames, \"ls\")\n\t\tassert.Contains(t, toolNames, \"read_file\")\n\t\tassert.Contains(t, toolNames, \"write_file\")\n\t\tassert.Contains(t, toolNames, \"edit_file\")\n\t\tassert.Contains(t, toolNames, \"glob\")\n\t\tassert.Contains(t, toolNames, \"grep\")\n\t})\n\n\tt.Run(\"returns 7 tools for Shell\", func(t *testing.T) {\n\t\tshellBackend := &mockShellBackend{\n\t\t\tBackend: backend,\n\t\t\tresp:    &filesystem.ExecuteResponse{Output: \"ok\"},\n\t\t}\n\t\ttools, err := getFilesystemTools(ctx, &MiddlewareConfig{Backend: shellBackend, Shell: shellBackend})\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 7)\n\n\t\t// Verify execute tool is included\n\t\ttoolNames := make([]string, 0, len(tools))\n\t\tfor _, to := range tools {\n\t\t\tinfo, _ := to.Info(ctx)\n\t\t\ttoolNames = append(toolNames, info.Name)\n\t\t}\n\t\tassert.Contains(t, toolNames, \"execute\")\n\t})\n\n\tt.Run(\"custom tool descriptions\", func(t *testing.T) {\n\t\tcustomLsDesc := \"Custom ls description\"\n\t\tcustomReadDesc := \"Custom read description\"\n\n\t\ttools, err := getFilesystemTools(ctx, &MiddlewareConfig{\n\t\t\tBackend:                backend,\n\t\t\tCustomLsToolDesc:       &customLsDesc,\n\t\t\tCustomReadFileToolDesc: &customReadDesc,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 6)\n\n\t\t// Verify custom descriptions are applied\n\t\tfor _, to := range tools {\n\t\t\tinfo, _ := to.Info(ctx)\n\t\t\tif info.Name == \"ls\" {\n\t\t\t\tassert.Equal(t, customLsDesc, info.Desc)\n\t\t\t}\n\t\t\tif info.Name == \"read_file\" {\n\t\t\t\tassert.Equal(t, customReadDesc, info.Desc)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestNew(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tt.Run(\"nil config returns error\", func(t *testing.T) {\n\t\t_, err := New(ctx, nil)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"config should not be nil\")\n\t})\n\n\tt.Run(\"nil backend returns error\", func(t *testing.T) {\n\t\t_, err := New(ctx, &MiddlewareConfig{Backend: nil})\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"backend should not be nil\")\n\t})\n\n\tt.Run(\"valid config with default settings\", func(t *testing.T) {\n\t\tm, err := New(ctx, &MiddlewareConfig{Backend: backend})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, m)\n\n\t\tfm, ok := m.(*filesystemMiddleware)\n\t\tassert.True(t, ok)\n\t\tassert.Len(t, fm.additionalTools, 6)\n\t})\n\n\tt.Run(\"custom system prompt\", func(t *testing.T) {\n\t\tcustomPrompt := \"Custom system prompt\"\n\t\tm, err := New(ctx, &MiddlewareConfig{\n\t\t\tBackend:            backend,\n\t\t\tCustomSystemPrompt: &customPrompt,\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tfm, ok := m.(*filesystemMiddleware)\n\t\tassert.True(t, ok)\n\t\tassert.Equal(t, customPrompt, fm.additionalInstruction)\n\t})\n\n\tt.Run(\"ShellBackend adds execute tool\", func(t *testing.T) {\n\t\tshellBackend := &mockShellBackend{\n\t\t\tBackend: backend,\n\t\t\tresp:    &filesystem.ExecuteResponse{Output: \"ok\"},\n\t\t}\n\t\tm, err := New(ctx, &MiddlewareConfig{Backend: shellBackend, Shell: shellBackend})\n\t\tassert.NoError(t, err)\n\n\t\tfm, ok := m.(*filesystemMiddleware)\n\t\tassert.True(t, ok)\n\t\tassert.Len(t, fm.additionalTools, 7)\n\t})\n}\n\nfunc TestFilesystemMiddleware_BeforeAgent(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tt.Run(\"adds instruction and tools to context\", func(t *testing.T) {\n\t\tm, err := New(ctx, &MiddlewareConfig{Backend: backend})\n\t\tassert.NoError(t, err)\n\n\t\trunCtx := &adk.ChatModelAgentContext{\n\t\t\tInstruction: \"Original instruction\",\n\t\t\tTools:       nil,\n\t\t}\n\n\t\tnewCtx, newRunCtx, err := m.BeforeAgent(ctx, runCtx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, newCtx)\n\t\tassert.NotNil(t, newRunCtx)\n\t\tassert.Contains(t, newRunCtx.Instruction, \"Original instruction\")\n\t\tassert.Len(t, newRunCtx.Tools, 6)\n\t})\n\n\tt.Run(\"nil runCtx returns nil\", func(t *testing.T) {\n\t\tm, err := New(ctx, &MiddlewareConfig{Backend: backend})\n\t\tassert.NoError(t, err)\n\n\t\tnewCtx, newRunCtx, err := m.BeforeAgent(ctx, nil)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, newCtx)\n\t\tassert.Nil(t, newRunCtx)\n\t})\n}\n\nfunc TestFilesystemMiddleware_WrapInvokableToolCall(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tt.Run(\"small result passes through unchanged\", func(t *testing.T) {\n\t\tm, err := New(ctx, &MiddlewareConfig{Backend: backend})\n\t\tassert.NoError(t, err)\n\n\t\tendpoint := func(ctx context.Context, args string, opts ...tool.Option) (string, error) {\n\t\t\treturn \"small result\", nil\n\t\t}\n\n\t\ttCtx := &adk.ToolContext{Name: \"test_tool\", CallID: \"call-1\"}\n\t\twrapped, err := m.WrapInvokableToolCall(ctx, endpoint, tCtx)\n\t\tassert.NoError(t, err)\n\n\t\tresult, err := wrapped(ctx, \"{}\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"small result\", result)\n\t})\n\n}\n\nfunc TestGrepToolWithSortingAndPagination(t *testing.T) {\n\tbackend := filesystem.NewInMemoryBackend()\n\tctx := context.Background()\n\n\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\tFilePath: \"/zebra.txt\",\n\t\tContent:  \"match1\\nmatch2\\nmatch3\",\n\t})\n\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\tFilePath: \"/apple.txt\",\n\t\tContent:  \"match4\\nmatch5\",\n\t})\n\tbackend.Write(ctx, &filesystem.WriteRequest{\n\t\tFilePath: \"/banana.txt\",\n\t\tContent:  \"match6\\nmatch7\\nmatch8\",\n\t})\n\n\tgrepTool, err := newGrepTool(backend, \"\", \"\")\n\tassert.NoError(t, err)\n\n\tt.Run(\"files sorted by basename\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"files_with_matches\"}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 4, len(lines)) // 1 summary + 3 files\n\t\tassert.Contains(t, lines[0], \"Found 3 files\")\n\t\tassert.Contains(t, lines[1], \"apple.txt\")\n\t\tassert.Contains(t, lines[2], \"banana.txt\")\n\t\tassert.Contains(t, lines[3], \"zebra.txt\")\n\t})\n\n\tt.Run(\"files_with_matches with offset\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"files_with_matches\", \"offset\": 1}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 3, len(lines))                // 1 summary + 2 files (pagination applied)\n\t\tassert.Contains(t, lines[0], \"Found 3 files\") // total count before pagination\n\t\tassert.Contains(t, lines[1], \"banana.txt\")\n\t\tassert.Contains(t, lines[2], \"zebra.txt\")\n\t})\n\n\tt.Run(\"files_with_matches with head_limit\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"files_with_matches\", \"head_limit\": 2}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 3, len(lines))                // 1 summary + 2 files (pagination applied)\n\t\tassert.Contains(t, lines[0], \"Found 3 files\") // total count before pagination\n\t\tassert.Contains(t, lines[1], \"apple.txt\")\n\t\tassert.Contains(t, lines[2], \"banana.txt\")\n\t})\n\n\tt.Run(\"files_with_matches with offset and head_limit\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"files_with_matches\", \"offset\": 1, \"head_limit\": 1}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 2, len(lines))                // 1 summary + 1 file (pagination applied)\n\t\tassert.Contains(t, lines[0], \"Found 3 files\") // total count before pagination\n\t\tassert.Contains(t, lines[1], \"banana.txt\")\n\t})\n\n\tt.Run(\"content mode sorted and paginated\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"content\", \"head_limit\": 3}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 3, len(lines))\n\t\tassert.Contains(t, lines[0], \"apple.txt\")\n\t})\n\n\tt.Run(\"content mode with offset\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"content\", \"offset\": 2, \"head_limit\": 2}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 2, len(lines))\n\t})\n\n\tt.Run(\"count mode sorted\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"count\"}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 5, len(lines)) // 3 file counts + 1 empty line + 1 summary line\n\t\tassert.Contains(t, lines[0], \"apple.txt:2\")\n\t\tassert.Contains(t, lines[1], \"banana.txt:3\")\n\t\tassert.Contains(t, lines[2], \"zebra.txt:3\")\n\t\tassert.Contains(t, lines[4], \"Found 8 total occurrences across 3 files.\")\n\t})\n\n\tt.Run(\"count mode with pagination\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"count\", \"offset\": 1, \"head_limit\": 1}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 3, len(lines)) // 1 file count + 1 empty line + 1 summary line\n\t\tassert.Contains(t, lines[0], \"banana.txt:3\")\n\t\tassert.Contains(t, lines[2], \"Found 8 total occurrences across 3 files.\") // summary shows total before pagination\n\t})\n\n\tt.Run(\"offset exceeds result count\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"files_with_matches\", \"offset\": 100}`)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, result, \"Found 3 files\") // still shows total count\n\t})\n\n\tt.Run(\"negative offset treated as zero\", func(t *testing.T) {\n\t\tresult, err := invokeTool(t, grepTool, `{\"pattern\": \"match\", \"output_mode\": \"files_with_matches\", \"offset\": -5}`)\n\t\tassert.NoError(t, err)\n\t\tlines := strings.Split(strings.TrimSpace(result), \"\\n\")\n\t\tassert.Equal(t, 4, len(lines)) // 1 summary + 3 files\n\t})\n}\n\nfunc TestApplyPagination(t *testing.T) {\n\tt.Run(\"basic pagination\", func(t *testing.T) {\n\t\titems := []string{\"a\", \"b\", \"c\", \"d\", \"e\"}\n\t\tresult := applyPagination(items, 0, 3)\n\t\tassert.Equal(t, []string{\"a\", \"b\", \"c\"}, result)\n\t})\n\n\tt.Run(\"with offset\", func(t *testing.T) {\n\t\titems := []string{\"a\", \"b\", \"c\", \"d\", \"e\"}\n\t\tresult := applyPagination(items, 2, 2)\n\t\tassert.Equal(t, []string{\"c\", \"d\"}, result)\n\t})\n\n\tt.Run(\"offset exceeds length\", func(t *testing.T) {\n\t\titems := []string{\"a\", \"b\", \"c\"}\n\t\tresult := applyPagination(items, 10, 5)\n\t\tassert.Equal(t, []string{}, result)\n\t})\n\n\tt.Run(\"negative offset\", func(t *testing.T) {\n\t\titems := []string{\"a\", \"b\", \"c\"}\n\t\tresult := applyPagination(items, -1, 2)\n\t\tassert.Equal(t, []string{\"a\", \"b\"}, result)\n\t})\n\n\tt.Run(\"zero head limit means no limit\", func(t *testing.T) {\n\t\titems := []string{\"a\", \"b\", \"c\", \"d\", \"e\"}\n\t\tresult := applyPagination(items, 1, 0)\n\t\tassert.Equal(t, []string{\"b\", \"c\", \"d\", \"e\"}, result)\n\t})\n}\n\nfunc TestCustomToolNames(t *testing.T) {\n\tbackend := setupTestBackend()\n\tctx := context.Background()\n\n\tt.Run(\"custom tool names applied to individual tools\", func(t *testing.T) {\n\t\tcustomLsName := \"list_files\"\n\t\tcustomReadName := \"read\"\n\t\tcustomWriteName := \"write\"\n\t\tcustomEditName := \"edit\"\n\t\tcustomGlobName := \"find_files\"\n\t\tcustomGrepName := \"search\"\n\n\t\tlsTool, err := newLsTool(backend, customLsName, \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ := lsTool.Info(ctx)\n\t\tassert.Equal(t, \"list_files\", info.Name)\n\n\t\treadTool, err := newReadFileTool(backend, customReadName, \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = readTool.Info(ctx)\n\t\tassert.Equal(t, \"read\", info.Name)\n\n\t\twriteTool, err := newWriteFileTool(backend, customWriteName, \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = writeTool.Info(ctx)\n\t\tassert.Equal(t, \"write\", info.Name)\n\n\t\teditTool, err := newEditFileTool(backend, customEditName, \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = editTool.Info(ctx)\n\t\tassert.Equal(t, \"edit\", info.Name)\n\n\t\tglobTool, err := newGlobTool(backend, customGlobName, \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = globTool.Info(ctx)\n\t\tassert.Equal(t, \"find_files\", info.Name)\n\n\t\tgrepTool, err := newGrepTool(backend, customGrepName, \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = grepTool.Info(ctx)\n\t\tassert.Equal(t, \"search\", info.Name)\n\t})\n\n\tt.Run(\"default tool names when custom names not provided\", func(t *testing.T) {\n\t\tlsTool, err := newLsTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ := lsTool.Info(ctx)\n\t\tassert.Equal(t, ToolNameLs, info.Name)\n\n\t\treadTool, err := newReadFileTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = readTool.Info(ctx)\n\t\tassert.Equal(t, ToolNameReadFile, info.Name)\n\n\t\twriteTool, err := newWriteFileTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = writeTool.Info(ctx)\n\t\tassert.Equal(t, ToolNameWriteFile, info.Name)\n\n\t\teditTool, err := newEditFileTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = editTool.Info(ctx)\n\t\tassert.Equal(t, ToolNameEditFile, info.Name)\n\n\t\tglobTool, err := newGlobTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = globTool.Info(ctx)\n\t\tassert.Equal(t, ToolNameGlob, info.Name)\n\n\t\tgrepTool, err := newGrepTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ = grepTool.Info(ctx)\n\t\tassert.Equal(t, ToolNameGrep, info.Name)\n\t})\n\n\tt.Run(\"custom execute tool name\", func(t *testing.T) {\n\t\tcustomExecuteName := \"run_command\"\n\t\tshellBackend := &mockShellBackend{\n\t\t\tBackend: backend,\n\t\t\tresp:    &filesystem.ExecuteResponse{Output: \"ok\"},\n\t\t}\n\n\t\texecuteTool, err := newExecuteTool(shellBackend, customExecuteName, \"\")\n\t\tassert.NoError(t, err)\n\t\tinfo, _ := executeTool.Info(ctx)\n\t\tassert.Equal(t, \"run_command\", info.Name)\n\t})\n\n\tt.Run(\"custom tool names via ToolConfig in getFilesystemTools\", func(t *testing.T) {\n\t\tcustomLsName := \"list_files\"\n\t\tcustomReadName := \"read\"\n\n\t\ttools, err := getFilesystemTools(ctx, &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tName: customLsName,\n\t\t\t},\n\t\t\tReadFileToolConfig: &ToolConfig{\n\t\t\t\tName: customReadName,\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\ttoolNames := make(map[string]bool)\n\t\tfor _, to := range tools {\n\t\t\tinfo, _ := to.Info(ctx)\n\t\t\ttoolNames[info.Name] = true\n\t\t}\n\n\t\tassert.True(t, toolNames[\"list_files\"])\n\t\tassert.True(t, toolNames[\"read\"])\n\t})\n\n\tt.Run(\"custom tool names via ToolConfig in middleware\", func(t *testing.T) {\n\t\tcustomLsName := \"list_files\"\n\t\tcustomReadName := \"read\"\n\n\t\tm, err := New(ctx, &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tName: customLsName,\n\t\t\t},\n\t\t\tReadFileToolConfig: &ToolConfig{\n\t\t\t\tName: customReadName,\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tfm, ok := m.(*filesystemMiddleware)\n\t\tassert.True(t, ok)\n\n\t\ttoolNames := make(map[string]bool)\n\t\tfor _, to := range fm.additionalTools {\n\t\t\tinfo, _ := to.Info(ctx)\n\t\t\ttoolNames[info.Name] = true\n\t\t}\n\n\t\tassert.True(t, toolNames[\"list_files\"])\n\t\tassert.True(t, toolNames[\"read\"])\n\t})\n}\n\nfunc TestSelectToolName(t *testing.T) {\n\tt.Run(\"returns custom name when provided\", func(t *testing.T) {\n\t\tcustomName := \"custom_tool\"\n\t\tresult := selectToolName(customName, \"default_tool\")\n\t\tassert.Equal(t, \"custom_tool\", result)\n\t})\n\n\tt.Run(\"returns default name when custom name is nil\", func(t *testing.T) {\n\t\tresult := selectToolName(\"\", \"default_tool\")\n\t\tassert.Equal(t, \"default_tool\", result)\n\t})\n}\n\nfunc TestGetOrCreateTool(t *testing.T) {\n\tbackend := setupTestBackend()\n\n\tt.Run(\"returns custom tool when provided\", func(t *testing.T) {\n\t\tcustomTool, err := newLsTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tresult, err := getOrCreateTool(customTool, func() (tool.BaseTool, error) {\n\t\t\tt.Fatal(\"createFunc should not be called when custom tool is provided\")\n\t\t\treturn nil, nil\n\t\t})\n\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, customTool, result)\n\t})\n\n\tt.Run(\"calls createFunc when custom tool is nil\", func(t *testing.T) {\n\t\texpectedTool, err := newReadFileTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tcreateFuncCalled := false\n\t\tresult, err := getOrCreateTool(nil, func() (tool.BaseTool, error) {\n\t\t\tcreateFuncCalled = true\n\t\t\treturn expectedTool, nil\n\t\t})\n\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, createFuncCalled, \"createFunc should be called when custom tool is nil\")\n\t\tassert.Equal(t, expectedTool, result)\n\t})\n\n\tt.Run(\"returns nil when custom tool is nil and createFunc returns nil\", func(t *testing.T) {\n\t\tresult, err := getOrCreateTool(nil, func() (tool.BaseTool, error) {\n\t\t\treturn nil, nil\n\t\t})\n\n\t\tassert.NoError(t, err)\n\t\tassert.Nil(t, result)\n\t})\n\n\tt.Run(\"propagates error from createFunc\", func(t *testing.T) {\n\t\texpectedErr := assert.AnError\n\n\t\tresult, err := getOrCreateTool(nil, func() (tool.BaseTool, error) {\n\t\t\treturn nil, expectedErr\n\t\t})\n\n\t\tassert.Error(t, err)\n\t\tassert.Equal(t, expectedErr, err)\n\t\tassert.Nil(t, result)\n\t})\n}\n\nfunc TestCustomTools(t *testing.T) {\n\tbackend := setupTestBackend()\n\tctx := context.Background()\n\n\tt.Run(\"custom ls tool is used via ToolConfig\", func(t *testing.T) {\n\t\tcustomLsTool, err := newLsTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customLsTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 1)\n\t\tassert.Equal(t, customLsTool, tools[0])\n\t})\n\n\tt.Run(\"custom read file tool is used via ToolConfig\", func(t *testing.T) {\n\t\tcustomReadTool, err := newReadFileTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tReadFileToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customReadTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 1)\n\t\tassert.Equal(t, customReadTool, tools[0])\n\t})\n\n\tt.Run(\"custom write file tool is used via ToolConfig\", func(t *testing.T) {\n\t\tcustomWriteTool, err := newWriteFileTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tWriteFileToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customWriteTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 1)\n\t\tassert.Equal(t, customWriteTool, tools[0])\n\t})\n\n\tt.Run(\"custom edit file tool is used via ToolConfig\", func(t *testing.T) {\n\t\tcustomEditTool, err := newEditFileTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tEditFileToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customEditTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 1)\n\t\tassert.Equal(t, customEditTool, tools[0])\n\t})\n\n\tt.Run(\"custom glob tool is used via ToolConfig\", func(t *testing.T) {\n\t\tcustomGlobTool, err := newGlobTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tGlobToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customGlobTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 1)\n\t\tassert.Equal(t, customGlobTool, tools[0])\n\t})\n\n\tt.Run(\"custom grep tool is used via ToolConfig\", func(t *testing.T) {\n\t\tcustomGrepTool, err := newGrepTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tGrepToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customGrepTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 1)\n\t\tassert.Equal(t, customGrepTool, tools[0])\n\t})\n\n\tt.Run(\"multiple custom tools can be used together\", func(t *testing.T) {\n\t\tcustomLsTool, err := newLsTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\t\tcustomReadTool, err := newReadFileTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\t\tcustomGlobTool, err := newGlobTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customLsTool,\n\t\t\t},\n\t\t\tReadFileToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customReadTool,\n\t\t\t},\n\t\t\tGlobToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customGlobTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 3)\n\t})\n\n\tt.Run(\"custom tools take precedence over backend\", func(t *testing.T) {\n\t\tcustomLsTool, err := newLsTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customLsTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tlsToolFound := false\n\t\tfor _, t := range tools {\n\t\t\tif t == customLsTool {\n\t\t\t\tlsToolFound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, lsToolFound, \"custom ls tool should be in the tools list\")\n\t})\n\n\tt.Run(\"backend tools are created when custom tools not provided\", func(t *testing.T) {\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Greater(t, len(tools), 0, \"should create backend tools when custom tools not provided\")\n\t})\n}\n\nfunc TestToolConfig(t *testing.T) {\n\tbackend := setupTestBackend()\n\tctx := context.Background()\n\n\tt.Run(\"use new ToolConfig\", func(t *testing.T) {\n\t\tcustomName := \"my_ls\"\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tName: customName,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 6)\n\n\t\tvar lsToolFound bool\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tif info.Name == \"my_ls\" {\n\t\t\t\tlsToolFound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, lsToolFound)\n\t})\n\n\tt.Run(\"ToolConfig disabled\", func(t *testing.T) {\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tDisable: true,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 5)\n\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tassert.NotEqual(t, ToolNameLs, info.Name)\n\t\t}\n\t})\n\n\tt.Run(\"ToolConfig with custom tool\", func(t *testing.T) {\n\t\tcustomLsTool, err := newLsTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customLsTool,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tvar lsToolFound bool\n\t\tfor _, tool := range tools {\n\t\t\tif tool == customLsTool {\n\t\t\t\tlsToolFound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, lsToolFound)\n\t})\n\n\tt.Run(\"ToolConfig Desc takes precedence over legacy Desc\", func(t *testing.T) {\n\t\tcustomDesc := \"new description\"\n\t\tlegacyDesc := \"old description\"\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tDesc: &customDesc,\n\t\t\t},\n\t\t\tCustomLsToolDesc: &legacyDesc,\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tvar found bool\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tif info.Name == ToolNameLs && info.Desc == \"new description\" {\n\t\t\t\tfound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, found)\n\t})\n\n\tt.Run(\"legacy Desc field still works\", func(t *testing.T) {\n\t\tlegacyDesc := \"legacy description\"\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend:          backend,\n\t\t\tCustomLsToolDesc: &legacyDesc,\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tvar found bool\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tif info.Name == ToolNameLs && info.Desc == \"legacy description\" {\n\t\t\t\tfound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, found)\n\t})\n\n\tt.Run(\"multiple ToolConfig\", func(t *testing.T) {\n\t\tlsName := \"my_ls\"\n\t\treadName := \"my_read\"\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tName: lsName,\n\t\t\t},\n\t\t\tReadFileToolConfig: &ToolConfig{\n\t\t\t\tName: readName,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\ttoolNames := make(map[string]bool)\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\ttoolNames[info.Name] = true\n\t\t}\n\n\t\tassert.True(t, toolNames[\"my_ls\"])\n\t\tassert.True(t, toolNames[\"my_read\"])\n\t})\n}\n\nfunc TestToolConfigEdgeCases(t *testing.T) {\n\tbackend := setupTestBackend()\n\tctx := context.Background()\n\n\tt.Run(\"nil ToolConfig.Desc with nil legacyDesc\", func(t *testing.T) {\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tDesc: nil,\n\t\t\t},\n\t\t\tCustomLsToolDesc: nil,\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tvar lsTool tool.BaseTool\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tif info.Name == ToolNameLs {\n\t\t\t\tlsTool = tool\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.NotNil(t, lsTool, \"ls tool should be created even with nil Desc\")\n\t})\n\n\tt.Run(\"nil ToolConfig.Desc falls back to legacyDesc\", func(t *testing.T) {\n\t\tlegacyDesc := \"legacy description from pointer\"\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tDesc: nil,\n\t\t\t},\n\t\t\tCustomLsToolDesc: &legacyDesc,\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tvar found bool\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tif info.Name == ToolNameLs && info.Desc == \"legacy description from pointer\" {\n\t\t\t\tfound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, found, \"nil ToolConfig.Desc should fall back to legacyDesc\")\n\t})\n\n\tt.Run(\"CustomTool with Disable flag should not create tool\", func(t *testing.T) {\n\t\tcustomLsTool, err := newLsTool(backend, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customLsTool,\n\t\t\t\tDisable:    true,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tassert.NotEqual(t, ToolNameLs, info.Name, \"disabled tool should not be created even if CustomTool is set\")\n\t\t}\n\t})\n\n\tt.Run(\"multiple ToolConfig with conflicting settings\", func(t *testing.T) {\n\t\tlegacyDesc := \"legacy ls desc\"\n\t\tcustomDesc := \"custom desc\"\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tName:    \"custom_ls\",\n\t\t\t\tDesc:    &customDesc,\n\t\t\t\tDisable: false,\n\t\t\t},\n\t\t\tCustomLsToolDesc: &legacyDesc,\n\t\t\tReadFileToolConfig: &ToolConfig{\n\t\t\t\tDisable: true,\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\thasLsTool := false\n\t\thasReadTool := false\n\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tif info.Name == \"custom_ls\" {\n\t\t\t\thasLsTool = true\n\t\t\t\tassert.Equal(t, \"custom desc\", info.Desc, \"ToolConfig.Desc should take precedence over legacy\")\n\t\t\t}\n\t\t\tif info.Name == ToolNameReadFile {\n\t\t\t\thasReadTool = true\n\t\t\t}\n\t\t}\n\n\t\tassert.True(t, hasLsTool, \"ls tool should be created\")\n\t\tassert.False(t, hasReadTool, \"read_file tool should be disabled\")\n\t})\n\n\tt.Run(\"nil ToolConfig with nil legacyDesc creates default tool\", func(t *testing.T) {\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend:          backend,\n\t\t\tLsToolConfig:     nil,\n\t\t\tCustomLsToolDesc: nil,\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tvar lsTool tool.BaseTool\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tif info.Name == ToolNameLs {\n\t\t\t\tlsTool = tool\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.NotNil(t, lsTool, \"tool should be created with backend even when config is nil\")\n\t})\n\n\tt.Run(\"empty Name in ToolConfig uses default name\", func(t *testing.T) {\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tName: \"\",\n\t\t\t},\n\t\t}\n\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tvar lsTool tool.BaseTool\n\t\tfor _, tool := range tools {\n\t\t\tinfo, _ := tool.Info(ctx)\n\t\t\tif info.Name == ToolNameLs {\n\t\t\t\tlsTool = tool\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.NotNil(t, lsTool, \"tool should use default name when Name is empty\")\n\t})\n}\n\nfunc TestGetFilesystemTools_DisableAllTools(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tconfig := &MiddlewareConfig{\n\t\tBackend:             backend,\n\t\tLsToolConfig:        &ToolConfig{Disable: true},\n\t\tReadFileToolConfig:  &ToolConfig{Disable: true},\n\t\tWriteFileToolConfig: &ToolConfig{Disable: true},\n\t\tEditFileToolConfig:  &ToolConfig{Disable: true},\n\t\tGlobToolConfig:      &ToolConfig{Disable: true},\n\t\tGrepToolConfig:      &ToolConfig{Disable: true},\n\t}\n\n\ttools, err := getFilesystemTools(ctx, config)\n\tassert.NoError(t, err)\n\tassert.Len(t, tools, 0)\n}\n\nfunc TestGetFilesystemTools_StreamingShell(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tt.Run(\"returns 7 tools with StreamingShell\", func(t *testing.T) {\n\t\tmockSS := &mockStreamingShell{}\n\t\ttools, err := getFilesystemTools(ctx, &MiddlewareConfig{\n\t\t\tBackend:        backend,\n\t\t\tStreamingShell: mockSS,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 7)\n\n\t\ttoolNames := make([]string, 0, len(tools))\n\t\tfor _, to := range tools {\n\t\t\tinfo, _ := to.Info(ctx)\n\t\t\ttoolNames = append(toolNames, info.Name)\n\t\t}\n\t\tassert.Contains(t, toolNames, ToolNameExecute)\n\t})\n\n\tt.Run(\"StreamingShell takes precedence over Shell\", func(t *testing.T) {\n\t\tmockSS := &mockStreamingShell{}\n\t\tshellBackend := &mockShellBackend{\n\t\t\tBackend: backend,\n\t\t\tresp:    &filesystem.ExecuteResponse{Output: \"ok\"},\n\t\t}\n\n\t\t// When both are set, Validate should fail\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend:        backend,\n\t\t\tShell:          shellBackend,\n\t\t\tStreamingShell: mockSS,\n\t\t}\n\t\terr := config.Validate()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"shell and streaming shell should not be both set\")\n\t})\n}\n\nfunc TestGetFilesystemTools_NilBackend(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"nil backend with shell only returns execute tool\", func(t *testing.T) {\n\t\tmockSS := &mockStreamingShell{}\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend:        nil,\n\t\t\tStreamingShell: mockSS,\n\t\t}\n\t\t// Validate should fail, but getFilesystemTools itself handles nil backend gracefully\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\t// Only execute tool should be returned since backend is nil\n\t\tassert.Len(t, tools, 1)\n\n\t\tinfo, _ := tools[0].Info(ctx)\n\t\tassert.Equal(t, ToolNameExecute, info.Name)\n\t})\n\n\tt.Run(\"nil backend with regular Shell returns execute tool\", func(t *testing.T) {\n\t\tmockShell := &mockShellBackend{\n\t\t\tresp: &filesystem.ExecuteResponse{Output: \"ok\"},\n\t\t}\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: nil,\n\t\t\tShell:   mockShell,\n\t\t}\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 1)\n\n\t\tinfo, _ := tools[0].Info(ctx)\n\t\tassert.Equal(t, ToolNameExecute, info.Name)\n\t})\n\n\tt.Run(\"nil backend and nil shell returns empty tools\", func(t *testing.T) {\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: nil,\n\t\t}\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tools, 0)\n\t})\n}\n\nfunc TestGetFilesystemTools_PartialDisable(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tconfig := &MiddlewareConfig{\n\t\tBackend:            backend,\n\t\tLsToolConfig:       &ToolConfig{Disable: true},\n\t\tReadFileToolConfig: &ToolConfig{Disable: true},\n\t}\n\n\ttools, err := getFilesystemTools(ctx, config)\n\tassert.NoError(t, err)\n\tassert.Len(t, tools, 4)\n\n\ttoolNames := make([]string, 0, len(tools))\n\tfor _, to := range tools {\n\t\tinfo, _ := to.Info(ctx)\n\t\ttoolNames = append(toolNames, info.Name)\n\t}\n\tassert.NotContains(t, toolNames, ToolNameLs)\n\tassert.NotContains(t, toolNames, ToolNameReadFile)\n\tassert.Contains(t, toolNames, ToolNameWriteFile)\n\tassert.Contains(t, toolNames, ToolNameEditFile)\n\tassert.Contains(t, toolNames, ToolNameGlob)\n\tassert.Contains(t, toolNames, ToolNameGrep)\n}\n\ntype mockStreamingShell struct{}\n\nfunc (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\tsr, sw := schema.Pipe[*filesystem.ExecuteResponse](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(&filesystem.ExecuteResponse{\n\t\t\tOutput:   \"streaming output\",\n\t\t\tExitCode: ptrOf(0),\n\t\t}, nil)\n\t}()\n\treturn sr, nil\n}\n\ntype mockStreamingShellWithError struct{}\n\nfunc (m *mockStreamingShellWithError) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\treturn nil, fmt.Errorf(\"streaming shell error\")\n}\n\ntype mockStreamingShellWithRecvError struct{}\n\nfunc (m *mockStreamingShellWithRecvError) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\tsr, sw := schema.Pipe[*filesystem.ExecuteResponse](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(nil, fmt.Errorf(\"recv error during streaming\"))\n\t}()\n\treturn sr, nil\n}\n\ntype mockStreamingShellWithExitCode struct {\n\texitCode int\n}\n\nfunc (m *mockStreamingShellWithExitCode) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\tsr, sw := schema.Pipe[*filesystem.ExecuteResponse](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(&filesystem.ExecuteResponse{\n\t\t\tOutput:   \"some output\",\n\t\t\tExitCode: ptrOf(m.exitCode),\n\t\t}, nil)\n\t}()\n\treturn sr, nil\n}\n\ntype mockStreamingShellNoOutput struct{}\n\nfunc (m *mockStreamingShellNoOutput) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\tsr, sw := schema.Pipe[*filesystem.ExecuteResponse](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(&filesystem.ExecuteResponse{\n\t\t\tExitCode: ptrOf(0),\n\t\t}, nil)\n\t}()\n\treturn sr, nil\n}\n\ntype mockStreamingShellTruncated struct{}\n\nfunc (m *mockStreamingShellTruncated) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\tsr, sw := schema.Pipe[*filesystem.ExecuteResponse](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(&filesystem.ExecuteResponse{\n\t\t\tOutput:    \"partial\",\n\t\t\tTruncated: true,\n\t\t\tExitCode:  ptrOf(0),\n\t\t}, nil)\n\t}()\n\treturn sr, nil\n}\n\ntype mockStreamingShellNilChunk struct{}\n\nfunc (m *mockStreamingShellNilChunk) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\tsr, sw := schema.Pipe[*filesystem.ExecuteResponse](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(nil, nil)\n\t\tsw.Send(&filesystem.ExecuteResponse{\n\t\t\tOutput:   \"after nil\",\n\t\t\tExitCode: ptrOf(0),\n\t\t}, nil)\n\t}()\n\treturn sr, nil\n}\n\nfunc TestNewStreamingExecuteTool(t *testing.T) {\n\tt.Run(\"successful streaming execution\", func(t *testing.T) {\n\t\texecuteTool, err := newStreamingExecuteTool(&mockStreamingShell{}, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tst := executeTool.(tool.StreamableTool)\n\t\tsr, err := st.StreamableRun(context.Background(), `{\"command\": \"echo hello\"}`)\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tvar chunks []string\n\t\tfor {\n\t\t\tchunk, recvErr := sr.Recv()\n\t\t\tif recvErr == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, recvErr)\n\t\t\tchunks = append(chunks, chunk)\n\t\t}\n\t\tassert.True(t, len(chunks) > 0)\n\t\tresult := strings.Join(chunks, \"\")\n\t\tassert.Contains(t, result, \"streaming output\")\n\t})\n\n\tt.Run(\"streaming execution with ExecuteStreaming error\", func(t *testing.T) {\n\t\texecuteTool, err := newStreamingExecuteTool(&mockStreamingShellWithError{}, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tst := executeTool.(tool.StreamableTool)\n\t\t_, err = st.StreamableRun(context.Background(), `{\"command\": \"fail\"}`)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"streaming shell error\")\n\t})\n\n\tt.Run(\"streaming execution with recv error\", func(t *testing.T) {\n\t\texecuteTool, err := newStreamingExecuteTool(&mockStreamingShellWithRecvError{}, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tst := executeTool.(tool.StreamableTool)\n\t\tsr, err := st.StreamableRun(context.Background(), `{\"command\": \"echo hello\"}`)\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tvar gotError bool\n\t\tfor {\n\t\t\t_, recvErr := sr.Recv()\n\t\t\tif recvErr == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif recvErr != nil {\n\t\t\t\tgotError = true\n\t\t\t\tassert.Contains(t, recvErr.Error(), \"recv error during streaming\")\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, gotError)\n\t})\n\n\tt.Run(\"streaming execution with non-zero exit code\", func(t *testing.T) {\n\t\texecuteTool, err := newStreamingExecuteTool(&mockStreamingShellWithExitCode{exitCode: 1}, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tst := executeTool.(tool.StreamableTool)\n\t\tsr, err := st.StreamableRun(context.Background(), `{\"command\": \"false\"}`)\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tvar chunks []string\n\t\tfor {\n\t\t\tchunk, recvErr := sr.Recv()\n\t\t\tif recvErr == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, recvErr)\n\t\t\tchunks = append(chunks, chunk)\n\t\t}\n\t\tresult := strings.Join(chunks, \"\")\n\t\tassert.Contains(t, result, \"[Command failed with exit code 1]\")\n\t})\n\n\tt.Run(\"streaming execution with zero exit code and no output\", func(t *testing.T) {\n\t\texecuteTool, err := newStreamingExecuteTool(&mockStreamingShellNoOutput{}, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tst := executeTool.(tool.StreamableTool)\n\t\tsr, err := st.StreamableRun(context.Background(), `{\"command\": \"true\"}`)\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tvar chunks []string\n\t\tfor {\n\t\t\tchunk, recvErr := sr.Recv()\n\t\t\tif recvErr == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, recvErr)\n\t\t\tchunks = append(chunks, chunk)\n\t\t}\n\t\tresult := strings.Join(chunks, \"\")\n\t\tassert.Contains(t, result, \"[Command executed successfully with no output]\")\n\t})\n\n\tt.Run(\"streaming execution with truncated output\", func(t *testing.T) {\n\t\texecuteTool, err := newStreamingExecuteTool(&mockStreamingShellTruncated{}, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tst := executeTool.(tool.StreamableTool)\n\t\tsr, err := st.StreamableRun(context.Background(), `{\"command\": \"cat largefile\"}`)\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tvar chunks []string\n\t\tfor {\n\t\t\tchunk, recvErr := sr.Recv()\n\t\t\tif recvErr == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, recvErr)\n\t\t\tchunks = append(chunks, chunk)\n\t\t}\n\t\tresult := strings.Join(chunks, \"\")\n\t\tassert.Contains(t, result, \"partial\")\n\t\tassert.Contains(t, result, \"[Output was truncated due to size limits]\")\n\t})\n\n\tt.Run(\"streaming execution with nil chunk skipped\", func(t *testing.T) {\n\t\texecuteTool, err := newStreamingExecuteTool(&mockStreamingShellNilChunk{}, \"\", \"\")\n\t\tassert.NoError(t, err)\n\n\t\tst := executeTool.(tool.StreamableTool)\n\t\tsr, err := st.StreamableRun(context.Background(), `{\"command\": \"echo test\"}`)\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tvar chunks []string\n\t\tfor {\n\t\t\tchunk, recvErr := sr.Recv()\n\t\t\tif recvErr == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, recvErr)\n\t\t\tchunks = append(chunks, chunk)\n\t\t}\n\t\tresult := strings.Join(chunks, \"\")\n\t\tassert.Contains(t, result, \"after nil\")\n\t})\n\n\tt.Run(\"streaming execution with custom name and desc\", func(t *testing.T) {\n\t\texecuteTool, err := newStreamingExecuteTool(&mockStreamingShell{}, \"custom_execute\", \"custom desc\")\n\t\tassert.NoError(t, err)\n\n\t\tinfo, err := executeTool.Info(context.Background())\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"custom_execute\", info.Name)\n\t\tassert.Equal(t, \"custom desc\", info.Desc)\n\t})\n}\n\nfunc TestNew_StreamingShell(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tt.Run(\"StreamingShell adds streaming execute tool\", func(t *testing.T) {\n\t\tm, err := New(ctx, &MiddlewareConfig{\n\t\t\tBackend:        backend,\n\t\t\tStreamingShell: &mockStreamingShell{},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tfm, ok := m.(*filesystemMiddleware)\n\t\tassert.True(t, ok)\n\t\tassert.Len(t, fm.additionalTools, 7)\n\t})\n\n\tt.Run(\"both Shell and StreamingShell returns error\", func(t *testing.T) {\n\t\t_, err := New(ctx, &MiddlewareConfig{\n\t\t\tBackend:        backend,\n\t\t\tShell:          &mockShellBackend{Backend: backend, resp: &filesystem.ExecuteResponse{Output: \"ok\"}},\n\t\t\tStreamingShell: &mockStreamingShell{},\n\t\t})\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"shell and streaming shell should not be both set\")\n\t})\n}\n\nfunc TestNewMiddleware_Validation(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"nil config returns error\", func(t *testing.T) {\n\t\t_, err := NewMiddleware(ctx, nil)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"config should not be nil\")\n\t})\n\n\tt.Run(\"nil backend returns error\", func(t *testing.T) {\n\t\t_, err := NewMiddleware(ctx, &Config{Backend: nil})\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"backend should not be nil\")\n\t})\n\n\tt.Run(\"both Shell and StreamingShell returns error\", func(t *testing.T) {\n\t\tbackend := setupTestBackend()\n\t\t_, err := NewMiddleware(ctx, &Config{\n\t\t\tBackend:        backend,\n\t\t\tShell:          &mockShellBackend{Backend: backend, resp: &filesystem.ExecuteResponse{Output: \"ok\"}},\n\t\t\tStreamingShell: &mockStreamingShell{},\n\t\t})\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"shell and streaming shell should not be both set\")\n\t})\n}\n\nfunc TestMiddlewareConfig_Validate(t *testing.T) {\n\tt.Run(\"nil config returns error\", func(t *testing.T) {\n\t\tvar c *MiddlewareConfig\n\t\terr := c.Validate()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"config should not be nil\")\n\t})\n\n\tt.Run(\"nil backend returns error\", func(t *testing.T) {\n\t\tc := &MiddlewareConfig{}\n\t\terr := c.Validate()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"backend should not be nil\")\n\t})\n\n\tt.Run(\"both shells returns error\", func(t *testing.T) {\n\t\tc := &MiddlewareConfig{\n\t\t\tBackend:        setupTestBackend(),\n\t\t\tShell:          &mockShellBackend{},\n\t\t\tStreamingShell: &mockStreamingShell{},\n\t\t}\n\t\terr := c.Validate()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"shell and streaming shell should not be both set\")\n\t})\n\n\tt.Run(\"valid config passes\", func(t *testing.T) {\n\t\tc := &MiddlewareConfig{\n\t\t\tBackend: setupTestBackend(),\n\t\t}\n\t\terr := c.Validate()\n\t\tassert.NoError(t, err)\n\t})\n}\n\nfunc TestNewStreamingExecuteTool_MultipleChunks(t *testing.T) {\n\tmockSS := &mockStreamingShellMultiChunk{}\n\texecuteTool, err := newStreamingExecuteTool(mockSS, \"\", \"\")\n\tassert.NoError(t, err)\n\n\tst := executeTool.(tool.StreamableTool)\n\tsr, err := st.StreamableRun(context.Background(), `{\"command\": \"long-running\"}`)\n\tassert.NoError(t, err)\n\tdefer sr.Close()\n\n\tvar chunks []string\n\tfor {\n\t\tchunk, recvErr := sr.Recv()\n\t\tif recvErr == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, recvErr)\n\t\tchunks = append(chunks, chunk)\n\t}\n\t// Should have received multiple chunks\n\tassert.True(t, len(chunks) >= 3)\n\tresult := strings.Join(chunks, \"\")\n\tassert.Contains(t, result, \"chunk1\")\n\tassert.Contains(t, result, \"chunk2\")\n\tassert.Contains(t, result, \"chunk3\")\n}\n\ntype mockStreamingShellMultiChunk struct{}\n\nfunc (m *mockStreamingShellMultiChunk) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\tsr, sw := schema.Pipe[*filesystem.ExecuteResponse](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(&filesystem.ExecuteResponse{Output: \"chunk1\\n\"}, nil)\n\t\tsw.Send(&filesystem.ExecuteResponse{Output: \"chunk2\\n\"}, nil)\n\t\tsw.Send(&filesystem.ExecuteResponse{Output: \"chunk3\\n\", ExitCode: ptrOf(0)}, nil)\n\t}()\n\treturn sr, nil\n}\n\nfunc TestNewStreamingExecuteTool_ExitCodeOnlyInLastChunk(t *testing.T) {\n\tmockSS := &mockStreamingShellExitCodeLast{exitCode: 2}\n\texecuteTool, err := newStreamingExecuteTool(mockSS, \"\", \"\")\n\tassert.NoError(t, err)\n\n\tst := executeTool.(tool.StreamableTool)\n\tsr, err := st.StreamableRun(context.Background(), `{\"command\": \"fail-at-end\"}`)\n\tassert.NoError(t, err)\n\tdefer sr.Close()\n\n\tvar chunks []string\n\tfor {\n\t\tchunk, recvErr := sr.Recv()\n\t\tif recvErr == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, recvErr)\n\t\tchunks = append(chunks, chunk)\n\t}\n\tresult := strings.Join(chunks, \"\")\n\tassert.Contains(t, result, \"output line\")\n\tassert.Contains(t, result, \"[Command failed with exit code 2]\")\n}\n\ntype mockStreamingShellExitCodeLast struct {\n\texitCode int\n}\n\nfunc (m *mockStreamingShellExitCodeLast) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {\n\tsr, sw := schema.Pipe[*filesystem.ExecuteResponse](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(&filesystem.ExecuteResponse{Output: \"output line\"}, nil)\n\t\tsw.Send(&filesystem.ExecuteResponse{ExitCode: ptrOf(m.exitCode)}, nil)\n\t}()\n\treturn sr, nil\n}\n\nfunc TestConvExecuteResponse_NilResponse(t *testing.T) {\n\tresult := convExecuteResponse(nil)\n\tassert.Equal(t, \"\", result)\n}\n\nfunc TestConvExecuteResponse_NilExitCode(t *testing.T) {\n\tresult := convExecuteResponse(&filesystem.ExecuteResponse{\n\t\tOutput: \"some output\",\n\t})\n\tassert.Equal(t, \"some output\", result)\n}\n\nfunc TestConfig_Validate(t *testing.T) {\n\tt.Run(\"nil config returns error\", func(t *testing.T) {\n\t\tvar c *Config\n\t\terr := c.Validate()\n\t\tassert.Error(t, err)\n\t})\n\n\tt.Run(\"nil backend returns error\", func(t *testing.T) {\n\t\tc := &Config{}\n\t\terr := c.Validate()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"backend should not be nil\")\n\t})\n\n\tt.Run(\"both shells returns error\", func(t *testing.T) {\n\t\tc := &Config{\n\t\t\tBackend:        setupTestBackend(),\n\t\t\tShell:          &mockShellBackend{},\n\t\t\tStreamingShell: &mockStreamingShell{},\n\t\t}\n\t\terr := c.Validate()\n\t\tassert.Error(t, err)\n\t})\n\n\tt.Run(\"valid config passes\", func(t *testing.T) {\n\t\tc := &Config{\n\t\t\tBackend: setupTestBackend(),\n\t\t}\n\t\terr := c.Validate()\n\t\tassert.NoError(t, err)\n\t})\n}\n\nfunc TestGetFilesystemTools_CustomToolWithShell(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tt.Run(\"custom tool replaces default for all disabled except custom\", func(t *testing.T) {\n\t\tcustomLs, err := newLsTool(backend, \"my_ls\", \"my ls desc\")\n\t\tassert.NoError(t, err)\n\n\t\tconfig := &MiddlewareConfig{\n\t\t\tBackend: backend,\n\t\t\tLsToolConfig: &ToolConfig{\n\t\t\t\tCustomTool: customLs,\n\t\t\t},\n\t\t}\n\t\ttools, err := getFilesystemTools(ctx, config)\n\t\tassert.NoError(t, err)\n\n\t\tvar found bool\n\t\tfor _, to := range tools {\n\t\t\tinfo, _ := to.Info(ctx)\n\t\t\tif info.Name == \"my_ls\" {\n\t\t\t\tfound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.True(t, found)\n\t})\n}\n\nfunc TestMergeToolConfigWithDesc(t *testing.T) {\n\tconfig := &MiddlewareConfig{Backend: setupTestBackend()}\n\n\tt.Run(\"both nil returns empty ToolConfig\", func(t *testing.T) {\n\t\tresult := config.mergeToolConfigWithDesc(nil, nil)\n\t\tassert.NotNil(t, result)\n\t\tassert.Equal(t, \"\", result.Name)\n\t\tassert.Nil(t, result.Desc)\n\t\tassert.False(t, result.Disable)\n\t})\n\n\tt.Run(\"nil toolConfig with legacyDesc\", func(t *testing.T) {\n\t\tdesc := \"legacy\"\n\t\tresult := config.mergeToolConfigWithDesc(nil, &desc)\n\t\tassert.NotNil(t, result)\n\t\tassert.Equal(t, \"legacy\", *result.Desc)\n\t})\n\n\tt.Run(\"toolConfig with Desc overrides legacyDesc\", func(t *testing.T) {\n\t\ttcDesc := \"tc desc\"\n\t\tlegacyDesc := \"legacy\"\n\t\ttc := &ToolConfig{Desc: &tcDesc}\n\t\tresult := config.mergeToolConfigWithDesc(tc, &legacyDesc)\n\t\tassert.Equal(t, \"tc desc\", *result.Desc)\n\t})\n\n\tt.Run(\"toolConfig with nil Desc falls back to legacyDesc\", func(t *testing.T) {\n\t\tlegacyDesc := \"legacy\"\n\t\ttc := &ToolConfig{Name: \"custom\"}\n\t\tresult := config.mergeToolConfigWithDesc(tc, &legacyDesc)\n\t\tassert.Equal(t, \"legacy\", *result.Desc)\n\t\tassert.Equal(t, \"custom\", result.Name)\n\t})\n\n\tt.Run(\"toolConfig with nil Desc and nil legacyDesc\", func(t *testing.T) {\n\t\ttc := &ToolConfig{Name: \"custom\"}\n\t\tresult := config.mergeToolConfigWithDesc(tc, nil)\n\t\tassert.Nil(t, result.Desc)\n\t\tassert.Equal(t, \"custom\", result.Name)\n\t})\n}\n\nfunc TestNewMiddleware_WithShell(t *testing.T) {\n\tctx := context.Background()\n\tbackend := setupTestBackend()\n\n\tt.Run(\"Shell backend creates execute tool\", func(t *testing.T) {\n\t\tshellBackend := &mockShellBackend{\n\t\t\tBackend: backend,\n\t\t\tresp:    &filesystem.ExecuteResponse{Output: \"ok\"},\n\t\t}\n\t\tm, err := NewMiddleware(ctx, &Config{\n\t\t\tBackend: backend,\n\t\t\tShell:   shellBackend,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, m.AdditionalTools, 7)\n\t})\n\n\tt.Run(\"StreamingShell backend creates streaming execute tool\", func(t *testing.T) {\n\t\tm, err := NewMiddleware(ctx, &Config{\n\t\t\tBackend:        backend,\n\t\t\tStreamingShell: &mockStreamingShell{},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, m.AdditionalTools, 7)\n\t})\n}\n\nfunc TestNewExecuteTool_ShellError(t *testing.T) {\n\tmockShell := &mockShellBackendWithError{}\n\texecuteTool, err := newExecuteTool(mockShell, \"\", \"\")\n\tassert.NoError(t, err)\n\n\tresult, err := invokeTool(t, executeTool, `{\"command\": \"fail\"}`)\n\tassert.Error(t, err)\n\tassert.Equal(t, \"\", result)\n\tassert.Contains(t, err.Error(), \"shell execution error\")\n}\n\ntype mockShellBackendWithError struct{}\n\nfunc (m *mockShellBackendWithError) Execute(ctx context.Context, req *filesystem.ExecuteRequest) (*filesystem.ExecuteResponse, error) {\n\treturn nil, errors.New(\"shell execution error\")\n}\n"
  },
  {
    "path": "adk/middlewares/filesystem/large_tool_result.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage filesystem\n\nimport (\n\t\"bufio\"\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"strings\"\n\t\"unicode/utf8\"\n\n\t\"github.com/slongfield/pyfmt\"\n\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype toolResultOffloadingConfig struct {\n\tBackend       filesystem.Backend\n\tTokenLimit    int\n\tPathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error)\n}\n\nfunc newToolResultOffloading(ctx context.Context, config *toolResultOffloadingConfig) compose.ToolMiddleware {\n\toffloading := &toolResultOffloading{\n\t\tbackend:       config.Backend,\n\t\ttokenLimit:    config.TokenLimit,\n\t\tpathGenerator: config.PathGenerator,\n\t}\n\n\tif offloading.tokenLimit == 0 {\n\t\toffloading.tokenLimit = 20000\n\t}\n\n\tif offloading.pathGenerator == nil {\n\t\toffloading.pathGenerator = func(ctx context.Context, input *compose.ToolInput) (string, error) {\n\t\t\treturn fmt.Sprintf(\"/large_tool_result/%s\", input.CallID), nil\n\t\t}\n\t}\n\n\treturn compose.ToolMiddleware{\n\t\tInvokable:  offloading.invoke,\n\t\tStreamable: offloading.stream,\n\t}\n}\n\ntype toolResultOffloading struct {\n\tbackend       filesystem.Backend\n\ttokenLimit    int\n\tpathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error)\n}\n\nfunc (t *toolResultOffloading) invoke(endpoint compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {\n\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\toutput, err := endpoint(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tresult, err := t.handleResult(ctx, output.Result, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &compose.ToolOutput{Result: result}, nil\n\t}\n}\n\nfunc (t *toolResultOffloading) stream(endpoint compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {\n\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\toutput, err := endpoint(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tresult, err := concatString(output.Result)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tresult, err = t.handleResult(ctx, result, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &compose.StreamToolOutput{Result: schema.StreamReaderFromArray([]string{result})}, nil\n\t}\n}\n\nfunc (t *toolResultOffloading) handleResult(ctx context.Context, result string, input *compose.ToolInput) (string, error) {\n\tif len(result) > t.tokenLimit*4 {\n\t\tpath, err := t.pathGenerator(ctx, input)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\tnResult := formatToolMessage(result)\n\t\tmsgTemplate := internal.SelectPrompt(internal.I18nPrompts{\n\t\t\tEnglish: tooLargeToolMessage,\n\t\t\tChinese: tooLargeToolMessageChinese,\n\t\t})\n\t\tnResult, err = pyfmt.Fmt(msgTemplate, map[string]any{\n\t\t\t\"tool_call_id\":   input.CallID,\n\t\t\t\"file_path\":      path,\n\t\t\t\"content_sample\": nResult,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\terr = t.backend.Write(ctx, &WriteRequest{\n\t\t\tFilePath: path,\n\t\t\tContent:  result,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn nResult, nil\n\t}\n\n\treturn result, nil\n}\n\nfunc concatString(sr *schema.StreamReader[string]) (string, error) {\n\tif sr == nil {\n\t\treturn \"\", errors.New(\"stream is nil\")\n\t}\n\tsb := strings.Builder{}\n\tfor {\n\t\tstr, err := sr.Recv()\n\t\tif errors.Is(err, io.EOF) {\n\t\t\treturn sb.String(), nil\n\t\t}\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tsb.WriteString(str)\n\t}\n}\n\nfunc formatToolMessage(s string) string {\n\treader := bufio.NewScanner(strings.NewReader(s))\n\tvar b strings.Builder\n\n\tlineNum := 1\n\tfor reader.Scan() {\n\t\tif lineNum > 10 {\n\t\t\tbreak\n\t\t}\n\t\tline := reader.Text()\n\n\t\tif utf8.RuneCountInString(line) > 1000 {\n\t\t\trunes := []rune(line)\n\t\t\tline = string(runes[:1000])\n\t\t}\n\n\t\tb.WriteString(fmt.Sprintf(\"%d: %s\\n\", lineNum, line))\n\n\t\tlineNum++\n\t}\n\n\treturn b.String()\n}\n"
  },
  {
    "path": "adk/middlewares/filesystem/large_tool_result_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage filesystem\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// mockBackend is a simple in-memory backend for testing\ntype mockBackend struct {\n\tfiles map[string]string\n}\n\nfunc newMockBackend() *mockBackend {\n\treturn &mockBackend{\n\t\tfiles: make(map[string]string),\n\t}\n}\n\nfunc (m *mockBackend) Write(ctx context.Context, req *WriteRequest) error {\n\tm.files[req.FilePath] = req.Content\n\treturn nil\n}\n\nfunc (m *mockBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) {\n\tcontent, ok := m.files[req.FilePath]\n\tif !ok {\n\t\treturn nil, errors.New(\"file not found\")\n\t}\n\treturn &FileContent{Content: content}, nil\n}\n\nfunc (m *mockBackend) LsInfo(ctx context.Context, _ *LsInfoRequest) ([]FileInfo, error) {\n\treturn nil, nil\n}\n\nfunc (m *mockBackend) GrepRaw(ctx context.Context, _ *GrepRequest) ([]GrepMatch, error) {\n\treturn nil, nil\n}\n\nfunc (m *mockBackend) GlobInfo(ctx context.Context, _ *GlobInfoRequest) ([]FileInfo, error) {\n\treturn nil, nil\n}\n\nfunc (m *mockBackend) Edit(ctx context.Context, _ *EditRequest) error {\n\treturn nil\n}\n\nfunc TestToolResultOffloading_SmallResult(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 100, // Small limit for testing\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\t// Create a mock endpoint that returns a small result\n\tsmallResult := \"This is a small result\"\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: smallResult}, nil\n\t}\n\n\t// Wrap the endpoint with the middleware\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\t// Execute\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_123\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Small result should pass through unchanged\n\tif output.Result != smallResult {\n\t\tt.Errorf(\"expected result %q, got %q\", smallResult, output.Result)\n\t}\n\n\t// No file should be written\n\tif len(backend.files) != 0 {\n\t\tt.Errorf(\"expected no files to be written, got %d files\", len(backend.files))\n\t}\n}\n\nfunc TestToolResultOffloading_LargeResult(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10, // Very small limit to trigger offloading\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\t// Create a large result (more than 10 * 4 = 40 bytes)\n\tlargeResult := strings.Repeat(\"This is a long line of text that will exceed the token limit.\\n\", 10)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: largeResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_456\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Result should be replaced with a message\n\tif !strings.Contains(output.Result, \"Tool result too large\") {\n\t\tt.Errorf(\"expected result to contain 'Tool result too large', got %q\", output.Result)\n\t}\n\n\tif !strings.Contains(output.Result, \"call_456\") {\n\t\tt.Errorf(\"expected result to contain call ID 'call_456', got %q\", output.Result)\n\t}\n\n\tif !strings.Contains(output.Result, \"/large_tool_result/call_456\") {\n\t\tt.Errorf(\"expected result to contain file path, got %q\", output.Result)\n\t}\n\n\t// File should be written\n\tif len(backend.files) != 1 {\n\t\tt.Fatalf(\"expected 1 file to be written, got %d files\", len(backend.files))\n\t}\n\n\tsavedContent, ok := backend.files[\"/large_tool_result/call_456\"]\n\tif !ok {\n\t\tt.Fatalf(\"expected file at /large_tool_result/call_456, got files: %v\", backend.files)\n\t}\n\n\tif savedContent != largeResult {\n\t\tt.Errorf(\"saved content doesn't match original result\")\n\t}\n}\n\nfunc TestToolResultOffloading_CustomPathGenerator(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tcustomPath := \"/custom/path/result.txt\"\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t\tPathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) {\n\t\t\treturn customPath, nil\n\t\t},\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\tlargeResult := strings.Repeat(\"Large content \", 100)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: largeResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_789\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Check custom path is used\n\tif !strings.Contains(output.Result, customPath) {\n\t\tt.Errorf(\"expected result to contain custom path %q, got %q\", customPath, output.Result)\n\t}\n\n\t// File should be written to custom path\n\tsavedContent, ok := backend.files[customPath]\n\tif !ok {\n\t\tt.Fatalf(\"expected file at %q, got files: %v\", customPath, backend.files)\n\t}\n\n\tif savedContent != largeResult {\n\t\tt.Errorf(\"saved content doesn't match original result\")\n\t}\n}\n\nfunc TestToolResultOffloading_PathGeneratorError(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\texpectedErr := errors.New(\"path generation failed\")\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t\tPathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) {\n\t\t\treturn \"\", expectedErr\n\t\t},\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\tlargeResult := strings.Repeat(\"Large content \", 100)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: largeResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_error\",\n\t}\n\t_, err := wrappedEndpoint(ctx, input)\n\n\tif err == nil {\n\t\tt.Fatal(\"expected error, got nil\")\n\t}\n\n\tif !errors.Is(err, expectedErr) {\n\t\tt.Errorf(\"expected error %v, got %v\", expectedErr, err)\n\t}\n}\n\nfunc TestToolResultOffloading_EndpointError(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 100,\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\texpectedErr := errors.New(\"endpoint execution failed\")\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn nil, expectedErr\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_endpoint_error\",\n\t}\n\t_, err := wrappedEndpoint(ctx, input)\n\n\tif err == nil {\n\t\tt.Fatal(\"expected error, got nil\")\n\t}\n\n\tif !errors.Is(err, expectedErr) {\n\t\tt.Errorf(\"expected error %v, got %v\", expectedErr, err)\n\t}\n}\n\nfunc TestToolResultOffloading_DefaultTokenLimit(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 0, // Should default to 20000\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\t// Create a result smaller than 20000 * 4 = 80000 bytes\n\tsmallResult := strings.Repeat(\"x\", 1000)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: smallResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_default\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Should pass through unchanged\n\tif output.Result != smallResult {\n\t\tt.Errorf(\"expected result to pass through unchanged\")\n\t}\n\n\t// No file should be written\n\tif len(backend.files) != 0 {\n\t\tt.Errorf(\"expected no files to be written, got %d files\", len(backend.files))\n\t}\n}\n\nfunc TestToolResultOffloading_Stream(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\t// Create a streaming endpoint that returns large content\n\tlargeResult := strings.Repeat(\"Large streaming content \", 100)\n\tmockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\t// Split the result into chunks\n\t\tchunks := []string{largeResult[:len(largeResult)/2], largeResult[len(largeResult)/2:]}\n\t\treturn &compose.StreamToolOutput{\n\t\t\tResult: schema.StreamReaderFromArray(chunks),\n\t\t}, nil\n\t}\n\n\twrappedEndpoint := middleware.Streamable(mockStreamEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_stream\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Read the stream\n\tvar result strings.Builder\n\tfor {\n\t\tchunk, err := output.Result.Recv()\n\t\tif errors.Is(err, io.EOF) {\n\t\t\tbreak\n\t\t}\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"error reading stream: %v\", err)\n\t\t}\n\t\tresult.WriteString(chunk)\n\t}\n\n\tresultStr := result.String()\n\n\t// Result should be replaced with a message\n\tif !strings.Contains(resultStr, \"Tool result too large\") {\n\t\tt.Errorf(\"expected result to contain 'Tool result too large', got %q\", resultStr)\n\t}\n\n\tif !strings.Contains(resultStr, \"call_stream\") {\n\t\tt.Errorf(\"expected result to contain call ID 'call_stream', got %q\", resultStr)\n\t}\n\n\t// File should be written\n\tif len(backend.files) != 1 {\n\t\tt.Fatalf(\"expected 1 file to be written, got %d files\", len(backend.files))\n\t}\n\n\tsavedContent, ok := backend.files[\"/large_tool_result/call_stream\"]\n\tif !ok {\n\t\tt.Fatalf(\"expected file at /large_tool_result/call_stream, got files: %v\", backend.files)\n\t}\n\n\tif savedContent != largeResult {\n\t\tt.Errorf(\"saved content doesn't match original result\")\n\t}\n}\n\nfunc TestToolResultOffloading_StreamError(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\texpectedErr := errors.New(\"stream endpoint failed\")\n\tmockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\treturn nil, expectedErr\n\t}\n\n\twrappedEndpoint := middleware.Streamable(mockStreamEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_stream_error\",\n\t}\n\t_, err := wrappedEndpoint(ctx, input)\n\n\tif err == nil {\n\t\tt.Fatal(\"expected error, got nil\")\n\t}\n\n\tif !errors.Is(err, expectedErr) {\n\t\tt.Errorf(\"expected error %v, got %v\", expectedErr, err)\n\t}\n}\n\nfunc TestFormatToolMessage(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\tinput    string\n\t\texpected string\n\t}{\n\t\t{\n\t\t\tname:     \"single line\",\n\t\t\tinput:    \"single line\",\n\t\t\texpected: \"1: single line\\n\",\n\t\t},\n\t\t{\n\t\t\tname:     \"multiple lines\",\n\t\t\tinput:    \"line1\\nline2\\nline3\",\n\t\t\texpected: \"1: line1\\n2: line2\\n3: line3\\n\",\n\t\t},\n\t\t{\n\t\t\tname:     \"more than 10 lines\",\n\t\t\tinput:    \"1\\n2\\n3\\n4\\n5\\n6\\n7\\n8\\n9\\n10\\n11\\n12\",\n\t\t\texpected: \"1: 1\\n2: 2\\n3: 3\\n4: 4\\n5: 5\\n6: 6\\n7: 7\\n8: 8\\n9: 9\\n10: 10\\n\",\n\t\t},\n\t\t{\n\t\t\tname:     \"long line truncation\",\n\t\t\tinput:    strings.Repeat(\"a\", 1500),\n\t\t\texpected: fmt.Sprintf(\"1: %s\\n\", strings.Repeat(\"a\", 1000)),\n\t\t},\n\t\t{\n\t\t\tname:     \"unicode characters\",\n\t\t\tinput:    \"你好世界\\n测试\",\n\t\t\texpected: \"1: 你好世界\\n2: 测试\\n\",\n\t\t},\n\t\t{\n\t\t\tname:     \"long unicode line\",\n\t\t\tinput:    strings.Repeat(\"你\", 1500),\n\t\t\texpected: fmt.Sprintf(\"1: %s\\n\", strings.Repeat(\"你\", 1000)),\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := formatToolMessage(tt.input)\n\t\t\tif result != tt.expected {\n\t\t\t\tt.Errorf(\"formatToolMessage() = %q, want %q\", result, tt.expected)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestConcatString(t *testing.T) {\n\ttests := []struct {\n\t\tname        string\n\t\tchunks      []string\n\t\texpected    string\n\t\texpectError bool\n\t}{\n\t\t{\n\t\t\tname:     \"single chunk\",\n\t\t\tchunks:   []string{\"hello\"},\n\t\t\texpected: \"hello\",\n\t\t},\n\t\t{\n\t\t\tname:     \"multiple chunks\",\n\t\t\tchunks:   []string{\"hello\", \" \", \"world\"},\n\t\t\texpected: \"hello world\",\n\t\t},\n\t\t{\n\t\t\tname:     \"empty chunks\",\n\t\t\tchunks:   []string{\"\", \"\", \"\"},\n\t\t\texpected: \"\",\n\t\t},\n\t\t{\n\t\t\tname:     \"mixed chunks\",\n\t\t\tchunks:   []string{\"a\", \"\", \"b\", \"c\"},\n\t\t\texpected: \"abc\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tsr := schema.StreamReaderFromArray(tt.chunks)\n\t\t\tresult, err := concatString(sr)\n\n\t\t\tif tt.expectError {\n\t\t\t\tif err == nil {\n\t\t\t\t\tt.Error(\"expected error, got nil\")\n\t\t\t\t}\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t\t\t}\n\n\t\t\tif result != tt.expected {\n\t\t\t\tt.Errorf(\"concatString() = %q, want %q\", result, tt.expected)\n\t\t\t}\n\t\t})\n\t}\n\n\t// Test nil stream\n\tt.Run(\"nil stream\", func(t *testing.T) {\n\t\t_, err := concatString(nil)\n\t\tif err == nil {\n\t\t\tt.Error(\"expected error for nil stream, got nil\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"stream is nil\") {\n\t\t\tt.Errorf(\"expected 'stream is nil' error, got %v\", err)\n\t\t}\n\t})\n}\n\nfunc TestToolResultOffloading_BackendWriteError(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a backend that fails on write\n\tbackend := &failingBackend{\n\t\twriteErr: errors.New(\"write failed\"),\n\t}\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\tlargeResult := strings.Repeat(\"Large content \", 100)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: largeResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_write_error\",\n\t}\n\t_, err := wrappedEndpoint(ctx, input)\n\n\tif err == nil {\n\t\tt.Fatal(\"expected error, got nil\")\n\t}\n\n\tif !strings.Contains(err.Error(), \"write failed\") {\n\t\tt.Errorf(\"expected 'write failed' error, got %v\", err)\n\t}\n}\n\n// failingBackend is a mock backend that can be configured to fail\ntype failingBackend struct {\n\twriteErr error\n}\n\nfunc (f *failingBackend) Write(ctx context.Context, req *WriteRequest) error {\n\tif f.writeErr != nil {\n\t\treturn f.writeErr\n\t}\n\treturn nil\n}\n\nfunc (f *failingBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) {\n\treturn &FileContent{}, nil\n}\n\nfunc (f *failingBackend) LsInfo(ctx context.Context, _ *LsInfoRequest) ([]FileInfo, error) {\n\treturn nil, nil\n}\n\nfunc (f *failingBackend) GrepRaw(ctx context.Context, _ *GrepRequest) ([]GrepMatch, error) {\n\treturn nil, nil\n}\n\nfunc (f *failingBackend) GlobInfo(ctx context.Context, _ *GlobInfoRequest) ([]FileInfo, error) {\n\treturn nil, nil\n}\n\nfunc (f *failingBackend) Edit(ctx context.Context, _ *EditRequest) error {\n\treturn nil\n}\n"
  },
  {
    "path": "adk/middlewares/filesystem/prompt.go",
    "content": "/*\n * Copyright (c) 2025 Harrison Chase\n * Copyright (c) 2025 CloudWeGo Authors\n * SPDX-License-Identifier: MIT\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage filesystem\n\n// This file contains prompt templates and tool descriptions adapted from the DeepAgents project.\n// Original source: https://github.com/langchain-ai/deepagents\n//\n// These prompts are used under the terms of the original project's open source license.\n// When using this code in your own open source project, ensure compliance with the original license requirements.\n\nconst (\n\ttooLargeToolMessage = `Tool result too large, the result of this tool call {tool_call_id} was saved in the filesystem at this path: {file_path}\nYou can read the result from the filesystem by using the read_file tool, but make sure to only read part of the result at a time.\nYou can do this by specifying an offset and limit in the read_file tool call.\nFor example, to read the first 100 lines, you can use the read_file tool with offset=0 and limit=100.\n\nHere are the first 10 lines of the result:\n{content_sample}`\n\n\ttooLargeToolMessageChinese = `工具结果过大，此工具调用 {tool_call_id} 的结果已保存到文件系统的以下路径：{file_path}\n你可以使用 read_file 工具从文件系统读取结果，但请确保每次只读取部分结果。\n你可以通过在 read_file 工具调用中指定 offset 和 limit 来实现。\n例如，要读取前 100 行，你可以使用 read_file 工具，设置 offset=0 和 limit=100。\n\n以下是结果的前 10 行：\n{content_sample}`\n\n\tListFilesToolDesc = `Lists all files in the filesystem, filtering by directory.\n\nUsage:\n- The path parameter must be an absolute path, not a relative path\n- The ls tool will return a list of all files in the specified directory.\n- This is very useful for exploring the file system and finding the right file to read or edit.\n- You should almost ALWAYS use this tool before using the read_file or edit_file tools.`\n\n\tListFilesToolDescChinese = `列出文件系统中的所有文件，按目录过滤。\n\n使用方法：\n- path 参数必须是绝对路径，不能是相对路径\n- ls 工具将返回指定目录中所有文件的列表\n- 这对于探索文件系统和找到要读取或编辑的正确文件非常有用\n- 在使用 read_file 或 edit_file 工具之前，你几乎总是应该先使用此工具`\n\n\tReadFileToolDesc = `Reads a file from the filesystem. You can access any file directly by using this tool.\nAssume this tool is able to read all files on the machine. If the User provides a path to a file assume that path is valid. It is okay to read a file that does not exist; an error will be returned.\n\nUsage:\n- The file_path parameter must be an absolute path, not a relative path\n- By default, it reads up to 2000 lines starting from the beginning of the file\n- **IMPORTANT for large files and codebase exploration**: Use pagination with offset and limit parameters to avoid context overflow\n\t- First scan: read_file(path, limit=100) to see file structure\n\t- Read more sections: read_file(path, offset=100, limit=200) for next 200 lines\n\t- Only omit limit (read full file) when necessary for editing\n- Specify offset and limit: read_file(path, offset=0, limit=100) reads first 100 lines\n- Results are returned using cat -n format, with line numbers starting at 1\n- You have the capability to call multiple tools in a single response. It is always better to speculatively read multiple files as a batch that are potentially useful.\n- If you read a file that exists but has empty contents you will receive a system reminder warning in place of file contents.\n- You should ALWAYS make sure a file has been read before editing it.`\n\n\tReadFileToolDescChinese = `从文件系统读取文件。你可以使用此工具直接访问任何文件。\n假设此工具能够读取机器上的所有文件。如果用户提供了文件路径，假设该路径是有效的。读取不存在的文件是可以的；将返回错误。\n\n使用方法：\n- file_path 参数必须是绝对路径，不能是相对路径\n- 默认情况下，从文件开头读取最多 2000 行\n- **大文件和代码库探索的重要提示**：使用 offset 和 limit 参数进行分页，以避免上下文溢出\n\t- 首次扫描：read_file(path, limit=100) 查看文件结构\n\t- 读取更多部分：read_file(path, offset=100, limit=200) 读取接下来的 200 行\n\t- 仅在编辑必要时才省略 limit（读取完整文件）\n- 指定 offset 和 limit：read_file(path, offset=0, limit=100) 读取前 100 行\n- 结果以 cat -n 格式返回，行号从 1 开始\n- 你可以在单个响应中调用多个工具。最好同时推测性地批量读取多个可能有用的文件\n- 如果你读取的文件存在但内容为空，你将收到系统提醒警告而不是文件内容\n- 在编辑文件之前，你应该始终确保已读取该文件`\n\n\tEditFileToolDesc = `Performs exact string replacements in files.\n\nUsage:\n- You must use your 'read_file' tool at least once in the conversation before editing. This tool will error if you attempt an edit without reading the file.\n- When editing text from Read tool output, ensure you preserve the exact indentation (tabs/spaces) as it appears AFTER the line number prefix. The line number prefix format is: spaces + line number + tab. Everything after that tab is the actual file content to match. Never include any part of the line number prefix in the old_string or new_string.\n- ALWAYS prefer editing existing files. NEVER write new files unless explicitly required.\n- Only use emojis if the user explicitly requests it. Avoid adding emojis to files unless asked.\n- The edit will FAIL if 'old_string' is not unique in the file. Either provide a larger string with more surrounding context to make it unique or use 'replace_all' to change every instance of 'old_string'.\n- Use 'replace_all' for replacing and renaming strings across the file. This parameter is useful if you want to rename a variable for instance.`\n\n\tEditFileToolDescChinese = `在文件中执行精确的字符串替换。\n\n使用方法：\n- 在编辑之前，你必须在对话中至少使用一次 'read_file' 工具。如果你在未读取文件的情况下尝试编辑，此工具将报错\n- 当从 Read 工具输出编辑文本时，请确保保留行号前缀之后的确切缩进（制表符/空格）。行号前缀格式为：空格 + 行号 + 制表符。制表符之后的所有内容都是要匹配的实际文件内容。永远不要在 old_string 或 new_string 中包含行号前缀的任何部分\n- 始终优先编辑现有文件。除非明确要求，否则不要创建新文件\n- 仅在用户明确要求时使用表情符号。除非被要求，否则避免在文件中添加表情符号\n- 如果 'old_string' 在文件中不唯一，编辑将失败。要么提供包含更多上下文的更长字符串使其唯一，要么使用 'replace_all' 更改 'old_string' 的每个实例\n- 使用 'replace_all' 在整个文件中替换和重命名字符串。例如，如果你想重命名变量，此参数很有用`\n\n\tWriteFileToolDesc = `Writes a file to the local filesystem.\n\nUsage:\n- This tool will overwrite the existing file if there is one at the provided path.\n- If this is an existing file, you MUST use the Read tool first to read the file's contents. This tool will fail if you did not read the file first.\n- ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required.\n- NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User.\n- Only use emojis if the user explicitly requests it. Avoid writing emojis to files unless asked.`\n\n\tWriteFileToolDescChinese = `将文件写入本地文件系统。\n\n使用方法：\n- 如果提供的路径已存在文件，此工具将覆盖现有文件\n- 如果这是一个现有文件，你必须先使用 Read 工具读取文件内容。如果你没有先读取文件，此工具将失败\n- 始终优先编辑代码库中的现有文件。除非明确要求，否则不要创建新文件\n- 不要主动创建文档文件（*.md）或 README 文件。仅在用户明确要求时才创建文档文件\n- 仅在用户明确要求时使用表情符号。除非被要求，否则避免在文件中写入表情符号`\n\n\tGlobToolDesc = `Fast file pattern matching tool that works with any codebase size\n- Supports glob patterns like \"**/*.js\" or \"src/**/*.ts\"\n- Returns matching file paths sorted by modification time\n- Use this tool when you need to find files by name patterns\n- You can call multiple tools in a single response. It is always better to speculatively perform multiple searches in parallel if they are potentially useful.\n\nExamples:\n- '**/*.py' - Find all Python files\n- '*.txt' - Find all text files in root\n- '/subdir/**/*.md' - Find all markdown files under /subdir`\n\n\tGlobToolDescChinese = `适用于任何代码库大小的快速文件模式匹配工具\n- 支持 glob 模式，如 \"**/*.js\" 或 \"src/**/*.ts\"\n- 返回按修改时间排序的匹配文件路径\n- 当你需要按名称模式查找文件时使用此工具\n- 你可以在单个响应中调用多个工具。最好同时并行执行多个可能有用的搜索\n\n示例：\n- '**/*.py' - 查找所有 Python 文件\n- '*.txt' - 查找根目录中的所有文本文件\n- '/subdir/**/*.md' - 查找 /subdir 下的所有 markdown 文件`\n\n\tGrepToolDesc = `\nA powerful search tool built on ripgrep\n\n  Usage:\n  - ALWAYS use Grep for search tasks. NEVER invoke 'grep' or 'rg' as a Bash command. The Grep tool has been optimized for correct permissions and access.\n  - Supports full regex syntax (e.g., \"log.*Error\", \"function\\s+\\w+\")\n  - Filter files with glob parameter (e.g., \"*.js\", \"**/*.tsx\") or type parameter (e.g., \"js\", \"py\", \"rust\")\n  - Output modes: \"content\" shows matching lines, \"files_with_matches\" shows only file paths (default), \"count\" shows match counts\n  - Use Task tool for open-ended searches requiring multiple rounds\n  - Pattern syntax: Uses ripgrep (not grep) - literal braces need escaping (use 'interface\\{\\}' to find 'interface{}' in Go code)\n  - Multiline matching: By default patterns match within single lines only. For cross-line patterns like 'struct \\{[\\s\\S]*?field', use 'multiline: true'`\n\n\tGrepToolDescChinese = `\n基于 ripgrep 的强大搜索工具\n\n  使用方法：\n  - 始终使用 Grep 进行搜索任务。不要将 'grep' 或 'rg' 作为 Bash 命令调用。Grep 工具已针对正确的权限和访问进行了优化\n  - 支持完整的正则表达式语法（例如，\"log.*Error\"，\"function\\s+\\w+\"）\n  - 使用 glob 参数（例如，\"*.js\"，\"**/*.tsx\"）或 type 参数（例如，\"js\"，\"py\"，\"rust\"）过滤文件\n  - 输出模式：\"content\" 显示匹配行，\"files_with_matches\" 仅显示文件路径（默认），\"count\" 显示匹配计数\n  - 对于需要多轮的开放式搜索，使用 Task 工具\n  - 模式语法：使用 ripgrep（不是 grep）- 字面大括号需要转义（使用 'interface\\{\\}' 在 Go 代码中查找 'interface{}'）\n  - 多行匹配：默认情况下，模式仅在单行内匹配。对于跨行模式如 'struct \\{[\\s\\S]*?field'，使用 'multiline: true'`\n\n\tExecuteToolDesc = `\nExecutes a given command in the sandbox environment with proper handling and security measures.\n\nBefore executing the command, please follow these steps:\n\n1. Directory Verification:\n- If the command will create new directories or files, first use the ls tool to verify the parent directory exists and is the correct location\n- For example, before running \"mkdir foo/bar\", first use ls to check that \"foo\" exists and is the intended parent directory\n\n2. Command Execution:\n- Always quote file paths that contain spaces with double quotes (e.g., cd \"path with spaces/file.txt\")\n- Examples of proper quoting:\n- cd \"/Users/name/My Documents\" (correct)\n- cd /Users/name/My Documents (incorrect - will fail)\n- python \"/path/with spaces/script.py\" (correct)\n- python /path/with spaces/script.py (incorrect - will fail)\n- After ensuring proper quoting, execute the command\n- Capture the output of the command\n\nUsage notes:\n- The command parameter is required\n- Commands run in an isolated sandbox environment\n- Returns combined stdout/stderr output with exit code\n- If the output is very large, it may be truncated\n- VERY IMPORTANT: You MUST avoid using search commands like find and grep. Instead use the grep, glob tools to search. You MUST avoid read tools like cat, head, tail, and use read_file to read files.\n- When issuing multiple commands, use the ';' or '&&' operator to separate them. DO NOT use newlines (newlines are ok in quoted strings)\n- Use '&&' when commands depend on each other (e.g., \"mkdir dir && cd dir\")\n- Use ';' only when you need to run commands sequentially but don't care if earlier commands fail\n- Try to maintain your current working directory throughout the session by using absolute paths and avoiding usage of cd\n\nExamples:\nGood examples:\n- execute(command=\"pytest /foo/bar/tests\")\n- execute(command=\"python /path/to/script.py\")\n- execute(command=\"npm install && npm test\")\n\nBad examples (avoid these):\n- execute(command=\"cd /foo/bar && pytest tests\")  # Use absolute path instead\n- execute(command=\"cat file.txt\")  # Use read_file tool instead\n- execute(command=\"find . -name '*.py'\")  # Use glob tool instead\n- execute(command=\"grep -r 'pattern' .\")  # Use grep tool instead\n`\n\n\tExecuteToolDescChinese = `\n在沙箱环境中执行给定命令，具有适当的处理和安全措施。\n\n执行命令前，请按照以下步骤操作：\n\n1. 目录验证：\n- 如果命令将创建新目录或文件，首先使用 ls 工具验证父目录是否存在且是正确的位置\n- 例如，在运行 \"mkdir foo/bar\" 之前，首先使用 ls 检查 \"foo\" 是否存在且是预期的父目录\n\n2. 命令执行：\n- 始终用双引号引用包含空格的文件路径（例如，cd \"path with spaces/file.txt\"）\n- 正确引用的示例：\n- cd \"/Users/name/My Documents\"（正确）\n- cd /Users/name/My Documents（错误 - 将失败）\n- python \"/path/with spaces/script.py\"（正确）\n- python /path/with spaces/script.py（错误 - 将失败）\n- 确保正确引用后，执行命令\n- 捕获命令的输出\n\n使用说明：\n- command 参数是必需的\n- 命令在隔离的沙箱环境中运行\n- 返回合并的 stdout/stderr 输出和退出代码\n- 如果输出非常大，可能会被截断\n- 非常重要：你必须避免使用 find 和 grep 等搜索命令。请改用 grep、glob 工具进行搜索。你必须避免使用 cat、head、tail 等读取工具，请使用 read_file 读取文件\n- 发出多个命令时，使用 ';' 或 '&&' 运算符分隔它们。不要使用换行符（引号字符串中的换行符是可以的）\n- 当命令相互依赖时使用 '&&'（例如，\"mkdir dir && cd dir\"）\n- 仅当你需要按顺序运行命令但不关心早期命令是否失败时使用 ';'\n- 尝试通过使用绝对路径并避免使用 cd 来在整个会话中保持当前工作目录\n\n示例：\n好的示例：\n- execute(command=\"pytest /foo/bar/tests\")\n- execute(command=\"python /path/to/script.py\")\n- execute(command=\"npm install && npm test\")\n\n不好的示例（避免这些）：\n- execute(command=\"cd /foo/bar && pytest tests\")  # 改用绝对路径\n- execute(command=\"cat file.txt\")  # 改用 read_file 工具\n- execute(command=\"find . -name '*.py'\")  # 改用 glob 工具\n- execute(command=\"grep -r 'pattern' .\")  # 改用 grep 工具\n`\n)\n"
  },
  {
    "path": "adk/middlewares/patchtoolcalls/patchtoolcalls.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package patchtoolcalls provides a middleware that patches dangling tool calls in the message history.\npackage patchtoolcalls\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Config defines the configuration options for the patch tool calls middleware.\ntype Config struct {\n\t// PatchedContentGenerator is an optional custom function to generate the content\n\t// of patched tool messages. If not provided, a default message will be used.\n\t//\n\t// Parameters:\n\t//   - ctx: the context for the operation\n\t//   - toolName: the name of the tool that was called\n\t//   - toolCallID: the id of the tool call\n\t//\n\t// Returns:\n\t//   - string: the content to use for the patched tool message\n\t//   - error: any error that occurred during generation\n\tPatchedContentGenerator func(ctx context.Context, toolName, toolCallID string) (string, error)\n}\n\n// New creates a new patch tool calls middleware with the given configuration.\n//\n// The middleware scans the message history before each model invocation and inserts\n// placeholder tool messages for any tool calls that don't have corresponding responses.\nfunc New(ctx context.Context, cfg *Config) (adk.ChatModelAgentMiddleware, error) {\n\tif cfg == nil {\n\t\tcfg = &Config{}\n\t}\n\treturn &middleware{\n\t\tBaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},\n\t\tgen:                          cfg.PatchedContentGenerator,\n\t}, nil\n}\n\ntype middleware struct {\n\t*adk.BaseChatModelAgentMiddleware\n\tgen func(ctx context.Context, toolName, toolCallID string) (string, error)\n}\n\nfunc (m *middleware) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState,\n\tmc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {\n\n\tif len(state.Messages) == 0 {\n\t\treturn ctx, state, nil\n\t}\n\n\tpatched := make([]adk.Message, 0, len(state.Messages))\n\n\tfor i, msg := range state.Messages {\n\t\tpatched = append(patched, msg)\n\n\t\tif msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 {\n\t\t\tcontinue\n\t\t}\n\n\t\tfor _, tc := range msg.ToolCalls {\n\t\t\tif hasCorrespondingToolMessage(state.Messages[i+1:], tc.ID) {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\ttoolMsg, err := m.createPatchedToolMessage(ctx, tc)\n\t\t\tif err != nil {\n\t\t\t\treturn ctx, nil, err\n\t\t\t}\n\t\t\tpatched = append(patched, toolMsg)\n\t\t}\n\t}\n\n\tnState := *state\n\tnState.Messages = patched\n\treturn ctx, &nState, nil\n}\n\nfunc hasCorrespondingToolMessage(messages []adk.Message, toolCallID string) bool {\n\tfor _, msg := range messages {\n\t\tif msg.Role == schema.Tool && msg.ToolCallID == toolCallID {\n\t\t\treturn true\n\t\t}\n\t}\n\treturn false\n}\n\nfunc (m *middleware) createPatchedToolMessage(ctx context.Context, tc schema.ToolCall) (adk.Message, error) {\n\tif m.gen != nil {\n\t\tcontent, err := m.gen(ctx, tc.Function.Name, tc.ID)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn schema.ToolMessage(content, tc.ID, schema.WithToolName(tc.Function.Name)), nil\n\t}\n\ttpl := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: defaultPatchedToolMessageTemplate,\n\t\tChinese: defaultPatchedToolMessageTemplateChinese,\n\t})\n\n\treturn schema.ToolMessage(fmt.Sprintf(tpl, tc.Function.Name, tc.ID), tc.ID, schema.WithToolName(tc.Function.Name)), nil\n}\n\nconst (\n\tdefaultPatchedToolMessageTemplate        = \"Tool call %s with id %s was cancelled - another message came in before it could be completed.\"\n\tdefaultPatchedToolMessageTemplateChinese = \"工具调用 %s（ID 为 %s）已被取消——在其完成之前收到了另一条消息。\"\n)\n"
  },
  {
    "path": "adk/middlewares/patchtoolcalls/patchtoolcalls_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage patchtoolcalls\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestPatchToolCalls(t *testing.T) {\n\tctx := context.Background()\n\tm, err := New(ctx, nil)\n\tassert.NoError(t, err)\n\n\t// empty messages\n\tstate := &adk.ChatModelAgentState{\n\t\tMessages: nil,\n\t}\n\t_, newState, err := m.BeforeModelRewriteState(ctx, state, nil)\n\tassert.NoError(t, err)\n\tassert.Len(t, newState.Messages, 0)\n\n\tstate = &adk.ChatModelAgentState{\n\t\tMessages: []adk.Message{\n\t\t\tschema.UserMessage(\"hello\"),\n\t\t\tschema.AssistantMessage(\"hi there\", nil),\n\t\t},\n\t}\n\t_, newState, err = m.BeforeModelRewriteState(ctx, state, nil)\n\tassert.NoError(t, err)\n\tassert.Len(t, newState.Messages, 2)\n\n\tstate = &adk.ChatModelAgentState{\n\t\tMessages: []adk.Message{\n\t\t\tschema.UserMessage(\"hello\"),\n\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t{ID: \"call_1\", Function: schema.FunctionCall{Name: \"tool_a\"}},\n\t\t\t\t{ID: \"call_2\", Function: schema.FunctionCall{Name: \"tool_b\"}},\n\t\t\t}),\n\t\t\tschema.ToolMessage(\"result_a\", \"call_1\", schema.WithToolName(\"tool_a\")),\n\t\t},\n\t}\n\t_, newState, err = m.BeforeModelRewriteState(ctx, state, nil)\n\tassert.NoError(t, err)\n\tpatchedMsg := newState.Messages[2]\n\tassert.Equal(t, schema.Tool, patchedMsg.Role)\n\tassert.Equal(t, \"call_2\", patchedMsg.ToolCallID)\n\tassert.Equal(t, \"tool_b\", patchedMsg.ToolName)\n\tassert.Equal(t, fmt.Sprintf(defaultPatchedToolMessageTemplate, \"tool_b\", \"call_2\"), patchedMsg.Content)\n\n\tm, err = New(ctx, &Config{\n\t\tPatchedContentGenerator: func(ctx context.Context, toolName, toolCallID string) (string, error) {\n\t\t\treturn fmt.Sprintf(\"123 %s %s\", toolName, toolCallID), nil\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\tstate = &adk.ChatModelAgentState{\n\t\tMessages: []adk.Message{\n\t\t\tschema.UserMessage(\"hello\"),\n\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t{ID: \"call_1\", Function: schema.FunctionCall{Name: \"tool_a\"}},\n\t\t\t\t{ID: \"call_2\", Function: schema.FunctionCall{Name: \"tool_b\"}},\n\t\t\t}),\n\t\t\tschema.ToolMessage(\"result_a\", \"call_1\", schema.WithToolName(\"tool_a\")),\n\t\t},\n\t}\n\t_, newState, err = m.BeforeModelRewriteState(ctx, state, nil)\n\tassert.NoError(t, err)\n\tpatchedMsg = newState.Messages[2]\n\tassert.Equal(t, schema.Tool, patchedMsg.Role)\n\tassert.Equal(t, \"call_2\", patchedMsg.ToolCallID)\n\tassert.Equal(t, \"tool_b\", patchedMsg.ToolName)\n\tassert.Equal(t, \"123 tool_b call_2\", patchedMsg.Content)\n}\n"
  },
  {
    "path": "adk/middlewares/plantask/backend_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"sync\"\n\n\tfspkg \"github.com/cloudwego/eino/adk/filesystem\"\n)\n\ntype inMemoryBackend struct {\n\tfiles map[string]string\n\tmu    sync.RWMutex\n}\n\nfunc newInMemoryBackend() *inMemoryBackend {\n\treturn &inMemoryBackend{\n\t\tfiles: make(map[string]string),\n\t}\n}\n\nfunc (b *inMemoryBackend) LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error) {\n\tb.mu.RLock()\n\tdefer b.mu.RUnlock()\n\n\treqPath := strings.TrimSuffix(req.Path, \"/\")\n\tvar result []FileInfo\n\tfor path := range b.files {\n\t\tdir := filepath.Dir(path)\n\t\tif dir == reqPath {\n\t\t\tresult = append(result, FileInfo{Path: path})\n\t\t}\n\t}\n\treturn result, nil\n}\n\nfunc (b *inMemoryBackend) Read(ctx context.Context, req *ReadRequest) (*fspkg.FileContent, error) {\n\tb.mu.RLock()\n\tdefer b.mu.RUnlock()\n\n\tcontent, ok := b.files[req.FilePath]\n\tif !ok {\n\t\treturn nil, errors.New(\"file not found\")\n\t}\n\treturn &fspkg.FileContent{Content: content}, nil\n}\n\nfunc (b *inMemoryBackend) Write(ctx context.Context, req *WriteRequest) error {\n\tb.mu.Lock()\n\tdefer b.mu.Unlock()\n\n\tb.files[req.FilePath] = req.Content\n\treturn nil\n}\n\nfunc (b *inMemoryBackend) Delete(ctx context.Context, req *DeleteRequest) error {\n\tb.mu.Lock()\n\tdefer b.mu.Unlock()\n\n\tdelete(b.files, req.FilePath)\n\treturn nil\n}\n"
  },
  {
    "path": "adk/middlewares/plantask/plantask.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"sync\"\n\n\t\"github.com/cloudwego/eino/adk\"\n)\n\n// Config is the configuration for the tool search middleware.\ntype Config struct {\n\tBackend Backend\n\tBaseDir string\n}\n\n// New creates a new plantask middleware that provides task management tools for agents.\n// It adds TaskCreate, TaskGet, TaskUpdate, and TaskList tools to the agent's tool set,\n// allowing agents to create and manage structured task lists during coding sessions.\nfunc New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) {\n\tif config == nil {\n\t\treturn nil, fmt.Errorf(\"config is required\")\n\t}\n\tif config.Backend == nil {\n\t\treturn nil, fmt.Errorf(\"backend is required\")\n\t}\n\tif config.BaseDir == \"\" {\n\t\treturn nil, fmt.Errorf(\"baseDir is required\")\n\t}\n\n\treturn &middleware{backend: config.Backend, baseDir: config.BaseDir}, nil\n}\n\ntype middleware struct {\n\tadk.BaseChatModelAgentMiddleware\n\tbackend Backend\n\tbaseDir string\n}\n\nfunc (m *middleware) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) {\n\tif runCtx == nil {\n\t\treturn ctx, runCtx, nil\n\t}\n\n\tnRunCtx := *runCtx\n\tlock := sync.Mutex{}\n\tnRunCtx.Tools = append(nRunCtx.Tools,\n\t\tnewTaskCreateTool(m.backend, m.baseDir, &lock),\n\t\tnewTaskGetTool(m.backend, m.baseDir, &lock),\n\t\tnewTaskUpdateTool(m.backend, m.baseDir, &lock),\n\t\tnewTaskListTool(m.backend, m.baseDir, &lock),\n\t)\n\n\treturn ctx, &nRunCtx, nil\n}\n"
  },
  {
    "path": "adk/middlewares/plantask/plantask_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/components/tool\"\n)\n\nfunc TestNew(t *testing.T) {\n\tctx := context.Background()\n\n\t_, err := New(ctx, nil)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"config is required\")\n\n\t_, err = New(ctx, &Config{})\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"backend is required\")\n\n\t_, err = New(ctx, &Config{Backend: newInMemoryBackend()})\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"baseDir is required\")\n\n\tm, err := New(ctx, &Config{Backend: newInMemoryBackend(), BaseDir: \"/tmp/tasks\"})\n\tassert.NoError(t, err)\n\tassert.NotNil(t, m)\n}\n\nfunc TestMiddlewareBeforeAgent(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\n\tm, err := New(ctx, &Config{Backend: backend, BaseDir: baseDir})\n\tassert.NoError(t, err)\n\n\tmw := m.(*middleware)\n\n\tctx, runCtx, err := mw.BeforeAgent(ctx, nil)\n\tassert.NoError(t, err)\n\tassert.Nil(t, runCtx)\n\n\trunCtx = &adk.ChatModelAgentContext{\n\t\tTools: []tool.BaseTool{},\n\t}\n\tctx, newRunCtx, err := mw.BeforeAgent(ctx, runCtx)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, newRunCtx)\n\tassert.Len(t, newRunCtx.Tools, 4)\n\n\ttoolNames := make([]string, 0, 4)\n\tfor _, t := range newRunCtx.Tools {\n\t\tinfo, _ := t.Info(ctx)\n\t\ttoolNames = append(toolNames, info.Name)\n\t}\n\tassert.Contains(t, toolNames, \"TaskCreate\")\n\tassert.Contains(t, toolNames, \"TaskGet\")\n\tassert.Contains(t, toolNames, \"TaskUpdate\")\n\tassert.Contains(t, toolNames, \"TaskList\")\n}\n\nfunc TestIntegration(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\tcreateTool := newTaskCreateTool(backend, baseDir, lock)\n\tgetTool := newTaskGetTool(backend, baseDir, lock)\n\tupdateTool := newTaskUpdateTool(backend, baseDir, lock)\n\tlistTool := newTaskListTool(backend, baseDir, lock)\n\n\tresult, err := createTool.InvokableRun(ctx, `{\"subject\": \"Task 1\", \"description\": \"First task\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"Task #1\")\n\n\tresult, err = createTool.InvokableRun(ctx, `{\"subject\": \"Task 2\", \"description\": \"Second task\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"Task #2\")\n\n\t_, err = updateTool.InvokableRun(ctx, `{\"taskId\": \"2\", \"addBlockedBy\": [\"1\"]}`)\n\tassert.NoError(t, err)\n\n\tresult, err = listTool.InvokableRun(ctx, `{}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"#1 [pending] Task 1\")\n\tassert.Contains(t, result, \"#2 [pending] Task 2\")\n\tassert.Contains(t, result, \"[blocked by #1]\")\n\n\t_, err = updateTool.InvokableRun(ctx, `{\"taskId\": \"1\", \"status\": \"in_progress\"}`)\n\tassert.NoError(t, err)\n\n\tresult, err = getTool.InvokableRun(ctx, `{\"taskId\": \"1\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"Status: in_progress\")\n\n\t_, err = updateTool.InvokableRun(ctx, `{\"taskId\": \"1\", \"status\": \"completed\"}`)\n\tassert.NoError(t, err)\n\n\tresult, err = listTool.InvokableRun(ctx, `{}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"#1 [completed] Task 1\")\n}\n"
  },
  {
    "path": "adk/middlewares/plantask/task.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"regexp\"\n\n\t\"github.com/cloudwego/eino/adk/middlewares/filesystem\"\n)\n\nvar validTaskIDRegex = regexp.MustCompile(`^\\d+$`)\n\nconst highWatermarkFileName = \".highwatermark\"\n\ntype task struct {\n\tID          string         `json:\"id\"`\n\tSubject     string         `json:\"subject\"`\n\tDescription string         `json:\"description\"`\n\tStatus      string         `json:\"status\"`\n\tBlocks      []string       `json:\"blocks\"`\n\tBlockedBy   []string       `json:\"blockedBy\"`\n\tActiveForm  string         `json:\"activeForm,omitempty\"`\n\tOwner       string         `json:\"owner,omitempty\"`\n\tMetadata    map[string]any `json:\"metadata,omitempty\"`\n}\n\ntype taskOut struct {\n\tResult string `json:\"result\"`\n}\n\nconst (\n\ttaskStatusPending    = \"pending\"\n\ttaskStatusInProgress = \"in_progress\"\n\ttaskStatusCompleted  = \"completed\"\n\ttaskStatusDeleted    = \"deleted\"\n)\n\ntype FileInfo = filesystem.FileInfo\ntype LsInfoRequest = filesystem.LsInfoRequest\ntype ReadRequest = filesystem.ReadRequest\ntype WriteRequest = filesystem.WriteRequest\n\ntype DeleteRequest struct {\n\tFilePath string\n}\n\n// Backend defines the storage interface for task persistence.\n// Implementations can use local filesystem, remote storage, or any other storage backend.\ntype Backend interface {\n\t// LsInfo lists file information in the specified directory.\n\tLsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error)\n\t// Read reads the content of a file.\n\tRead(ctx context.Context, req *ReadRequest) (*filesystem.FileContent, error)\n\t// Write writes content to a file, creating it if it doesn't exist.\n\tWrite(ctx context.Context, req *WriteRequest) error\n\t// Delete removes a file from storage.\n\tDelete(ctx context.Context, req *DeleteRequest) error\n}\n\nfunc isValidTaskID(taskID string) bool {\n\treturn validTaskIDRegex.MatchString(taskID)\n}\n\nfunc appendUnique(slice []string, items ...string) []string {\n\tseen := make(map[string]struct{}, len(slice))\n\tfor _, s := range slice {\n\t\tseen[s] = struct{}{}\n\t}\n\tfor _, item := range items {\n\t\tif _, exists := seen[item]; !exists {\n\t\t\tslice = append(slice, item)\n\t\t\tseen[item] = struct{}{}\n\t\t}\n\t}\n\treturn slice\n}\n\nfunc hasCyclicDependency(taskMap map[string]*task, blockerID, blockedID string) bool {\n\tif blockerID == blockedID {\n\t\treturn true\n\t}\n\n\tvisited := make(map[string]bool)\n\treturn canReach(taskMap, blockedID, blockerID, visited)\n}\n\nfunc canReach(taskMap map[string]*task, fromID, toID string, visited map[string]bool) bool {\n\tif fromID == toID {\n\t\treturn true\n\t}\n\tif visited[fromID] {\n\t\treturn false\n\t}\n\tvisited[fromID] = true\n\n\tfromTask, exists := taskMap[fromID]\n\tif !exists {\n\t\treturn false\n\t}\n\n\tfor _, blockedID := range fromTask.Blocks {\n\t\tif canReach(taskMap, blockedID, toID, visited) {\n\t\t\treturn true\n\t\t}\n\t}\n\n\treturn false\n}\n"
  },
  {
    "path": "adk/middlewares/plantask/task_create.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"sync\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc newTaskCreateTool(backend Backend, baseDir string, lock *sync.Mutex) *taskCreateTool {\n\treturn &taskCreateTool{Backend: backend, BaseDir: baseDir, lock: lock}\n}\n\ntype taskCreateTool struct {\n\tBackend Backend\n\tBaseDir string\n\tlock    *sync.Mutex\n}\n\ntype taskCreateArgs struct {\n\tSubject     string         `json:\"subject\"`\n\tDescription string         `json:\"description\"`\n\tActiveForm  string         `json:\"activeForm,omitempty\"`\n\tMetadata    map[string]any `json:\"metadata,omitempty\"`\n}\n\nfunc (t *taskCreateTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\tdesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: taskCreateToolDesc,\n\t\tChinese: taskCreateToolDescChinese,\n\t})\n\n\treturn &schema.ToolInfo{\n\t\tName: TaskCreateToolName,\n\t\tDesc: desc,\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"subject\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"A brief title for the task\",\n\t\t\t\tRequired: true,\n\t\t\t},\n\t\t\t\"description\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"A detailed description of what needs to be done\",\n\t\t\t\tRequired: true,\n\t\t\t},\n\t\t\t\"activeForm\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"Present continuous form shown in spinner when in_progress (e.g., \\\"Running tests\\\")\",\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t\t\"metadata\": {\n\t\t\t\tType: schema.Object,\n\t\t\t\tDesc: \"Arbitrary metadata to attach to the task\",\n\t\t\t\tSubParams: map[string]*schema.ParameterInfo{\n\t\t\t\t\t\"propertyNames\": {\n\t\t\t\t\t\tType: schema.String,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t}),\n\t}, nil\n}\n\nfunc (t *taskCreateTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tt.lock.Lock()\n\tdefer t.lock.Unlock()\n\n\tparams := &taskCreateArgs{}\n\terr := sonic.UnmarshalString(argumentsInJSON, params)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tfiles, err := t.Backend.LsInfo(ctx, &LsInfoRequest{\n\t\tPath: t.BaseDir,\n\t})\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s list files in %s failed, err: %w\", TaskCreateToolName, t.BaseDir, err)\n\t}\n\n\thighwatermark := int64(0)\n\tfor _, file := range files {\n\t\tfileName := filepath.Base(file.Path)\n\t\tif fileName == highWatermarkFileName {\n\t\t\tcontent, readErr := t.Backend.Read(ctx, &ReadRequest{\n\t\t\t\tFilePath: file.Path,\n\t\t\t})\n\t\t\tif readErr != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"%s read highwatermark file %s failed, err: %w\", TaskCreateToolName, file.Path, readErr)\n\t\t\t}\n\t\t\tif content.Content != \"\" {\n\t\t\t\tvar val int64\n\t\t\t\tif _, scanErr := fmt.Sscanf(content.Content, \"%d\", &val); scanErr == nil {\n\t\t\t\t\thighwatermark = val\n\t\t\t\t}\n\t\t\t}\n\t\t\tbreak\n\t\t}\n\t}\n\n\ttaskID := highwatermark + 1\n\ttaskFileName := fmt.Sprintf(\"%d.json\", taskID)\n\n\tfor _, file := range files {\n\t\tfileName := filepath.Base(file.Path)\n\t\tif fileName == taskFileName {\n\t\t\treturn \"\", fmt.Errorf(\"Task #%d already exists\", taskID)\n\t\t}\n\t}\n\n\tnewTask := &task{\n\t\tID:          fmt.Sprintf(\"%d\", taskID),\n\t\tSubject:     params.Subject,\n\t\tDescription: params.Description,\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t\tActiveForm:  params.ActiveForm,\n\t\tMetadata:    params.Metadata,\n\t}\n\n\ttaskData, err := sonic.MarshalString(newTask)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s marshal task #%d failed, err: %w\", TaskCreateToolName, taskID, err)\n\t}\n\n\t//  Write highwatermark file first\n\thighwatermarkPath := filepath.Join(t.BaseDir, highWatermarkFileName)\n\terr = t.Backend.Write(ctx, &WriteRequest{\n\t\tFilePath: highwatermarkPath,\n\t\tContent:  fmt.Sprintf(\"%d\", taskID),\n\t})\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s update highwatermark file %s failed, err: %w\", TaskCreateToolName, highwatermarkPath, err)\n\t}\n\n\ttaskFilePath := filepath.Join(t.BaseDir, taskFileName)\n\terr = t.Backend.Write(ctx, &WriteRequest{\n\t\tFilePath: taskFilePath,\n\t\tContent:  taskData,\n\t})\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s create Task #%d failed, err: %w\", TaskCreateToolName, taskID, err)\n\t}\n\n\tresp := &taskOut{\n\t\tResult: fmt.Sprintf(\"Task #%d created successfully: %s\", taskID, params.Subject),\n\t}\n\n\tjsonResp, err := sonic.MarshalString(resp)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s marshal taskOut failed, err: %w\", TaskCreateToolName, err)\n\t}\n\n\treturn jsonResp, nil\n}\n\nconst TaskCreateToolName = \"TaskCreate\"\nconst taskCreateToolDesc = `Use this tool to create a structured task list for your current coding session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user.\nIt also helps the user understand the progress of the task and overall progress of their requests.\n\n## When to Use This Tool\n\nUse this tool proactively in these scenarios:\n\n- Complex multi-step tasks - When a task requires 3 or more distinct steps or actions\n- Non-trivial and complex tasks - Tasks that require careful planning or multiple operations\n- Plan mode - When using plan mode, create a task list to track the work\n- User explicitly requests todo list - When the user directly asks you to use the todo list\n- User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated)\n- After receiving new instructions - Immediately capture user requirements as tasks\n- When you start working on a task - Mark it as in_progress BEFORE beginning work\n- After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation\n\n## When NOT to Use This Tool\n\nSkip using this tool when:\n- There is only a single, straightforward task\n- The task is trivial and tracking it provides no organizational benefit\n- The task can be completed in less than 3 trivial steps\n- The task is purely conversational or informational\n\nNOTE that you should not use this tool if there is only one trivial task to do. In this case you are better off just doing the task directly.\n\n## Task Fields\n\n- **subject**: A brief, actionable title in imperative form (e.g., \"Fix authentication bug in login flow\")\n- **description**: Detailed description of what needs to be done, including context and acceptance criteria\n- **activeForm**: Present continuous form shown in spinner when task is in_progress (e.g., \"Fixing authentication bug\"). This is displayed to the user while you work on the task.\n\n**IMPORTANT**: Always provide activeForm when creating tasks. The subject should be imperative (\"Run tests\") while activeForm should be present continuous (\"Running tests\"). All tasks are created with status \"pending\".\n\n## Tips\n\n- Create tasks with clear, specific subjects that describe the outcome\n- Include enough detail in the description for another agent to understand and complete the task\n- After creating tasks, use TaskUpdate to set up dependencies (blocks/blockedBy) if needed\n- Check TaskList first to avoid creating duplicate tasks\n`\n\nconst taskCreateToolDescChinese = `使用此工具为当前编码会话创建结构化的任务列表。这有助于跟踪进度、组织复杂任务，并向用户展示工作的完整性。\n它还帮助用户了解任务的进度和请求的整体进展。\n\n## 何时使用此工具\n\n在以下场景中主动使用此工具：\n\n- 复杂的多步骤任务 - 当任务需要 3 个或更多不同的步骤或操作时\n- 非简单的复杂任务 - 需要仔细规划或多个操作的任务\n- 计划模式 - 使用计划模式时，创建任务列表来跟踪工作\n- 用户明确要求待办列表 - 当用户直接要求使用待办列表时\n- 用户提供多个任务 - 当用户提供待办事项列表时（编号或逗号分隔）\n- 收到新指令后 - 立即将用户需求记录为任务\n- 开始处理任务时 - 在开始工作之前将其标记为 in_progress\n- 完成任务后 - 将其标记为已完成，并添加实施过程中发现的任何后续任务\n\n## 何时不使用此工具\n\n在以下情况下跳过使用此工具：\n- 只有一个简单直接的任务\n- 任务很简单，跟踪它没有组织上的好处\n- 任务可以在少于 3 个简单步骤内完成\n- 任务纯粹是对话性或信息性的\n\n注意：如果只有一个简单任务要做，不应该使用此工具。在这种情况下，直接完成任务更好。\n\n## 任务字段\n\n- **subject**：简短的、可操作的标题，使用祈使句形式（例如，\"修复登录流程中的认证错误\"）\n- **description**：需要完成的工作的详细描述，包括上下文和验收标准\n- **activeForm**：任务处于 in_progress 状态时在加载动画中显示的现在进行时形式（例如，\"正在修复认证错误\"）。这会在你处理任务时显示给用户。\n\n**重要**：创建任务时始终提供 activeForm。subject 应该是祈使句（\"运行测试\"），而 activeForm 应该是现在进行时（\"正在运行测试\"）。所有任务创建时状态为 \"pending\"。\n\n## 提示\n\n- 创建具有清晰、具体主题的任务，描述预期结果\n- 在描述中包含足够的细节，以便其他代理能够理解并完成任务\n- 创建任务后，如果需要，使用 TaskUpdate 设置依赖关系（blocks/blockedBy）\n- 先检查 TaskList 以避免创建重复任务\n`\n"
  },
  {
    "path": "adk/middlewares/plantask/task_create_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"path/filepath\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestTaskCreateTool(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttool := newTaskCreateTool(backend, baseDir, lock)\n\n\tinfo, err := tool.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, TaskCreateToolName, info.Name)\n\tassert.Equal(t, taskCreateToolDesc, info.Desc)\n\n\tresult, err := tool.InvokableRun(ctx, `{\"subject\": \"Test Task\", \"description\": \"Test description\", \"activeForm\": \"Testing\"}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, `{\"result\":\"Task #1 created successfully: Test Task\"}`, result)\n\n\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tassert.NoError(t, err)\n\n\tvar taskData task\n\terr = sonic.UnmarshalString(content.Content, &taskData)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"1\", taskData.ID)\n\tassert.Equal(t, \"Test Task\", taskData.Subject)\n\tassert.Equal(t, \"Test description\", taskData.Description)\n\tassert.Equal(t, taskStatusPending, taskData.Status)\n\tassert.Equal(t, \"Testing\", taskData.ActiveForm)\n\n\thwContent, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, highWatermarkFileName)})\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"1\", hwContent.Content)\n\n\tresult, err = tool.InvokableRun(ctx, `{\"subject\": \"Second Task\", \"description\": \"Second description\"}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, `{\"result\":\"Task #2 created successfully: Second Task\"}`, result)\n\n\thwContent, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, highWatermarkFileName)})\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"2\", hwContent.Content)\n}\n\nfunc TestTaskCreateToolWithMetadata(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttool := newTaskCreateTool(backend, baseDir, lock)\n\n\tresult, err := tool.InvokableRun(ctx, `{\"subject\": \"Task with metadata\", \"description\": \"Has metadata\", \"metadata\": {\"key1\": \"value1\", \"key2\": \"value2\"}}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"Task #1 created successfully\")\n\n\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tassert.NoError(t, err)\n\n\tvar taskData task\n\terr = sonic.UnmarshalString(content.Content, &taskData)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"value1\", taskData.Metadata[\"key1\"])\n\tassert.Equal(t, \"value2\", taskData.Metadata[\"key2\"])\n}\n"
  },
  {
    "path": "adk/middlewares/plantask/task_get.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"sync\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc newTaskGetTool(backend Backend, baseDir string, lock *sync.Mutex) *taskGetTool {\n\treturn &taskGetTool{Backend: backend, BaseDir: baseDir, lock: lock}\n}\n\ntype taskGetTool struct {\n\tBackend Backend\n\tBaseDir string\n\tlock    *sync.Mutex\n}\n\nfunc (t *taskGetTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\tdesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: taskGetToolDesc,\n\t\tChinese: taskGetToolDescChinese,\n\t})\n\n\treturn &schema.ToolInfo{\n\t\tName: TaskGetToolName,\n\t\tDesc: desc,\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"taskId\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"The ID of the task to retrieve\",\n\t\t\t\tRequired: true,\n\t\t\t},\n\t\t}),\n\t}, nil\n}\n\ntype taskGetArgs struct {\n\tTaskID string `json:\"taskId\"`\n}\n\nfunc (t *taskGetTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tt.lock.Lock()\n\tdefer t.lock.Unlock()\n\n\tparams := &taskGetArgs{}\n\terr := sonic.UnmarshalString(argumentsInJSON, params)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tif !isValidTaskID(params.TaskID) {\n\t\treturn \"\", fmt.Errorf(\"%s validate task ID failed, err: invalid format: %s\", TaskGetToolName, params.TaskID)\n\t}\n\n\ttaskFileName := fmt.Sprintf(\"%s.json\", params.TaskID)\n\ttaskFilePath := filepath.Join(t.BaseDir, taskFileName)\n\n\tcontent, err := t.Backend.Read(ctx, &ReadRequest{\n\t\tFilePath: taskFilePath,\n\t})\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s get Task #%s failed, err: %w\", TaskGetToolName, params.TaskID, err)\n\t}\n\n\ttaskData := &task{}\n\terr = sonic.UnmarshalString(content.Content, taskData)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s get Task #%s failed, err: %w\", TaskGetToolName, params.TaskID, err)\n\t}\n\n\tvar result strings.Builder\n\tresult.WriteString(fmt.Sprintf(\"Task #%s: %s\\n\", taskData.ID, taskData.Subject))\n\tresult.WriteString(fmt.Sprintf(\"Status: %s\\n\", taskData.Status))\n\tresult.WriteString(fmt.Sprintf(\"Description: %s\\n\", taskData.Description))\n\n\tif len(taskData.BlockedBy) > 0 {\n\t\tblockedByIDs := make([]string, len(taskData.BlockedBy))\n\t\tfor i, id := range taskData.BlockedBy {\n\t\t\tblockedByIDs[i] = \"#\" + id\n\t\t}\n\t\tresult.WriteString(fmt.Sprintf(\"Blocked by: %s\\n\", strings.Join(blockedByIDs, \", \")))\n\t}\n\tif len(taskData.Blocks) > 0 {\n\t\tblocksIDs := make([]string, len(taskData.Blocks))\n\t\tfor i, id := range taskData.Blocks {\n\t\t\tblocksIDs[i] = \"#\" + id\n\t\t}\n\t\tresult.WriteString(fmt.Sprintf(\"Blocks: %s\\n\", strings.Join(blocksIDs, \", \")))\n\t}\n\tif taskData.Owner != \"\" {\n\t\tresult.WriteString(fmt.Sprintf(\"Owner: %s\\n\", taskData.Owner))\n\t}\n\n\tresp := &taskOut{\n\t\tResult: result.String(),\n\t}\n\n\tjsonResp, err := sonic.MarshalString(resp)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s marshal taskOut failed, err: %w\", TaskGetToolName, err)\n\t}\n\n\treturn jsonResp, nil\n}\n\nconst TaskGetToolName = \"TaskGet\"\nconst taskGetToolDesc = `Use this tool to retrieve a task by its ID from the task list.\n\n## When to Use This Tool\n\n- When you need the full description and context before starting work on a task\n- To understand task dependencies (what it blocks, what blocks it)\n- After being assigned a task, to get complete requirements\n\n## Output\n\nReturns full task details:\n- **subject**: Task title\n- **description**: Detailed requirements and context\n- **status**: 'pending', 'in_progress', or 'completed'\n- **blocks**: Tasks waiting on this one to complete\n- **blockedBy**: Tasks that must complete before this one can start\n\n## Tips\n\n- After fetching a task, verify its blockedBy list is empty before beginning work.\n- Use TaskList to see all tasks in summary form.\n`\n\nconst taskGetToolDescChinese = `使用此工具通过任务 ID 从任务列表中获取任务。\n\n## 何时使用此工具\n\n- 当你需要在开始处理任务之前获取完整的描述和上下文时\n- 了解任务依赖关系（它阻塞什么，什么阻塞它）\n- 被分配任务后，获取完整的需求\n\n## 输出\n\n返回完整的任务详情：\n- **subject**：任务标题\n- **description**：详细的需求和上下文\n- **status**：'pending'、'in_progress' 或 'completed'\n- **blocks**：等待此任务完成的任务\n- **blockedBy**：必须在此任务开始之前完成的任务\n\n## 提示\n\n- 获取任务后，在开始工作之前验证其 blockedBy 列表是否为空。\n- 使用 TaskList 查看所有任务的摘要形式。\n`\n"
  },
  {
    "path": "adk/middlewares/plantask/task_get_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"path/filepath\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestTaskGetTool(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttaskData := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Test Task\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{\"2\", \"3\"},\n\t\tBlockedBy:   []string{\"4\"},\n\t}\n\ttaskJSON, _ := sonic.MarshalString(taskData)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: taskJSON})\n\n\ttool := newTaskGetTool(backend, baseDir, lock)\n\n\tinfo, err := tool.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, TaskGetToolName, info.Name)\n\tassert.Equal(t, taskGetToolDesc, info.Desc)\n\n\tresult, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"Task #1: Test Task\")\n\tassert.Contains(t, result, \"Status: \"+taskStatusPending)\n\tassert.Contains(t, result, \"Description: Test description\")\n\tassert.Contains(t, result, \"Blocked by: #4\")\n\tassert.Contains(t, result, \"Blocks: #2, #3\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"999\"}`)\n\tassert.Error(t, err)\n}\n\nfunc TestTaskGetToolInvalidTaskID(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttool := newTaskGetTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"../../../etc/passwd\"}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"validate task ID failed\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"abc\"}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"validate task ID failed\")\n}\n"
  },
  {
    "path": "adk/middlewares/plantask/task_list.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"sort\"\n\t\"strings\"\n\t\"sync\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc newTaskListTool(backend Backend, baseDir string, lock *sync.Mutex) *taskListTool {\n\treturn &taskListTool{Backend: backend, BaseDir: baseDir, lock: lock}\n}\n\ntype taskListTool struct {\n\tBackend Backend\n\tBaseDir string\n\tlock    *sync.Mutex\n}\n\nfunc (t *taskListTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\tdesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: taskListToolDesc,\n\t\tChinese: taskListToolDescChinese,\n\t})\n\n\treturn &schema.ToolInfo{\n\t\tName:        TaskListToolName,\n\t\tDesc:        desc,\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{}),\n\t}, nil\n}\n\nfunc listTasks(ctx context.Context, backend Backend, baseDir string) ([]*task, error) {\n\tfiles, err := backend.LsInfo(ctx, &LsInfoRequest{\n\t\tPath: baseDir,\n\t})\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"%s list files in %s failed, err: %w\", TaskListToolName, baseDir, err)\n\t}\n\n\tvar tasks []*task\n\tfor _, file := range files {\n\t\tfileName := filepath.Base(file.Path)\n\t\tif !strings.HasSuffix(fileName, \".json\") {\n\t\t\tcontinue\n\t\t}\n\n\t\ttaskID := strings.TrimSuffix(fileName, \".json\")\n\t\tif !isValidTaskID(taskID) {\n\t\t\tcontinue\n\t\t}\n\n\t\tcontent, err := backend.Read(ctx, &ReadRequest{\n\t\t\tFilePath: file.Path,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"%s read task file %s failed, err: %w\", TaskListToolName, file.Path, err)\n\t\t}\n\n\t\ttaskData := &task{}\n\t\terr = sonic.UnmarshalString(content.Content, taskData)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"%s parse task file %s failed, err: %w\", TaskListToolName, file.Path, err)\n\t\t}\n\n\t\ttasks = append(tasks, taskData)\n\t}\n\n\t// sort tasks by ID\n\tsort.Slice(tasks, func(i, j int) bool {\n\t\treturn tasks[i].ID < tasks[j].ID\n\t})\n\n\treturn tasks, nil\n}\n\nfunc (t *taskListTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tt.lock.Lock()\n\tdefer t.lock.Unlock()\n\n\ttasks, err := listTasks(ctx, t.Backend, t.BaseDir)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tif len(tasks) == 0 {\n\t\tresp := &taskOut{\n\t\t\tResult: \"No tasks found.\",\n\t\t}\n\t\tjsonResp, marshalErr := sonic.MarshalString(resp)\n\t\tif marshalErr != nil {\n\t\t\treturn \"\", fmt.Errorf(\"%s marshal taskOut failed, err: %w\", TaskListToolName, marshalErr)\n\t\t}\n\t\treturn jsonResp, nil\n\t}\n\n\tvar result strings.Builder\n\tfor i, taskData := range tasks {\n\t\tif i > 0 {\n\t\t\tresult.WriteString(\"\\n\")\n\t\t}\n\t\tresult.WriteString(fmt.Sprintf(\"#%s [%s] %s\", taskData.ID, taskData.Status, taskData.Subject))\n\t\tif taskData.Owner != \"\" {\n\t\t\tresult.WriteString(fmt.Sprintf(\" [owner: %s]\", taskData.Owner))\n\t\t}\n\t\tif len(taskData.BlockedBy) > 0 {\n\t\t\tblockedByIDs := make([]string, len(taskData.BlockedBy))\n\t\t\tfor j, id := range taskData.BlockedBy {\n\t\t\t\tblockedByIDs[j] = \"#\" + id\n\t\t\t}\n\t\t\tresult.WriteString(fmt.Sprintf(\" [blocked by %s]\", strings.Join(blockedByIDs, \", \")))\n\t\t}\n\t}\n\n\tresp := &taskOut{\n\t\tResult: result.String(),\n\t}\n\n\tjsonResp, err := sonic.MarshalString(resp)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s marshal taskOut failed, err: %w\", TaskListToolName, err)\n\t}\n\n\treturn jsonResp, nil\n}\n\nconst TaskListToolName = \"TaskList\"\nconst taskListToolDesc = `Use this tool to list all tasks in the task list.\n\n## When to Use This Tool\n\n- To see what tasks are available to work on (status: 'pending', no owner, not blocked)\n- To check overall progress on the project\n- To find tasks that are blocked and need dependencies resolved\n- After completing a task, to check for newly unblocked work or claim the next available task\n- **Prefer working on tasks in ID order** (lowest ID first) when multiple tasks are available, as earlier tasks often set up context for later ones\n\n## Output\n\nReturns a summary of each task:\n- **id**: Task identifier (use with TaskGet, TaskUpdate)\n- **subject**: Brief description of the task\n- **status**: 'pending', 'in_progress', or 'completed'\n- **owner**: Agent ID if assigned, empty if available\n- **blockedBy**: List of open task IDs that must be resolved first (tasks with blockedBy cannot be claimed until dependencies resolve)\n\nUse TaskGet with a specific task ID to view full details including description and comments.\n`\n\nconst taskListToolDescChinese = `使用此工具列出任务列表中的所有任务。\n\n## 何时使用此工具\n\n- 查看可以处理的任务（状态：'pending'，无所有者，未被阻塞）\n- 检查项目的整体进度\n- 查找被阻塞且需要解决依赖关系的任务\n- 完成任务后，检查新解除阻塞的工作或认领下一个可用任务\n- **优先按 ID 顺序处理任务**（最小 ID 优先），当有多个任务可用时，因为较早的任务通常为后续任务建立上下文\n\n## 输出\n\n返回每个任务的摘要：\n- **id**：任务标识符（与 TaskGet、TaskUpdate 一起使用）\n- **subject**：任务的简要描述\n- **status**：'pending'、'in_progress' 或 'completed'\n- **owner**：如果已分配则为代理 ID，如果可用则为空\n- **blockedBy**：必须首先解决的开放任务 ID 列表（具有 blockedBy 的任务在依赖关系解决之前无法被认领）\n\n使用 TaskGet 配合特定任务 ID 查看完整详情，包括描述和评论。\n`\n"
  },
  {
    "path": "adk/middlewares/plantask/task_list_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"path/filepath\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestTaskListTool(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttool := newTaskListTool(backend, baseDir, lock)\n\n\tinfo, err := tool.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, TaskListToolName, info.Name)\n\tassert.Equal(t, taskListToolDesc, info.Desc)\n\n\tresult, err := tool.InvokableRun(ctx, `{}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, `{\"result\":\"No tasks found.\"}`, result)\n\n\ttask1 := &task{ID: \"1\", Subject: \"Task 1\", Status: taskStatusPending, BlockedBy: []string{\"2\"}}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{ID: \"2\", Subject: \"Task 2\", Status: taskStatusInProgress, Owner: \"agent1\"}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\tresult, err = tool.InvokableRun(ctx, `{}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"#1 [\"+taskStatusPending+\"] Task 1\")\n\tassert.Contains(t, result, \"[blocked by #2]\")\n\tassert.Contains(t, result, \"#2 [\"+taskStatusInProgress+\"] Task 2\")\n\tassert.Contains(t, result, \"[owner: agent1]\")\n}\n"
  },
  {
    "path": "adk/middlewares/plantask/task_update.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"sync\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc newTaskUpdateTool(backend Backend, baseDir string, lock *sync.Mutex) *taskUpdateTool {\n\treturn &taskUpdateTool{Backend: backend, BaseDir: baseDir, lock: lock}\n}\n\ntype taskUpdateTool struct {\n\tBackend Backend\n\tBaseDir string\n\tlock    *sync.Mutex\n}\n\ntype taskUpdateArgs struct {\n\tTaskID       string         `json:\"taskId\"`\n\tSubject      string         `json:\"subject,omitempty\"`\n\tDescription  string         `json:\"description,omitempty\"`\n\tActiveForm   string         `json:\"activeForm,omitempty\"`\n\tStatus       string         `json:\"status,omitempty\"`\n\tAddBlocks    []string       `json:\"addBlocks,omitempty\"`\n\tAddBlockedBy []string       `json:\"addBlockedBy,omitempty\"`\n\tOwner        string         `json:\"owner,omitempty\"`\n\tMetadata     map[string]any `json:\"metadata,omitempty\"`\n}\n\nfunc (t *taskUpdateTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\tdesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: taskUpdateToolDesc,\n\t\tChinese: taskUpdateToolDescChinese,\n\t})\n\n\treturn &schema.ToolInfo{\n\t\tName: TaskUpdateToolName,\n\t\tDesc: desc,\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"taskId\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"The ID of the task to update\",\n\t\t\t\tRequired: true,\n\t\t\t},\n\t\t\t\"subject\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"New subject for the task\",\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t\t\"description\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"New description for the task\",\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t\t\"activeForm\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"Present continuous form shown in spinner when in_progress (e.g., \\\"Running tests\\\")\",\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t\t\"status\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"New status for the task: 'pending', 'in_progress', 'completed', or 'deleted'\",\n\t\t\t\tEnum:     []string{\"pending\", \"in_progress\", \"completed\", \"deleted\"},\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t\t\"addBlocks\": {\n\t\t\t\tType:     schema.Array,\n\t\t\t\tDesc:     \"Task IDs that this task blocks\",\n\t\t\t\tElemInfo: &schema.ParameterInfo{Type: schema.String},\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t\t\"addBlockedBy\": {\n\t\t\t\tType:     schema.Array,\n\t\t\t\tDesc:     \"Task IDs that block this task\",\n\t\t\t\tElemInfo: &schema.ParameterInfo{Type: schema.String},\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t\t\"owner\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"New owner for the task\",\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t\t\"metadata\": {\n\t\t\t\tType: schema.Object,\n\t\t\t\tDesc: \"Metadata keys to merge into the task. Set a key to null to delete it.\",\n\t\t\t\tSubParams: map[string]*schema.ParameterInfo{\n\t\t\t\t\t\"propertyNames\": {\n\t\t\t\t\t\tType: schema.String,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\tRequired: false,\n\t\t\t},\n\t\t}),\n\t}, nil\n}\n\nfunc (t *taskUpdateTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tt.lock.Lock()\n\tdefer t.lock.Unlock()\n\n\tparams := &taskUpdateArgs{}\n\terr := sonic.UnmarshalString(argumentsInJSON, params)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tif !isValidTaskID(params.TaskID) {\n\t\treturn \"\", fmt.Errorf(\"%s validate task ID failed, err: invalid format: %s\", TaskUpdateToolName, params.TaskID)\n\t}\n\n\ttaskFileName := fmt.Sprintf(\"%s.json\", params.TaskID)\n\ttaskFilePath := filepath.Join(t.BaseDir, taskFileName)\n\n\tif params.Status == taskStatusDeleted {\n\t\tif removeErr := t.removeTaskFromDependencies(ctx, params.TaskID); removeErr != nil {\n\t\t\treturn \"\", fmt.Errorf(\"%s remove Task #%s from dependencies failed, err: %w\", TaskUpdateToolName, params.TaskID, removeErr)\n\t\t}\n\n\t\terr = t.Backend.Delete(ctx, &DeleteRequest{\n\t\t\tFilePath: taskFilePath,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"%s delete Task #%s failed, err: %w\", TaskUpdateToolName, params.TaskID, err)\n\t\t}\n\n\t\tresp := &taskOut{\n\t\t\tResult: fmt.Sprintf(\"Updated task #%s deleted\", params.TaskID),\n\t\t}\n\t\tjsonResp, marshalErr := sonic.MarshalString(resp)\n\t\tif marshalErr != nil {\n\t\t\treturn \"\", fmt.Errorf(\"%s marshal taskOut failed, err: %w\", TaskUpdateToolName, marshalErr)\n\t\t}\n\t\treturn jsonResp, nil\n\t}\n\n\tcontent, err := t.Backend.Read(ctx, &ReadRequest{\n\t\tFilePath: taskFilePath,\n\t})\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s read Task #%s failed, err: %w\", TaskUpdateToolName, params.TaskID, err)\n\t}\n\n\ttaskData := &task{}\n\terr = sonic.UnmarshalString(content.Content, taskData)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s parse Task #%s failed, err: %w\", TaskUpdateToolName, params.TaskID, err)\n\t}\n\n\tvar updatedFields []string\n\n\tif params.Subject != \"\" {\n\t\ttaskData.Subject = params.Subject\n\t\tupdatedFields = append(updatedFields, \"subject\")\n\t}\n\tif params.Description != \"\" {\n\t\ttaskData.Description = params.Description\n\t\tupdatedFields = append(updatedFields, \"description\")\n\t}\n\tif params.ActiveForm != \"\" {\n\t\ttaskData.ActiveForm = params.ActiveForm\n\t\tupdatedFields = append(updatedFields, \"activeForm\")\n\t}\n\tif params.Status != \"\" {\n\t\ttaskData.Status = params.Status\n\t\tupdatedFields = append(updatedFields, \"status\")\n\t}\n\tif len(params.AddBlocks) > 0 || len(params.AddBlockedBy) > 0 {\n\t\ttasks, listErr := listTasks(ctx, t.Backend, t.BaseDir)\n\t\tif listErr != nil {\n\t\t\treturn \"\", fmt.Errorf(\"%s list tasks failed, err: %w\", TaskUpdateToolName, listErr)\n\t\t}\n\t\ttaskMap := make(map[string]*task, len(tasks))\n\t\tfor _, tk := range tasks {\n\t\t\ttaskMap[tk.ID] = tk\n\t\t}\n\n\t\tif len(params.AddBlocks) > 0 {\n\t\t\tfor _, blockedTaskID := range params.AddBlocks {\n\t\t\t\tif !isValidTaskID(blockedTaskID) {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"%s validate blocked task ID failed, err: invalid format: %s\", TaskUpdateToolName, blockedTaskID)\n\t\t\t\t}\n\t\t\t\tif hasCyclicDependency(taskMap, params.TaskID, blockedTaskID) {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"%s adding Task #%s to blocks of Task #%s would create a cyclic dependency\", TaskUpdateToolName, blockedTaskID, params.TaskID)\n\t\t\t\t}\n\t\t\t}\n\t\t\tfor _, blockedTaskID := range params.AddBlocks {\n\t\t\t\tif addErr := t.addBlockedByToTask(ctx, blockedTaskID, params.TaskID); addErr != nil {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"%s update Task #%s blocks failed, err: %w\", TaskUpdateToolName, params.TaskID, addErr)\n\t\t\t\t}\n\t\t\t}\n\t\t\ttaskData.Blocks = appendUnique(taskData.Blocks, params.AddBlocks...)\n\t\t\tupdatedFields = append(updatedFields, \"blocks\")\n\t\t}\n\t\tif len(params.AddBlockedBy) > 0 {\n\t\t\tfor _, blockingTaskID := range params.AddBlockedBy {\n\t\t\t\tif !isValidTaskID(blockingTaskID) {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"%s validate blocking task ID failed, err: invalid format: %s\", TaskUpdateToolName, blockingTaskID)\n\t\t\t\t}\n\t\t\t\tif hasCyclicDependency(taskMap, blockingTaskID, params.TaskID) {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"%s adding Task #%s to blockedBy of Task #%s would create a cyclic dependency\", TaskUpdateToolName, blockingTaskID, params.TaskID)\n\t\t\t\t}\n\t\t\t}\n\t\t\tfor _, blockingTaskID := range params.AddBlockedBy {\n\t\t\t\tif addErr := t.addBlocksToTask(ctx, blockingTaskID, params.TaskID); addErr != nil {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"%s update Task #%s blockedBy failed, err: %w\", TaskUpdateToolName, params.TaskID, addErr)\n\t\t\t\t}\n\t\t\t}\n\t\t\ttaskData.BlockedBy = appendUnique(taskData.BlockedBy, params.AddBlockedBy...)\n\t\t\tupdatedFields = append(updatedFields, \"blockedBy\")\n\t\t}\n\t}\n\tif params.Owner != \"\" {\n\t\ttaskData.Owner = params.Owner\n\t\tupdatedFields = append(updatedFields, \"owner\")\n\t}\n\tif params.Metadata != nil {\n\t\tif taskData.Metadata == nil {\n\t\t\ttaskData.Metadata = make(map[string]any)\n\t\t}\n\t\tfor k, v := range params.Metadata {\n\t\t\tif v == nil {\n\t\t\t\tdelete(taskData.Metadata, k)\n\t\t\t} else {\n\t\t\t\ttaskData.Metadata[k] = v\n\t\t\t}\n\t\t}\n\t\tupdatedFields = append(updatedFields, \"metadata\")\n\t}\n\n\tupdatedContent, err := sonic.MarshalString(taskData)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s marshal Task #%s failed, err: %w\", TaskUpdateToolName, params.TaskID, err)\n\t}\n\n\terr = t.Backend.Write(ctx, &WriteRequest{\n\t\tFilePath: taskFilePath,\n\t\tContent:  updatedContent,\n\t})\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s write Task #%s failed, err: %w\", TaskUpdateToolName, params.TaskID, err)\n\t}\n\n\tif params.Status == taskStatusCompleted {\n\t\tif checkErr := t.checkIfNeedDeleteAllTasks(ctx); checkErr != nil {\n\t\t\treturn \"\", fmt.Errorf(\"%s check and delete all tasks failed, err: %w\", TaskUpdateToolName, checkErr)\n\t\t}\n\t}\n\n\tresp := &taskOut{\n\t\tResult: fmt.Sprintf(\"Updated task #%s %s\", params.TaskID, strings.Join(updatedFields, \", \")),\n\t}\n\n\tjsonResp, err := sonic.MarshalString(resp)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s marshal taskOut failed, err: %w\", TaskUpdateToolName, err)\n\t}\n\n\treturn jsonResp, nil\n}\n\nfunc (t *taskUpdateTool) removeTaskFromDependencies(ctx context.Context, deletedTaskID string) error {\n\ttasks, err := listTasks(ctx, t.Backend, t.BaseDir)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tfor _, taskData := range tasks {\n\t\tif taskData.ID == deletedTaskID {\n\t\t\tcontinue\n\t\t}\n\n\t\tmodified := false\n\t\tnewBlocks := make([]string, 0, len(taskData.Blocks))\n\t\tfor _, id := range taskData.Blocks {\n\t\t\tif id != deletedTaskID {\n\t\t\t\tnewBlocks = append(newBlocks, id)\n\t\t\t} else {\n\t\t\t\tmodified = true\n\t\t\t}\n\t\t}\n\n\t\tnewBlockedBy := make([]string, 0, len(taskData.BlockedBy))\n\t\tfor _, id := range taskData.BlockedBy {\n\t\t\tif id != deletedTaskID {\n\t\t\t\tnewBlockedBy = append(newBlockedBy, id)\n\t\t\t} else {\n\t\t\t\tmodified = true\n\t\t\t}\n\t\t}\n\n\t\tif modified {\n\t\t\ttaskData.Blocks = newBlocks\n\t\t\ttaskData.BlockedBy = newBlockedBy\n\n\t\t\tupdatedContent, err := sonic.MarshalString(taskData)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"failed to marshal task #%s: %w\", taskData.ID, err)\n\t\t\t}\n\n\t\t\ttaskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf(\"%s.json\", taskData.ID))\n\t\t\tif err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil {\n\t\t\t\treturn fmt.Errorf(\"failed to write task #%s: %w\", taskData.ID, err)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn nil\n}\n\nfunc (t *taskUpdateTool) addBlockedByToTask(ctx context.Context, targetTaskID, blockerTaskID string) error {\n\ttaskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf(\"%s.json\", targetTaskID))\n\n\tcontent, err := t.Backend.Read(ctx, &ReadRequest{FilePath: taskFilePath})\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to read task #%s for updating blockedBy: %w\", targetTaskID, err)\n\t}\n\n\ttargetTask := &task{}\n\tif unmarshalErr := sonic.UnmarshalString(content.Content, targetTask); unmarshalErr != nil {\n\t\treturn fmt.Errorf(\"failed to parse task #%s: %w\", targetTaskID, unmarshalErr)\n\t}\n\n\ttargetTask.BlockedBy = appendUnique(targetTask.BlockedBy, blockerTaskID)\n\n\tupdatedContent, err := sonic.MarshalString(targetTask)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to marshal task #%s: %w\", targetTaskID, err)\n\t}\n\n\tif err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil {\n\t\treturn fmt.Errorf(\"failed to write task #%s: %w\", targetTaskID, err)\n\t}\n\n\treturn nil\n}\n\nfunc (t *taskUpdateTool) addBlocksToTask(ctx context.Context, targetTaskID, blockedTaskID string) error {\n\ttaskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf(\"%s.json\", targetTaskID))\n\n\tcontent, err := t.Backend.Read(ctx, &ReadRequest{FilePath: taskFilePath})\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to read task #%s for updating blocks: %w\", targetTaskID, err)\n\t}\n\n\ttargetTask := &task{}\n\tif unmarshalErr := sonic.UnmarshalString(content.Content, targetTask); unmarshalErr != nil {\n\t\treturn fmt.Errorf(\"failed to parse task #%s: %w\", targetTaskID, unmarshalErr)\n\t}\n\n\ttargetTask.Blocks = appendUnique(targetTask.Blocks, blockedTaskID)\n\n\tupdatedContent, err := sonic.MarshalString(targetTask)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to marshal task #%s: %w\", targetTaskID, err)\n\t}\n\n\tif err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil {\n\t\treturn fmt.Errorf(\"failed to write task #%s: %w\", targetTaskID, err)\n\t}\n\n\treturn nil\n}\n\n// checkIfNeedDeleteAllTasks checks if all tasks are completed, if so, it deletes all tasks\nfunc (t *taskUpdateTool) checkIfNeedDeleteAllTasks(ctx context.Context) error {\n\ttasks, err := listTasks(ctx, t.Backend, t.BaseDir)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tfor _, task := range tasks {\n\t\tif task.Status != taskStatusCompleted {\n\t\t\treturn nil\n\t\t}\n\t}\n\n\tfor _, task := range tasks {\n\t\terr := t.Backend.Delete(ctx, &DeleteRequest{\n\t\t\tFilePath: filepath.Join(t.BaseDir, task.ID+\".json\"),\n\t\t})\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\treturn nil\n}\n\nconst TaskUpdateToolName = \"TaskUpdate\"\nconst taskUpdateToolDesc = `Use this tool to update a task in the task list.\n\n## When to Use This Tool\n\n**Mark tasks as resolved:**\n- When you have completed the work described in a task\n- When a task is no longer needed or has been superseded\n- IMPORTANT: Always mark your assigned tasks as resolved when you finish them\n- After resolving, call TaskList to find your next task\n\n- ONLY mark a task as completed when you have FULLY accomplished it\n- If you encounter errors, blockers, or cannot finish, keep the task as in_progress\n- When blocked, create a new task describing what needs to be resolved\n- Never mark a task as completed if:\n  - Tests are failing\n  - Implementation is partial\n  - You encountered unresolved errors\n  - You couldn't find necessary files or dependencies\n\n**Delete tasks:**\n- When a task is no longer relevant or was created in error\n- Setting status to ` + \"`deleted`\" + ` permanently removes the task\n\n**Update task details:**\n- When requirements change or become clearer\n- When establishing dependencies between tasks\n\n## Fields You Can Update\n\n- **status**: The task status (see Status Workflow below)\n- **subject**: Change the task title (imperative form, e.g., \"Run tests\")\n- **description**: Change the task description\n- **activeForm**: Present continuous form shown in spinner when in_progress (e.g., \"Running tests\")\n- **owner**: Change the task owner (agent name)\n- **metadata**: Merge metadata keys into the task (set a key to null to delete it)\n- **addBlocks**: Mark tasks that cannot start until this one completes\n- **addBlockedBy**: Mark tasks that must complete before this one can start\n\n## Status Workflow\n\nStatus progresses: ` + \"`pending`\" + ` → ` + \"`in_progress`\" + ` → ` + \"`completed`\" + `\n\nUse ` + \"`deleted`\" + ` to permanently remove a task.\n\n## Staleness\n\nMake sure to read a task's latest state using ` + \"`TaskGet`\" + ` before updating it.\n\n## Examples\n\nMark task as in progress when starting work:\n` + \"```json\" + `\n{\"taskId\": \"1\", \"status\": \"in_progress\"}\n` + \"```\" + `\n\nMark task as completed after finishing work:\n` + \"```json\" + `\n{\"taskId\": \"1\", \"status\": \"completed\"}\n` + \"```\" + `\n\nDelete a task:\n` + \"```json\" + `\n{\"taskId\": \"1\", \"status\": \"deleted\"}\n` + \"```\" + `\n\nClaim a task by setting owner:\n` + \"```json\" + `\n{\"taskId\": \"1\", \"owner\": \"my-name\"}\n` + \"```\" + `\n\nSet up task dependencies:\n` + \"```json\" + `\n{\"taskId\": \"2\", \"addBlockedBy\": [\"1\"]}\n` + \"```\" + `\n`\n\nconst taskUpdateToolDescChinese = `使用此工具更新任务列表中的任务。\n\n## 何时使用此工具\n\n**将任务标记为已完成：**\n- 当你完成了任务中描述的工作时\n- 当任务不再需要或已被取代时\n- 重要：完成分配给你的任务后，务必将其标记为已完成\n- 完成后，调用 TaskList 查找下一个任务\n\n- 只有在完全完成任务时才将其标记为已完成\n- 如果遇到错误、阻塞或无法完成，请保持任务为 in_progress 状态\n- 当被阻塞时，创建一个新任务描述需要解决的问题\n- 在以下情况下不要将任务标记为已完成：\n  - 测试失败\n  - 实现不完整\n  - 遇到未解决的错误\n  - 找不到必要的文件或依赖项\n\n**删除任务：**\n- 当任务不再相关或创建错误时\n- 将状态设置为 ` + \"`deleted`\" + ` 会永久删除任务\n\n**更新任务详情：**\n- 当需求变更或变得更清晰时\n- 当建立任务之间的依赖关系时\n\n## 可更新的字段\n\n- **status**：任务状态（参见下方状态流程）\n- **subject**：更改任务标题（使用祈使句形式，例如\"运行测试\"）\n- **description**：更改任务描述\n- **activeForm**：in_progress 状态时在加载动画中显示的现在进行时形式（例如\"正在运行测试\"）\n- **owner**：更改任务所有者（代理名称）\n- **metadata**：将元数据键合并到任务中（将键设置为 null 可删除它）\n- **addBlocks**：标记在此任务完成之前无法开始的任务\n- **addBlockedBy**：标记必须在此任务开始之前完成的任务\n\n## 状态流程\n\n状态进展：` + \"`pending`\" + ` → ` + \"`in_progress`\" + ` → ` + \"`completed`\" + `\n\n使用 ` + \"`deleted`\" + ` 永久删除任务。\n\n## 过期性\n\n更新任务前，请确保使用 ` + \"`TaskGet`\" + ` 读取任务的最新状态。\n\n## 示例\n\n开始工作时将任务标记为进行中：\n` + \"```json\" + `\n{\"taskId\": \"1\", \"status\": \"in_progress\"}\n` + \"```\" + `\n\n完成工作后将任务标记为已完成：\n` + \"```json\" + `\n{\"taskId\": \"1\", \"status\": \"completed\"}\n` + \"```\" + `\n\n删除任务：\n` + \"```json\" + `\n{\"taskId\": \"1\", \"status\": \"deleted\"}\n` + \"```\" + `\n\n通过设置 owner 认领任务：\n` + \"```json\" + `\n{\"taskId\": \"1\", \"owner\": \"my-name\"}\n` + \"```\" + `\n\n设置任务依赖关系：\n` + \"```json\" + `\n{\"taskId\": \"2\", \"addBlockedBy\": [\"1\"]}\n` + \"```\" + `\n`\n"
  },
  {
    "path": "adk/middlewares/plantask/task_update_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage plantask\n\nimport (\n\t\"context\"\n\t\"path/filepath\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestTaskUpdateTool(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttaskData := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Original Subject\",\n\t\tDescription: \"Original description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttaskJSON, _ := sonic.MarshalString(taskData)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: taskJSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\tinfo, err := tool.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, TaskUpdateToolName, info.Name)\n\tassert.Equal(t, taskUpdateToolDesc, info.Desc)\n\n\tresult, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"status\": \"in_progress\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"Updated task #1\")\n\tassert.Contains(t, result, \"status\")\n\n\tcontent, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tassert.NoError(t, err)\n\tvar updated task\n\t_ = sonic.UnmarshalString(content.Content, &updated)\n\tassert.Equal(t, taskStatusInProgress, updated.Status)\n\n\tresult, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"subject\": \"New Subject\", \"description\": \"New description\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"subject\")\n\tassert.Contains(t, result, \"description\")\n\n\tcontent, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\t_ = sonic.UnmarshalString(content.Content, &updated)\n\tassert.Equal(t, \"New Subject\", updated.Subject)\n\tassert.Equal(t, \"New description\", updated.Description)\n}\n\nfunc TestTaskUpdateToolOwnerAndMetadata(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttaskData := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Test Task\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttaskJSON, _ := sonic.MarshalString(taskData)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: taskJSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\tresult, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"owner\": \"agent1\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"owner\")\n\n\tcontent, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tvar updated task\n\t_ = sonic.UnmarshalString(content.Content, &updated)\n\tassert.Equal(t, \"agent1\", updated.Owner)\n\n\tresult, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"metadata\": {\"key1\": \"value1\", \"key2\": \"value2\"}}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"metadata\")\n\n\tcontent, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\t_ = sonic.UnmarshalString(content.Content, &updated)\n\tassert.Equal(t, \"value1\", updated.Metadata[\"key1\"])\n\tassert.Equal(t, \"value2\", updated.Metadata[\"key2\"])\n\n\tresult, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"metadata\": {\"key1\": null, \"key3\": \"value3\"}}`)\n\tassert.NoError(t, err)\n\n\tcontent, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tvar updated2 task\n\t_ = sonic.UnmarshalString(content.Content, &updated2)\n\t_, key1Exists := updated2.Metadata[\"key1\"]\n\tassert.False(t, key1Exists)\n\tassert.Equal(t, \"value2\", updated2.Metadata[\"key2\"])\n\tassert.Equal(t, \"value3\", updated2.Metadata[\"key3\"])\n}\n\nfunc TestTaskUpdateToolBlocks(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Test Task\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{\n\t\tID:          \"2\",\n\t\tSubject:     \"Task 2\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\ttask3 := &task{\n\t\tID:          \"3\",\n\t\tSubject:     \"Task 3\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask3JSON, _ := sonic.MarshalString(task3)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"3.json\"), Content: task3JSON})\n\n\ttask4 := &task{\n\t\tID:          \"4\",\n\t\tSubject:     \"Task 4\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask4JSON, _ := sonic.MarshalString(task4)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"4.json\"), Content: task4JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\tresult, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlocks\": [\"2\", \"3\"]}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"blocks\")\n\n\tcontent, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tvar updated task\n\t_ = sonic.UnmarshalString(content.Content, &updated)\n\tassert.Equal(t, []string{\"2\", \"3\"}, updated.Blocks)\n\n\tresult, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlockedBy\": [\"4\"]}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"blockedBy\")\n\n\tcontent, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\t_ = sonic.UnmarshalString(content.Content, &updated)\n\tassert.Equal(t, []string{\"4\"}, updated.BlockedBy)\n}\n\nfunc TestTaskUpdateToolDelete(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttaskData := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Test Task\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t}\n\ttaskJSON, _ := sonic.MarshalString(taskData)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: taskJSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\tresult, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"status\": \"deleted\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"deleted\")\n\n\t_, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tassert.Error(t, err)\n}\n\nfunc TestTaskUpdateToolInvalidTaskID(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"../../../etc/passwd\", \"status\": \"in_progress\"}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"validate task ID failed\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"abc\", \"status\": \"in_progress\"}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"validate task ID failed\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1.5\", \"status\": \"in_progress\"}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"validate task ID failed\")\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlocks\": [\"invalid\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"validate blocked task ID failed\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlockedBy\": [\"invalid\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"validate blocking task ID failed\")\n}\n\nfunc TestTaskUpdateToolBlocksDeduplication(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{\n\t\tID:          \"2\",\n\t\tSubject:     \"Task 2\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{\"1\"},\n\t}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\ttask3 := &task{\n\t\tID:          \"3\",\n\t\tSubject:     \"Task 3\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{\"1\"},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask3JSON, _ := sonic.MarshalString(task3)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"3.json\"), Content: task3JSON})\n\n\ttask4 := &task{\n\t\tID:          \"4\",\n\t\tSubject:     \"Task 4\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask4JSON, _ := sonic.MarshalString(task4)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"4.json\"), Content: task4JSON})\n\n\ttask5 := &task{\n\t\tID:          \"5\",\n\t\tSubject:     \"Task 5\",\n\t\tDescription: \"Test description\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask5JSON, _ := sonic.MarshalString(task5)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"5.json\"), Content: task5JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlocks\": [\"2\", \"4\", \"4\"]}`)\n\tassert.NoError(t, err)\n\n\tcontent, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tvar updated task\n\t_ = sonic.UnmarshalString(content.Content, &updated)\n\tassert.Equal(t, []string{\"2\", \"4\"}, updated.Blocks)\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlockedBy\": [\"3\", \"5\", \"5\"]}`)\n\tassert.NoError(t, err)\n\n\tcontent, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\t_ = sonic.UnmarshalString(content.Content, &updated)\n\tassert.Equal(t, []string{\"3\", \"5\"}, updated.BlockedBy)\n}\n\nfunc TestTaskUpdateToolBidirectionalBlocks(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"First task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{\n\t\tID:          \"2\",\n\t\tSubject:     \"Task 2\",\n\t\tDescription: \"Second task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\ttask3 := &task{\n\t\tID:          \"3\",\n\t\tSubject:     \"Task 3\",\n\t\tDescription: \"Third task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask3JSON, _ := sonic.MarshalString(task3)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"3.json\"), Content: task3JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlocks\": [\"2\", \"3\"]}`)\n\tassert.NoError(t, err)\n\n\tcontent1, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tvar updatedTask1 task\n\t_ = sonic.UnmarshalString(content1.Content, &updatedTask1)\n\tassert.Equal(t, []string{\"2\", \"3\"}, updatedTask1.Blocks)\n\tassert.Empty(t, updatedTask1.BlockedBy)\n\n\tcontent2, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"2.json\")})\n\tvar updatedTask2 task\n\t_ = sonic.UnmarshalString(content2.Content, &updatedTask2)\n\tassert.Empty(t, updatedTask2.Blocks)\n\tassert.Equal(t, []string{\"1\"}, updatedTask2.BlockedBy)\n\n\tcontent3, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"3.json\")})\n\tvar updatedTask3 task\n\t_ = sonic.UnmarshalString(content3.Content, &updatedTask3)\n\tassert.Empty(t, updatedTask3.Blocks)\n\tassert.Equal(t, []string{\"1\"}, updatedTask3.BlockedBy)\n}\n\nfunc TestTaskUpdateToolBidirectionalBlockedBy(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"First task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{\n\t\tID:          \"2\",\n\t\tSubject:     \"Task 2\",\n\t\tDescription: \"Second task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\ttask3 := &task{\n\t\tID:          \"3\",\n\t\tSubject:     \"Task 3\",\n\t\tDescription: \"Third task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask3JSON, _ := sonic.MarshalString(task3)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"3.json\"), Content: task3JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"3\", \"addBlockedBy\": [\"1\", \"2\"]}`)\n\tassert.NoError(t, err)\n\n\tcontent3, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"3.json\")})\n\tvar updatedTask3 task\n\t_ = sonic.UnmarshalString(content3.Content, &updatedTask3)\n\tassert.Empty(t, updatedTask3.Blocks)\n\tassert.Equal(t, []string{\"1\", \"2\"}, updatedTask3.BlockedBy)\n\n\tcontent1, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tvar updatedTask1 task\n\t_ = sonic.UnmarshalString(content1.Content, &updatedTask1)\n\tassert.Equal(t, []string{\"3\"}, updatedTask1.Blocks)\n\tassert.Empty(t, updatedTask1.BlockedBy)\n\n\tcontent2, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"2.json\")})\n\tvar updatedTask2 task\n\t_ = sonic.UnmarshalString(content2.Content, &updatedTask2)\n\tassert.Equal(t, []string{\"3\"}, updatedTask2.Blocks)\n\tassert.Empty(t, updatedTask2.BlockedBy)\n}\n\nfunc TestTaskUpdateToolBidirectionalWithNonExistentTask(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"First task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlocks\": [\"999\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"update Task #1 blocks failed\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlockedBy\": [\"999\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"update Task #1 blockedBy failed\")\n}\n\nfunc TestTaskUpdateToolCyclicDependencyDetection(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"First task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{\n\t\tID:          \"2\",\n\t\tSubject:     \"Task 2\",\n\t\tDescription: \"Second task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\ttask3 := &task{\n\t\tID:          \"3\",\n\t\tSubject:     \"Task 3\",\n\t\tDescription: \"Third task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask3JSON, _ := sonic.MarshalString(task3)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"3.json\"), Content: task3JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlocks\": [\"1\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"cyclic dependency\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlockedBy\": [\"1\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"cyclic dependency\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlocks\": [\"2\"]}`)\n\tassert.NoError(t, err)\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"2\", \"addBlocks\": [\"1\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"cyclic dependency\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlockedBy\": [\"2\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"cyclic dependency\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"2\", \"addBlocks\": [\"3\"]}`)\n\tassert.NoError(t, err)\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"3\", \"addBlocks\": [\"1\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"cyclic dependency\")\n\n\t_, err = tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"addBlockedBy\": [\"3\"]}`)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"cyclic dependency\")\n\n\tcontent1, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tvar updatedTask1 task\n\t_ = sonic.UnmarshalString(content1.Content, &updatedTask1)\n\tassert.Equal(t, []string{\"2\"}, updatedTask1.Blocks)\n\tassert.Empty(t, updatedTask1.BlockedBy)\n\n\tcontent2, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"2.json\")})\n\tvar updatedTask2 task\n\t_ = sonic.UnmarshalString(content2.Content, &updatedTask2)\n\tassert.Equal(t, []string{\"3\"}, updatedTask2.Blocks)\n\tassert.Equal(t, []string{\"1\"}, updatedTask2.BlockedBy)\n\n\tcontent3, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"3.json\")})\n\tvar updatedTask3 task\n\t_ = sonic.UnmarshalString(content3.Content, &updatedTask3)\n\tassert.Empty(t, updatedTask3.Blocks)\n\tassert.Equal(t, []string{\"2\"}, updatedTask3.BlockedBy)\n}\n\nfunc TestTaskUpdateToolDeleteCleansDependencies(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"First task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{\"2\", \"3\"},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{\n\t\tID:          \"2\",\n\t\tSubject:     \"Task 2\",\n\t\tDescription: \"Second task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{\"3\"},\n\t\tBlockedBy:   []string{\"1\"},\n\t}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\ttask3 := &task{\n\t\tID:          \"3\",\n\t\tSubject:     \"Task 3\",\n\t\tDescription: \"Third task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{\"1\", \"2\"},\n\t}\n\ttask3JSON, _ := sonic.MarshalString(task3)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"3.json\"), Content: task3JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\tresult, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"status\": \"deleted\"}`)\n\tassert.NoError(t, err)\n\tassert.Contains(t, result, \"deleted\")\n\n\t_, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tassert.Error(t, err)\n\n\tcontent2, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"2.json\")})\n\tassert.NoError(t, err)\n\tvar updatedTask2 task\n\t_ = sonic.UnmarshalString(content2.Content, &updatedTask2)\n\tassert.Equal(t, []string{\"3\"}, updatedTask2.Blocks)\n\tassert.Empty(t, updatedTask2.BlockedBy)\n\n\tcontent3, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"3.json\")})\n\tassert.NoError(t, err)\n\tvar updatedTask3 task\n\t_ = sonic.UnmarshalString(content3.Content, &updatedTask3)\n\tassert.Empty(t, updatedTask3.Blocks)\n\tassert.Equal(t, []string{\"2\"}, updatedTask3.BlockedBy)\n}\n\nfunc TestTaskUpdateToolAutoDeleteAllTasksWhenAllCompleted(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"First task\",\n\t\tStatus:      taskStatusCompleted,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{\n\t\tID:          \"2\",\n\t\tSubject:     \"Task 2\",\n\t\tDescription: \"Second task\",\n\t\tStatus:      taskStatusCompleted,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\ttask3 := &task{\n\t\tID:          \"3\",\n\t\tSubject:     \"Task 3\",\n\t\tDescription: \"Third task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask3JSON, _ := sonic.MarshalString(task3)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"3.json\"), Content: task3JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"3\", \"status\": \"completed\"}`)\n\tassert.NoError(t, err)\n\n\t_, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tassert.Error(t, err)\n\t_, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"2.json\")})\n\tassert.Error(t, err)\n\t_, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"3.json\")})\n\tassert.Error(t, err)\n}\n\nfunc TestTaskUpdateToolNoDeleteWhenNotAllCompleted(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newInMemoryBackend()\n\tbaseDir := \"/tmp/tasks\"\n\tlock := &sync.Mutex{}\n\n\ttask1 := &task{\n\t\tID:          \"1\",\n\t\tSubject:     \"Task 1\",\n\t\tDescription: \"First task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask1JSON, _ := sonic.MarshalString(task1)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"1.json\"), Content: task1JSON})\n\n\ttask2 := &task{\n\t\tID:          \"2\",\n\t\tSubject:     \"Task 2\",\n\t\tDescription: \"Second task\",\n\t\tStatus:      taskStatusPending,\n\t\tBlocks:      []string{},\n\t\tBlockedBy:   []string{},\n\t}\n\ttask2JSON, _ := sonic.MarshalString(task2)\n\t_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, \"2.json\"), Content: task2JSON})\n\n\ttool := newTaskUpdateTool(backend, baseDir, lock)\n\n\t_, err := tool.InvokableRun(ctx, `{\"taskId\": \"1\", \"status\": \"completed\"}`)\n\tassert.NoError(t, err)\n\n\t_, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tassert.NoError(t, err)\n\t_, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"2.json\")})\n\tassert.NoError(t, err)\n\n\tcontent1, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, \"1.json\")})\n\tvar updatedTask1 task\n\t_ = sonic.UnmarshalString(content1.Content, &updatedTask1)\n\tassert.Equal(t, taskStatusCompleted, updatedTask1.Status)\n}\n"
  },
  {
    "path": "adk/middlewares/reduction/consts.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package reduction provides middlewares to trim context and clear tool results.\npackage reduction\n\nimport \"github.com/cloudwego/eino/adk/internal\"\n\nconst (\n\ttruncFmt = `<persisted-output>\nOutput too large ({original_size}). Full output saved to: {file_path}\nPreview (first {preview_size}):\n{preview_first}\n\nPreview (last {preview_size}):\n{preview_last}\n\n</persisted-output>`\n\ttruncFmtZh = `<persisted-output>\n输出结果过大 ({original_size}). 完整输出保存到: {file_path}\n预览 (前 {preview_size}):\n{preview_first}\n\n预览 (后 {preview_size}):\n{preview_last}\n\n</persisted-output>`\n)\n\nconst (\n\tclearWithOffloadingFmt = `<persisted-output>Tool result saved to: {file_path}\nUse {read_tool_name} to view</persisted-output>`\n\tclearWithOffloadingFmtZh = `<persisted-output>工具结果已保存至: {file_path}\n使用 {read_tool_name} 进行查看</persisted-output>`\n\n\tclearWithoutOffloadingFmt   = `[Old tool result content cleared]`\n\tclearWithoutOffloadingFmtZh = `[工具输出结果已清理]`\n)\n\nconst (\n\tmsgReducedFlag   = \"_reduction_mw_processed\"\n\tmsgReducedTokens = \"_reduction_mw_tokens\"\n)\n\nfunc getTruncFmt() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: truncFmt,\n\t\tChinese: truncFmtZh,\n\t})\n}\n\nfunc getClearWithOffloadingFmt() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: clearWithOffloadingFmt,\n\t\tChinese: clearWithOffloadingFmtZh,\n\t})\n}\n\nfunc getClearWithoutOffloadingFmt() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: clearWithoutOffloadingFmt,\n\t\tChinese: clearWithoutOffloadingFmtZh,\n\t})\n}\n\ntype scene int\n\nconst (\n\tsceneTruncation scene = 1\n\tsceneClear      scene = 2\n)\n"
  },
  {
    "path": "adk/middlewares/reduction/internal/clear_tool_result.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package internal provides middlewares to trim context and clear tool results.\npackage internal\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// ClearToolResultConfig configures the tool result clearing middleware.\n// This middleware clears old tool results when their total token count exceeds a threshold,\n// while protecting recent messages within a token budget.\ntype ClearToolResultConfig struct {\n\t// ToolResultTokenThreshold is the threshold for total tool result tokens.\n\t// When the sum of all tool result tokens exceeds this threshold, old tool results\n\t// (outside the KeepRecentTokens range) will be replaced with a placeholder.\n\t// Token estimation uses a simple heuristic: character count / 4.\n\t// If 0, defaults to 20000.\n\tToolResultTokenThreshold int\n\n\t// KeepRecentTokens is the token budget for recent messages to keep intact.\n\t// Messages within this token budget from the end will not have their tool results cleared,\n\t// even if the total tool result tokens exceed the threshold.\n\t// If 0, defaults to 40000.\n\tKeepRecentTokens int\n\n\t// ClearToolResultPlaceholder is the text to replace old tool results with.\n\t// If empty, defaults to \"[Old tool result content cleared]\".\n\tClearToolResultPlaceholder string\n\n\t// TokenCounter is a custom function to estimate token count for a message.\n\t// If nil, uses the default counter (character count / 4).\n\tTokenCounter func(msg *schema.Message) int\n\n\t// ExcludeTools is a list of tool names whose results should never be cleared.\n\tExcludeTools []string\n}\n\n// NewClearToolResult creates a new middleware that clears old tool results\n// based on token thresholds while protecting recent messages.\nfunc NewClearToolResult(ctx context.Context, config *ClearToolResultConfig) (adk.AgentMiddleware, error) {\n\treturn adk.AgentMiddleware{\n\t\tBeforeChatModel: newClearToolResult(ctx, config),\n\t}, nil\n}\n\nfunc newClearToolResult(ctx context.Context, config *ClearToolResultConfig) func(ctx context.Context, state *adk.ChatModelAgentState) error {\n\tif config == nil {\n\t\tconfig = &ClearToolResultConfig{}\n\t}\n\n\t// Set defaults\n\ttoolResultTokenThreshold := config.ToolResultTokenThreshold\n\tif toolResultTokenThreshold == 0 {\n\t\ttoolResultTokenThreshold = 20000\n\t}\n\n\tkeepRecentTokens := config.KeepRecentTokens\n\tif keepRecentTokens == 0 {\n\t\tkeepRecentTokens = 40000\n\t}\n\n\tplaceholder := config.ClearToolResultPlaceholder\n\tif placeholder == \"\" {\n\t\tplaceholder = \"[Old tool result content cleared]\"\n\t}\n\n\t// Set token estimator\n\tcounter := config.TokenCounter\n\tif counter == nil {\n\t\tcounter = defaultTokenCounter\n\t}\n\treturn func(ctx context.Context, state *adk.ChatModelAgentState) error {\n\t\treturn reduceByTokens(state, toolResultTokenThreshold, keepRecentTokens, placeholder, counter, config.ExcludeTools)\n\t}\n}\n\n// defaultTokenCounter estimates token count using character count / 4\n// This is a simple heuristic that works reasonably well for most languages\nfunc defaultTokenCounter(msg *schema.Message) int {\n\tcount := len(msg.Content)\n\n\t// Also count tool call arguments if present\n\tfor _, tc := range msg.ToolCalls {\n\t\tcount += len(tc.Function.Arguments)\n\t}\n\n\t// Estimate: roughly 4 characters per token\n\treturn (count + 3) / 4\n}\n\n// reduceByTokens reduces context based on tool result token threshold and recent message protection.\n// It clears old tool results when:\n// 1. The total tokens of all tool results exceed toolResultTokenThreshold\n// 2. Only tool results outside the keepRecentTokens range (from the end) are cleared\nfunc reduceByTokens(state *adk.ChatModelAgentState, toolResultTokenThreshold, keepRecentTokens int, placeholder string, counter func(*schema.Message) int, excludedTools []string) error {\n\tif len(state.Messages) == 0 {\n\t\treturn nil\n\t}\n\n\t// Step 1: Calculate total tool result tokens\n\ttotalToolResultTokens := 0\n\tfor _, msg := range state.Messages {\n\t\tif msg.Role == schema.Tool && msg.Content != placeholder {\n\t\t\ttotalToolResultTokens += counter(msg)\n\t\t}\n\t}\n\n\t// If total tool result tokens are under the threshold, no reduction needed\n\tif totalToolResultTokens <= toolResultTokenThreshold {\n\t\treturn nil\n\t}\n\n\t// Step 2: Calculate the index from which to protect recent messages\n\t// We need to find the starting index where cumulative tokens from the end <= keepRecentTokens\n\trecentStartIdx := len(state.Messages)\n\tcumulativeTokens := 0\n\n\tfor i := len(state.Messages) - 1; i >= 0; i-- {\n\t\tmsgTokens := counter(state.Messages[i])\n\t\tif cumulativeTokens+msgTokens > keepRecentTokens {\n\t\t\t// Adding this message would exceed the budget, so stop here\n\t\t\trecentStartIdx = i\n\t\t\tbreak\n\t\t}\n\t\tcumulativeTokens += msgTokens\n\t\trecentStartIdx = i\n\t}\n\n\t// Step 3: Clear tool results outside the protected range (before recentStartIdx)\n\tfor i := 0; i < recentStartIdx; i++ {\n\t\tmsg := state.Messages[i]\n\t\tif msg.Role == schema.Tool && msg.Content != placeholder && !excluded(msg.ToolName, excludedTools) {\n\t\t\tmsg.Content = placeholder\n\t\t}\n\t}\n\n\treturn nil\n}\n\nfunc excluded(name string, exclude []string) bool {\n\tfor _, ex := range exclude {\n\t\tif name == ex {\n\t\t\treturn true\n\t\t}\n\t}\n\treturn false\n}\n"
  },
  {
    "path": "adk/middlewares/reduction/internal/clear_tool_result_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc Test_reduceByTokens(t *testing.T) {\n\ttype args struct {\n\t\tstate                    *adk.ChatModelAgentState\n\t\ttoolResultTokenThreshold int\n\t\tkeepRecentTokens         int\n\t\tplaceholder              string\n\t\testimator                func(*schema.Message) int\n\t}\n\ttests := []struct {\n\t\tname          string\n\t\targs          args\n\t\twantErr       assert.ErrorAssertionFunc\n\t\tvalidateState func(*testing.T, *adk.ChatModelAgentState)\n\t}{\n\t\t{\n\t\t\tname: \"no reduction when tool result tokens under threshold\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{\n\t\t\t\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t\t\t\tschema.AssistantMessage(\"hi\", nil),\n\t\t\t\t\t\tschema.ToolMessage(\"short tool result\", \"call-1\", schema.WithToolName(\"tool1\")),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 100,\n\t\t\t\tkeepRecentTokens:         500,\n\t\t\t\tplaceholder:              \"[Old tool result content cleared]\",\n\t\t\t\testimator:                defaultTokenCounter,\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\tassert.Equal(t, \"short tool result\", state.Messages[2].Content)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"clear old tool results when total exceeds threshold\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{\n\t\t\t\t\t\tschema.UserMessage(\"msg1\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"a\", 40), \"call-1\", schema.WithToolName(\"tool1\")), // ~10 tokens (old)\n\t\t\t\t\t\tschema.UserMessage(\"msg2\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"b\", 40), \"call-2\", schema.WithToolName(\"tool2\")), // ~10 tokens (old)\n\t\t\t\t\t\tschema.UserMessage(\"msg3\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"c\", 40), \"call-3\", schema.WithToolName(\"tool3\")), // ~10 tokens (recent, protected)\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 20,\n\t\t\t\tkeepRecentTokens:         10,\n\t\t\t\tplaceholder:              \"[Old tool result content cleared]\",\n\t\t\t\testimator:                defaultTokenCounter,\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\tassert.Equal(t, \"[Old tool result content cleared]\", state.Messages[1].Content)\n\t\t\t\tassert.Equal(t, \"[Old tool result content cleared]\", state.Messages[3].Content)\n\t\t\t\tassert.Equal(t, strings.Repeat(\"c\", 40), state.Messages[5].Content)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"protect recent messages even when tool results exceed threshold\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{\n\t\t\t\t\t\tschema.UserMessage(\"old msg\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"x\", 100), \"call-1\", schema.WithToolName(\"tool1\")), // ~25 tokens (old)\n\t\t\t\t\t\tschema.UserMessage(\"recent msg\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"x\", 100), \"call-2\", schema.WithToolName(\"tool2\")), // ~25 tokens (recent, protected)\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 10,\n\t\t\t\tkeepRecentTokens:         20,\n\t\t\t\tplaceholder:              \"[Old tool result content cleared]\",\n\t\t\t\testimator:                defaultTokenCounter,\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\t// Total tool result tokens = 50, exceeds threshold of 10\n\t\t\t\t// But last 200 tokens are protected (includes last 2 messages)\n\t\t\t\t// So only the first tool result should be cleared\n\t\t\t\tassert.Equal(t, \"[Old tool result content cleared]\", state.Messages[1].Content)\n\t\t\t\tassert.Equal(t, strings.Repeat(\"x\", 100), state.Messages[3].Content)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"custom placeholder text\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{\n\t\t\t\t\t\tschema.UserMessage(\"msg\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"x\", 100), \"call-1\", schema.WithToolName(\"tool1\")),\n\t\t\t\t\t\tschema.UserMessage(strings.Repeat(\"x\", 100)),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 10,\n\t\t\t\tkeepRecentTokens:         20,\n\t\t\t\tplaceholder:              \"[历史工具结果已清除]\",\n\t\t\t\testimator:                defaultTokenCounter,\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\tassert.Equal(t, \"[历史工具结果已清除]\", state.Messages[1].Content)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"no tool messages\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{\n\t\t\t\t\t\tschema.UserMessage(\"msg 1\"),\n\t\t\t\t\t\tschema.AssistantMessage(\"response 1\", nil),\n\t\t\t\t\t\tschema.UserMessage(\"msg 2\"),\n\t\t\t\t\t\tschema.AssistantMessage(\"response 2\", nil),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 10,\n\t\t\t\tkeepRecentTokens:         10,\n\t\t\t\tplaceholder:              \"[Old tool result content cleared]\",\n\t\t\t\testimator:                defaultTokenCounter,\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\t// All messages should remain unchanged\n\t\t\t\tassert.Equal(t, \"msg 1\", state.Messages[0].Content)\n\t\t\t\tassert.Equal(t, \"response 1\", state.Messages[1].Content)\n\t\t\t\tassert.Equal(t, \"msg 2\", state.Messages[2].Content)\n\t\t\t\tassert.Equal(t, \"response 2\", state.Messages[3].Content)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"empty messages\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 100,\n\t\t\t\tkeepRecentTokens:         500,\n\t\t\t\tplaceholder:              \"[Old tool result content cleared]\",\n\t\t\t\testimator:                defaultTokenCounter,\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\tassert.Empty(t, state.Messages)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"custom token estimator - word count\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{\n\t\t\t\t\t\tschema.UserMessage(\"hello world\"),\n\t\t\t\t\t\tschema.ToolMessage(\"this is a long tool result\", \"call-1\", schema.WithToolName(\"tool1\")), // 6 words (old)\n\t\t\t\t\t\tschema.UserMessage(\"another message\"),\n\t\t\t\t\t\tschema.ToolMessage(\"recent tool result here\", \"call-2\", schema.WithToolName(\"tool2\")), // 4 words (recent)\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 9, // 10 words total threshold\n\t\t\t\tkeepRecentTokens:         5, // 15 words protection budget\n\t\t\t\tplaceholder:              \"[Old tool result content cleared]\",\n\t\t\t\testimator: func(msg *schema.Message) int {\n\t\t\t\t\tif msg.Content == \"\" {\n\t\t\t\t\t\treturn 0\n\t\t\t\t\t}\n\t\t\t\t\twords := 1\n\t\t\t\t\tfor _, ch := range msg.Content {\n\t\t\t\t\t\tif ch == ' ' {\n\t\t\t\t\t\t\twords++\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\treturn words\n\t\t\t\t},\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\tassert.Equal(t, \"[Old tool result content cleared]\", state.Messages[1].Content)\n\t\t\t\tassert.Equal(t, \"recent tool result here\", state.Messages[3].Content)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"already cleared results are not counted\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{\n\t\t\t\t\t\tschema.UserMessage(\"msg1\"),\n\t\t\t\t\t\tschema.ToolMessage(\"[Old tool result content cleared]\", \"call-1\", schema.WithToolName(\"tool1\")), // Already cleared\n\t\t\t\t\t\tschema.UserMessage(\"msg2\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"a\", 100), \"call-2\", schema.WithToolName(\"tool2\")), // New long result\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 10,\n\t\t\t\tkeepRecentTokens:         20,\n\t\t\t\tplaceholder:              \"[Old tool result content cleared]\",\n\t\t\t\testimator:                defaultTokenCounter,\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\t// Only the new long result counts toward the threshold\n\t\t\t\t// Both should have placeholder\n\t\t\t\tassert.Equal(t, \"[Old tool result content cleared]\", state.Messages[1].Content)\n\t\t\t\tassert.Equal(t, strings.Repeat(\"a\", 100), state.Messages[3].Content)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"all tool results within protected range\",\n\t\t\targs: args{\n\t\t\t\tstate: &adk.ChatModelAgentState{\n\t\t\t\t\tMessages: []adk.Message{\n\t\t\t\t\t\tschema.UserMessage(\"msg1\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"a\", 40), \"call-1\", schema.WithToolName(\"tool1\")), // ~10 tokens\n\t\t\t\t\t\tschema.UserMessage(\"msg2\"),\n\t\t\t\t\t\tschema.ToolMessage(strings.Repeat(\"b\", 40), \"call-2\", schema.WithToolName(\"tool2\")), // ~10 tokens\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\ttoolResultTokenThreshold: 10,   // Low threshold (will exceed)\n\t\t\t\tkeepRecentTokens:         1000, // Very high protection (protects all)\n\t\t\t\tplaceholder:              \"[Old tool result content cleared]\",\n\t\t\t\testimator:                defaultTokenCounter,\n\t\t\t},\n\t\t\twantErr: assert.NoError,\n\t\t\tvalidateState: func(t *testing.T, state *adk.ChatModelAgentState) {\n\t\t\t\t// All messages are within protected range, nothing should be cleared\n\t\t\t\tassert.Equal(t, strings.Repeat(\"a\", 40), state.Messages[1].Content)\n\t\t\t\tassert.Equal(t, strings.Repeat(\"b\", 40), state.Messages[3].Content)\n\t\t\t},\n\t\t},\n\t}\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\terr := reduceByTokens(tt.args.state, tt.args.toolResultTokenThreshold, tt.args.keepRecentTokens, tt.args.placeholder, tt.args.estimator, []string{})\n\t\t\ttt.wantErr(t, err, fmt.Sprintf(\"reduceByTokens(%v, %v, %v, %v)\", tt.args.state, tt.args.toolResultTokenThreshold, tt.args.keepRecentTokens, tt.args.placeholder))\n\t\t\tif tt.validateState != nil {\n\t\t\t\ttt.validateState(t, tt.args.state)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc Test_newClearToolResult(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"nil config uses defaults\", func(t *testing.T) {\n\t\tfn := newClearToolResult(ctx, nil)\n\t\tassert.NotNil(t, fn)\n\n\t\t// Test that function works with nil config (uses defaults)\n\t\tstate := &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t\tschema.ToolMessage(\"short result\", \"call-1\", schema.WithToolName(\"tool1\")),\n\t\t\t},\n\t\t}\n\t\terr := fn(ctx, state)\n\t\tassert.NoError(t, err)\n\t\t// Default threshold is 20000, so short result should not be cleared\n\t\tassert.Equal(t, \"short result\", state.Messages[1].Content)\n\t})\n\n\tt.Run(\"empty config uses defaults\", func(t *testing.T) {\n\t\tfn := newClearToolResult(ctx, &ClearToolResultConfig{})\n\t\tassert.NotNil(t, fn)\n\n\t\tstate := &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t\tschema.ToolMessage(\"short result\", \"call-1\", schema.WithToolName(\"tool1\")),\n\t\t\t},\n\t\t}\n\t\terr := fn(ctx, state)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"short result\", state.Messages[1].Content)\n\t})\n}\n"
  },
  {
    "path": "adk/middlewares/reduction/internal/large_tool_result.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport (\n\t\"bufio\"\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"strings\"\n\t\"unicode/utf8\"\n\n\t\"github.com/slongfield/pyfmt\"\n\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nconst (\n\ttooLargeToolMessage = `Tool result too large, the result of this tool call {tool_call_id} was saved in the filesystem at this path: {file_path}\nYou can read the result from the filesystem by using the '{read_file_tool_name}' tool, but make sure to only read part of the result at a time.\nYou can do this by specifying an offset and limit in the '{read_file_tool_name}' tool call.\nFor example, to read the first 100 lines, you can use the '{read_file_tool_name}' tool with offset=0 and limit=100.\n\nHere are the first 10 lines of the result:\n{content_sample}`\n\n\ttooLargeToolMessageChinese = `工具结果过大，此工具调用 {tool_call_id} 的结果已保存到文件系统的以下路径：{file_path}\n你可以使用 '{read_file_tool_name}' 工具从文件系统读取结果，但请确保每次只读取部分结果。\n你可以通过在 '{read_file_tool_name}' 工具调用中指定 offset 和 limit 来实现。\n例如，要读取前 100 行，你可以使用 '{read_file_tool_name}' 工具，设置 offset=0 和 limit=100。\n\n以下是结果的前 10 行：\n{content_sample}`\n)\n\ntype toolResultOffloadingConfig struct {\n\tBackend          Backend\n\tReadFileToolName string\n\tTokenLimit       int\n\tPathGenerator    func(ctx context.Context, input *compose.ToolInput) (string, error)\n\tTokenCounter     func(msg *schema.Message) int\n}\n\nfunc newToolResultOffloading(_ context.Context, config *toolResultOffloadingConfig) compose.ToolMiddleware {\n\toffloading := &toolResultOffloading{\n\t\tbackend:       config.Backend,\n\t\ttokenLimit:    config.TokenLimit,\n\t\tpathGenerator: config.PathGenerator,\n\t\ttoolName:      config.ReadFileToolName,\n\t\tcounter:       config.TokenCounter,\n\t}\n\n\tif offloading.tokenLimit == 0 {\n\t\toffloading.tokenLimit = 20000\n\t}\n\n\tif offloading.pathGenerator == nil {\n\t\toffloading.pathGenerator = func(ctx context.Context, input *compose.ToolInput) (string, error) {\n\t\t\treturn fmt.Sprintf(\"/large_tool_result/%s\", input.CallID), nil\n\t\t}\n\t}\n\n\tif len(offloading.toolName) == 0 {\n\t\toffloading.toolName = \"read_file\"\n\t}\n\n\tif offloading.counter == nil {\n\t\toffloading.counter = defaultTokenCounter\n\t}\n\n\treturn compose.ToolMiddleware{\n\t\tInvokable:  offloading.invoke,\n\t\tStreamable: offloading.stream,\n\t}\n}\n\ntype toolResultOffloading struct {\n\tbackend       Backend\n\ttokenLimit    int\n\tpathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error)\n\ttoolName      string\n\tcounter       func(msg *schema.Message) int\n}\n\nfunc (t *toolResultOffloading) invoke(endpoint compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {\n\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\toutput, err := endpoint(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tresult, err := t.handleResult(ctx, output.Result, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &compose.ToolOutput{Result: result}, nil\n\t}\n}\n\nfunc (t *toolResultOffloading) stream(endpoint compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {\n\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\toutput, err := endpoint(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tresult, err := concatString(output.Result)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tresult, err = t.handleResult(ctx, result, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &compose.StreamToolOutput{Result: schema.StreamReaderFromArray([]string{result})}, nil\n\t}\n}\n\nfunc (t *toolResultOffloading) handleResult(ctx context.Context, result string, input *compose.ToolInput) (string, error) {\n\tif t.counter(schema.ToolMessage(result, input.CallID, schema.WithToolName(input.Name))) > t.tokenLimit*4 {\n\t\tpath, err := t.pathGenerator(ctx, input)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\tnResult := formatToolMessage(result)\n\t\ttpl := internal.SelectPrompt(internal.I18nPrompts{\n\t\t\tEnglish: tooLargeToolMessage,\n\t\t\tChinese: tooLargeToolMessageChinese,\n\t\t})\n\t\tnResult, err = pyfmt.Fmt(tpl, map[string]any{\n\t\t\t\"tool_call_id\":        input.CallID,\n\t\t\t\"file_path\":           path,\n\t\t\t\"content_sample\":      nResult,\n\t\t\t\"read_file_tool_name\": t.toolName,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\terr = t.backend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: path,\n\t\t\tContent:  result,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn nResult, nil\n\t}\n\n\treturn result, nil\n}\n\nfunc concatString(sr *schema.StreamReader[string]) (string, error) {\n\tif sr == nil {\n\t\treturn \"\", errors.New(\"stream is nil\")\n\t}\n\tsb := strings.Builder{}\n\tfor {\n\t\tstr, err := sr.Recv()\n\t\tif errors.Is(err, io.EOF) {\n\t\t\treturn sb.String(), nil\n\t\t}\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tsb.WriteString(str)\n\t}\n}\n\nfunc formatToolMessage(s string) string {\n\treader := bufio.NewScanner(strings.NewReader(s))\n\tvar b strings.Builder\n\n\tlineNum := 1\n\tfor reader.Scan() {\n\t\tif lineNum > 10 {\n\t\t\tbreak\n\t\t}\n\t\tline := reader.Text()\n\n\t\tif utf8.RuneCountInString(line) > 1000 {\n\t\t\trunes := []rune(line)\n\t\t\tline = string(runes[:1000])\n\t\t}\n\n\t\tb.WriteString(fmt.Sprintf(\"%d: %s\\n\", lineNum, line))\n\n\t\tlineNum++\n\t}\n\n\treturn b.String()\n}\n"
  },
  {
    "path": "adk/middlewares/reduction/internal/large_tool_result_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// mockBackend is a simple in-memory backend for testing\ntype mockBackend struct {\n\tfiles map[string]string\n}\n\nfunc newMockBackend() *mockBackend {\n\treturn &mockBackend{\n\t\tfiles: make(map[string]string),\n\t}\n}\n\nfunc (m *mockBackend) Write(_ context.Context, wr *filesystem.WriteRequest) error {\n\tm.files[wr.FilePath] = wr.Content\n\treturn nil\n}\n\nfunc TestToolResultOffloading_SmallResult(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 100, // Small limit for testing\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\t// Create a mock endpoint that returns a small result\n\tsmallResult := \"This is a small result\"\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: smallResult}, nil\n\t}\n\n\t// Wrap the endpoint with the middleware\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\t// Execute\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_123\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Small result should pass through unchanged\n\tif output.Result != smallResult {\n\t\tt.Errorf(\"expected result %q, got %q\", smallResult, output.Result)\n\t}\n\n\t// No file should be written\n\tif len(backend.files) != 0 {\n\t\tt.Errorf(\"expected no files to be written, got %d files\", len(backend.files))\n\t}\n}\n\nfunc TestToolResultOffloading_LargeResult(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10, // Very small limit to trigger offloading\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\t// Create a large result (more than 10 * 4 = 40 bytes)\n\tlargeResult := strings.Repeat(\"This is a long line of text that will exceed the token limit.\\n\", 10)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: largeResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_456\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Result should be replaced with a message\n\tif !strings.Contains(output.Result, \"Tool result too large\") {\n\t\tt.Errorf(\"expected result to contain 'Tool result too large', got %q\", output.Result)\n\t}\n\n\tif !strings.Contains(output.Result, \"call_456\") {\n\t\tt.Errorf(\"expected result to contain call ID 'call_456', got %q\", output.Result)\n\t}\n\n\tif !strings.Contains(output.Result, \"/large_tool_result/call_456\") {\n\t\tt.Errorf(\"expected result to contain file path, got %q\", output.Result)\n\t}\n\n\t// File should be written\n\tif len(backend.files) != 1 {\n\t\tt.Fatalf(\"expected 1 file to be written, got %d files\", len(backend.files))\n\t}\n\n\tsavedContent, ok := backend.files[\"/large_tool_result/call_456\"]\n\tif !ok {\n\t\tt.Fatalf(\"expected file at /large_tool_result/call_456, got files: %v\", backend.files)\n\t}\n\n\tif savedContent != largeResult {\n\t\tt.Errorf(\"saved content doesn't match original result\")\n\t}\n}\n\nfunc TestToolResultOffloading_CustomPathGenerator(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tcustomPath := \"/custom/path/result.txt\"\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t\tPathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) {\n\t\t\treturn customPath, nil\n\t\t},\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\tlargeResult := strings.Repeat(\"Large content \", 100)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: largeResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_789\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Check custom path is used\n\tif !strings.Contains(output.Result, customPath) {\n\t\tt.Errorf(\"expected result to contain custom path %q, got %q\", customPath, output.Result)\n\t}\n\n\t// File should be written to custom path\n\tsavedContent, ok := backend.files[customPath]\n\tif !ok {\n\t\tt.Fatalf(\"expected file at %q, got files: %v\", customPath, backend.files)\n\t}\n\n\tif savedContent != largeResult {\n\t\tt.Errorf(\"saved content doesn't match original result\")\n\t}\n}\n\nfunc TestToolResultOffloading_PathGeneratorError(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\texpectedErr := errors.New(\"path generation failed\")\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t\tPathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) {\n\t\t\treturn \"\", expectedErr\n\t\t},\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\tlargeResult := strings.Repeat(\"Large content \", 100)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: largeResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_error\",\n\t}\n\t_, err := wrappedEndpoint(ctx, input)\n\n\tif err == nil {\n\t\tt.Fatal(\"expected error, got nil\")\n\t}\n\n\tif !errors.Is(err, expectedErr) {\n\t\tt.Errorf(\"expected error %v, got %v\", expectedErr, err)\n\t}\n}\n\nfunc TestToolResultOffloading_EndpointError(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 100,\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\texpectedErr := errors.New(\"endpoint execution failed\")\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn nil, expectedErr\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_endpoint_error\",\n\t}\n\t_, err := wrappedEndpoint(ctx, input)\n\n\tif err == nil {\n\t\tt.Fatal(\"expected error, got nil\")\n\t}\n\n\tif !errors.Is(err, expectedErr) {\n\t\tt.Errorf(\"expected error %v, got %v\", expectedErr, err)\n\t}\n}\n\nfunc TestToolResultOffloading_DefaultTokenLimit(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 0, // Should default to 20000\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\t// Create a result smaller than 20000 * 4 = 80000 bytes\n\tsmallResult := strings.Repeat(\"x\", 1000)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: smallResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_default\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Should pass through unchanged\n\tif output.Result != smallResult {\n\t\tt.Errorf(\"expected result to pass through unchanged\")\n\t}\n\n\t// No file should be written\n\tif len(backend.files) != 0 {\n\t\tt.Errorf(\"expected no files to be written, got %d files\", len(backend.files))\n\t}\n}\n\nfunc TestToolResultOffloading_Stream(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\t// Create a streaming endpoint that returns large content\n\tlargeResult := strings.Repeat(\"Large streaming content \", 100)\n\tmockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\t// Split the result into chunks\n\t\tchunks := []string{largeResult[:len(largeResult)/2], largeResult[len(largeResult)/2:]}\n\t\treturn &compose.StreamToolOutput{\n\t\t\tResult: schema.StreamReaderFromArray(chunks),\n\t\t}, nil\n\t}\n\n\twrappedEndpoint := middleware.Streamable(mockStreamEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_stream\",\n\t}\n\toutput, err := wrappedEndpoint(ctx, input)\n\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t}\n\n\t// Read the stream\n\tvar result strings.Builder\n\tfor {\n\t\tchunk, err := output.Result.Recv()\n\t\tif errors.Is(err, io.EOF) {\n\t\t\tbreak\n\t\t}\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"error reading stream: %v\", err)\n\t\t}\n\t\tresult.WriteString(chunk)\n\t}\n\n\tresultStr := result.String()\n\n\t// Result should be replaced with a message\n\tif !strings.Contains(resultStr, \"Tool result too large\") {\n\t\tt.Errorf(\"expected result to contain 'Tool result too large', got %q\", resultStr)\n\t}\n\n\tif !strings.Contains(resultStr, \"call_stream\") {\n\t\tt.Errorf(\"expected result to contain call ID 'call_stream', got %q\", resultStr)\n\t}\n\n\t// File should be written\n\tif len(backend.files) != 1 {\n\t\tt.Fatalf(\"expected 1 file to be written, got %d files\", len(backend.files))\n\t}\n\n\tsavedContent, ok := backend.files[\"/large_tool_result/call_stream\"]\n\tif !ok {\n\t\tt.Fatalf(\"expected file at /large_tool_result/call_stream, got files: %v\", backend.files)\n\t}\n\n\tif savedContent != largeResult {\n\t\tt.Errorf(\"saved content doesn't match original result\")\n\t}\n}\n\nfunc TestToolResultOffloading_StreamError(t *testing.T) {\n\tctx := context.Background()\n\tbackend := newMockBackend()\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\texpectedErr := errors.New(\"stream endpoint failed\")\n\tmockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\treturn nil, expectedErr\n\t}\n\n\twrappedEndpoint := middleware.Streamable(mockStreamEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_stream_error\",\n\t}\n\t_, err := wrappedEndpoint(ctx, input)\n\n\tif err == nil {\n\t\tt.Fatal(\"expected error, got nil\")\n\t}\n\n\tif !errors.Is(err, expectedErr) {\n\t\tt.Errorf(\"expected error %v, got %v\", expectedErr, err)\n\t}\n}\n\nfunc TestFormatToolMessage(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\tinput    string\n\t\texpected string\n\t}{\n\t\t{\n\t\t\tname:     \"single line\",\n\t\t\tinput:    \"single line\",\n\t\t\texpected: \"1: single line\\n\",\n\t\t},\n\t\t{\n\t\t\tname:     \"multiple lines\",\n\t\t\tinput:    \"line1\\nline2\\nline3\",\n\t\t\texpected: \"1: line1\\n2: line2\\n3: line3\\n\",\n\t\t},\n\t\t{\n\t\t\tname:     \"more than 10 lines\",\n\t\t\tinput:    \"1\\n2\\n3\\n4\\n5\\n6\\n7\\n8\\n9\\n10\\n11\\n12\",\n\t\t\texpected: \"1: 1\\n2: 2\\n3: 3\\n4: 4\\n5: 5\\n6: 6\\n7: 7\\n8: 8\\n9: 9\\n10: 10\\n\",\n\t\t},\n\t\t{\n\t\t\tname:     \"long line truncation\",\n\t\t\tinput:    strings.Repeat(\"a\", 1500),\n\t\t\texpected: fmt.Sprintf(\"1: %s\\n\", strings.Repeat(\"a\", 1000)),\n\t\t},\n\t\t{\n\t\t\tname:     \"unicode characters\",\n\t\t\tinput:    \"你好世界\\n测试\",\n\t\t\texpected: \"1: 你好世界\\n2: 测试\\n\",\n\t\t},\n\t\t{\n\t\t\tname:     \"long unicode line\",\n\t\t\tinput:    strings.Repeat(\"你\", 1500),\n\t\t\texpected: fmt.Sprintf(\"1: %s\\n\", strings.Repeat(\"你\", 1000)),\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := formatToolMessage(tt.input)\n\t\t\tif result != tt.expected {\n\t\t\t\tt.Errorf(\"formatToolMessage() = %q, want %q\", result, tt.expected)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestConcatString(t *testing.T) {\n\ttests := []struct {\n\t\tname        string\n\t\tchunks      []string\n\t\texpected    string\n\t\texpectError bool\n\t}{\n\t\t{\n\t\t\tname:     \"single chunk\",\n\t\t\tchunks:   []string{\"hello\"},\n\t\t\texpected: \"hello\",\n\t\t},\n\t\t{\n\t\t\tname:     \"multiple chunks\",\n\t\t\tchunks:   []string{\"hello\", \" \", \"world\"},\n\t\t\texpected: \"hello world\",\n\t\t},\n\t\t{\n\t\t\tname:     \"empty chunks\",\n\t\t\tchunks:   []string{\"\", \"\", \"\"},\n\t\t\texpected: \"\",\n\t\t},\n\t\t{\n\t\t\tname:     \"mixed chunks\",\n\t\t\tchunks:   []string{\"a\", \"\", \"b\", \"c\"},\n\t\t\texpected: \"abc\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tsr := schema.StreamReaderFromArray(tt.chunks)\n\t\t\tresult, err := concatString(sr)\n\n\t\t\tif tt.expectError {\n\t\t\t\tif err == nil {\n\t\t\t\t\tt.Error(\"expected error, got nil\")\n\t\t\t\t}\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\tif err != nil {\n\t\t\t\tt.Fatalf(\"unexpected error: %v\", err)\n\t\t\t}\n\n\t\t\tif result != tt.expected {\n\t\t\t\tt.Errorf(\"concatString() = %q, want %q\", result, tt.expected)\n\t\t\t}\n\t\t})\n\t}\n\n\t// Test nil stream\n\tt.Run(\"nil stream\", func(t *testing.T) {\n\t\t_, err := concatString(nil)\n\t\tif err == nil {\n\t\t\tt.Error(\"expected error for nil stream, got nil\")\n\t\t}\n\t\tif !strings.Contains(err.Error(), \"stream is nil\") {\n\t\t\tt.Errorf(\"expected 'stream is nil' error, got %v\", err)\n\t\t}\n\t})\n}\n\nfunc TestToolResultOffloading_BackendWriteError(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a backend that fails on write\n\tbackend := &failingBackend{\n\t\twriteErr: errors.New(\"write failed\"),\n\t}\n\n\tconfig := &toolResultOffloadingConfig{\n\t\tBackend:    backend,\n\t\tTokenLimit: 10,\n\t}\n\n\tmiddleware := newToolResultOffloading(ctx, config)\n\n\tlargeResult := strings.Repeat(\"Large content \", 100)\n\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\treturn &compose.ToolOutput{Result: largeResult}, nil\n\t}\n\n\twrappedEndpoint := middleware.Invokable(mockEndpoint)\n\n\tinput := &compose.ToolInput{\n\t\tName:   \"test_tool\",\n\t\tCallID: \"call_write_error\",\n\t}\n\t_, err := wrappedEndpoint(ctx, input)\n\n\tif err == nil {\n\t\tt.Fatal(\"expected error, got nil\")\n\t}\n\n\tif !strings.Contains(err.Error(), \"write failed\") {\n\t\tt.Errorf(\"expected 'write failed' error, got %v\", err)\n\t}\n}\n\n// failingBackend is a mock backend that can be configured to fail\ntype failingBackend struct {\n\twriteErr error\n}\n\nfunc (f *failingBackend) Write(context.Context, *filesystem.WriteRequest) error {\n\tif f.writeErr != nil {\n\t\treturn f.writeErr\n\t}\n\treturn nil\n}\n"
  },
  {
    "path": "adk/middlewares/reduction/internal/tool_result.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Backend defines the interface provided by the user to implement file storage.\n// It is used to save the content of large tool results to a persistent storage.\ntype Backend interface {\n\tWrite(context.Context, *filesystem.WriteRequest) error\n}\n\n// ToolResultConfig configures the tool result reduction middleware.\ntype ToolResultConfig struct {\n\t// ClearingTokenThreshold is the threshold for the total token count of all tool results.\n\t// When the sum of all tool result tokens exceeds this threshold, old tool results\n\t// (outside the KeepRecentTokens range) will be replaced with a placeholder.\n\t// Token estimation uses a simple heuristic: character count / 4.\n\t// optional, 20000 by default\n\tClearingTokenThreshold int\n\n\t// KeepRecentTokens is the token budget for recent messages to keep intact.\n\t// Messages within this token budget from the end will not have their tool results cleared,\n\t// even if the total tool result tokens exceed the threshold.\n\t// optional, 40000 by default\n\tKeepRecentTokens int\n\n\t// ClearToolResultPlaceholder is the text to replace old tool results with.\n\t// optional, \"[Old tool result content cleared]\" by default\n\tClearToolResultPlaceholder string\n\n\t// TokenCounter is a custom function to estimate token count for a message.\n\t// optional, uses the default counter (character count / 4) if nil\n\tTokenCounter func(msg *schema.Message) int\n\n\t// ExcludeTools is a list of tool names whose results should never be cleared.\n\t// optional\n\tExcludeTools []string\n\n\t// Backend is the storage backend for offloaded tool results.\n\t// required\n\tBackend Backend\n\n\t// OffloadingTokenLimit is the token threshold for a single tool result to trigger offloading.\n\t// When a single tool result exceeds OffloadingTokenLimit * 4 characters, it will be\n\t// offloaded to the filesystem.\n\t// optional, 20000 by default\n\tOffloadingTokenLimit int\n\n\t// ReadFileToolName is the name of the tool that LLM should use to read offloaded content.\n\t// This name will be included in the summary message sent to the LLM.\n\t// optional, \"read_file\" by default\n\t//\n\t// NOTE: If you are using the filesystem middleware, the read_file tool name\n\t// is exactly \"read_file\", which matches the default value.\n\tReadFileToolName string\n\n\t// PathGenerator generates the write path for offloaded results.\n\t// optional, \"/large_tool_result/{ToolCallID}\" by default\n\tPathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error)\n}\n\n// NewToolResultMiddleware creates a tool result reduction middleware.\n// This middleware combines two strategies to manage tool result tokens:\n//\n//  1. Clearing: Replaces old tool results with a placeholder when the total\n//     tool result tokens exceed the threshold, while protecting recent messages.\n//\n//  2. Offloading: Writes large individual tool results to the filesystem and\n//     returns a summary message guiding the LLM to read the full content.\n//\n// NOTE: If you are using the filesystem middleware (github.com/cloudwego/eino/adk/middlewares/filesystem),\n// this functionality is already included by default. Set Config.WithoutLargeToolResultOffloading = true\n// in the filesystem middleware if you want to use this middleware separately instead.\n//\n// NOTE: This middleware only handles offloading results to the filesystem.\n// You MUST also provide a read_file tool to your agent, otherwise the agent\n// will not be able to read the offloaded content. You can either:\n//   - Use the filesystem middleware (github.com/cloudwego/eino/adk/middlewares/filesystem)\n//     which provides the read_file tool automatically, OR\n//   - Implement your own read_file tool that reads from the same Backend\nfunc NewToolResultMiddleware(ctx context.Context, cfg *ToolResultConfig) (adk.AgentMiddleware, error) {\n\tbc := newClearToolResult(ctx, &ClearToolResultConfig{\n\t\tToolResultTokenThreshold:   cfg.ClearingTokenThreshold,\n\t\tKeepRecentTokens:           cfg.KeepRecentTokens,\n\t\tClearToolResultPlaceholder: cfg.ClearToolResultPlaceholder,\n\t\tTokenCounter:               cfg.TokenCounter,\n\t\tExcludeTools:               cfg.ExcludeTools,\n\t})\n\ttm := newToolResultOffloading(ctx, &toolResultOffloadingConfig{\n\t\tBackend:          cfg.Backend,\n\t\tReadFileToolName: cfg.ReadFileToolName,\n\t\tTokenLimit:       cfg.OffloadingTokenLimit,\n\t\tPathGenerator:    cfg.PathGenerator,\n\t})\n\treturn adk.AgentMiddleware{\n\t\tBeforeChatModel: bc,\n\t\tWrapToolCall:    tm,\n\t}, nil\n}\n"
  },
  {
    "path": "adk/middlewares/reduction/legacy.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage reduction\n\nimport \"github.com/cloudwego/eino/adk/middlewares/reduction/internal\"\n\n// Package reduction provides historical compatibility exports for reduction middleware APIs.\n//\n// DEPRECATED: All top-level exports in this file are maintained exclusively for backward compatibility.\n// New reduction middleware implementations are now developed and maintained in this package.\n// It is STRONGLY RECOMMENDED that new code directly use the New instead.\n//\n// Existing code relying on these exports will continue to work indefinitely,\n// but no new features or bug fixes will be backported to this compatibility layer.\n\ntype (\n\tClearToolResultConfig = internal.ClearToolResultConfig\n\tToolResultConfig      = internal.ToolResultConfig\n\tBackend               = internal.Backend\n)\n\nvar (\n\t// NewClearToolResult creates a new middleware that clears old tool results\n\t// based on token thresholds while protecting recent messages.\n\t//\n\t// Deprecated: Use New instead, which provides a more comprehensive tool result reduction\n\t// middleware with both truncation and clearing strategies. New returns a ChatModelAgentMiddleware\n\t// for better context propagation through wrapper methods.\n\tNewClearToolResult = internal.NewClearToolResult\n\n\t// NewToolResultMiddleware creates a tool result reduction middleware.\n\t// This middleware combines two strategies to manage tool result tokens:\n\t//\n\t//  1. Clearing: Replaces old tool results with a placeholder when the total\n\t//     tool result tokens exceed the threshold, while protecting recent messages.\n\t//\n\t//  2. Offloading: Writes large individual tool results to the filesystem and\n\t//     returns a summary message guiding the LLM to read the full content.\n\t//\n\t// NOTE: If you are using the filesystem middleware (github.com/cloudwego/eino/adk/middlewares/filesystem),\n\t// this functionality is already included by default. Set Config.WithoutLargeToolResultOffloading = true\n\t// in the filesystem middleware if you want to use this middleware separately instead.\n\t//\n\t// NOTE: This middleware only handles offloading results to the filesystem.\n\t// You MUST also provide a read_file tool to your agent, otherwise the agent\n\t// will not be able to read the offloaded content. You can either:\n\t//   - Use the filesystem middleware (github.com/cloudwego/eino/adk/middlewares/filesystem)\n\t//     which provides the read_file tool automatically, OR\n\t//   - Implement your own read_file tool that reads from the same Backend\n\t//\n\t// Deprecated: Use New instead, which provides a more comprehensive tool result reduction\n\t// middleware with both truncation and clearing strategies, per-tool configuration support,\n\t// and returns a ChatModelAgentMiddleware for better context propagation through wrapper methods.\n\tNewToolResultMiddleware = internal.NewToolResultMiddleware\n)\n"
  },
  {
    "path": "adk/middlewares/reduction/reduction.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage reduction\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"io\"\n\t\"path/filepath\"\n\t\"strings\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/google/uuid\"\n\t\"github.com/slongfield/pyfmt\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Config is the configuration for tool reduction middleware.\n// This middleware manages tool outputs in two phases to optimize context usage:\n//\n//  1. Truncation Phase:\n//     Triggered immediately after a tool execution completes.\n//     If the tool output length exceeds MaxLengthForTrunc, the full content is saved\n//     to the configured Backend, and the tool output is replaced with a truncated notice.\n//     This prevents immediate context overflow from a single large tool output.\n//\n//  2. Clear Phase:\n//     Triggered before sending messages to the model (in BeforeModelRewriteState).\n//     If the total token count exceeds MaxTokensForClear, the middleware iterates through\n//     historical messages. Based on RootDir and ClearRetentionSuffixLimit, it offloads tool call arguments and results\n//     to the Backend to reduce token usage, keeping the conversation within limits while retaining access to the\n//     important information. After all, ClearPostProcess will be called, which you could save or notify current state.\ntype Config struct {\n\t// Backend is the storage backend where truncated content will be saved.\n\t// Required.\n\tBackend Backend\n\n\t// SkipTruncation skip truncating.\n\tSkipTruncation bool\n\n\t// SkipClear skip clearing.\n\tSkipClear bool\n\n\t// ReadFileToolName is tool name used to retrieve from file.\n\t// After offloading content to file, you should give agent the same tool to retrieve content.\n\t// Required. Default is \"read_file\".\n\tReadFileToolName string\n\n\t// RootDir root dir to save truncated/cleared content.\n\t// Required.\n\t// Default is /tmp, truncated content saves to /tmp/trunc/{tool_call_id}, cleared content saves to  /tmp/clear/{tool_call_id}\n\tRootDir string\n\n\t// MaxLengthForTrunc is the maximum allowed length of the tool output.\n\t// If the output exceeds this length, it will be truncated.\n\t// Required. Default is 50000.\n\tMaxLengthForTrunc int\n\n\t// TokenCounter is used to count the number of tokens in the conversation messages.\n\t// It is used to determine when to trigger clearing based on token usage, and token usage after clearing.\n\t// Required.\n\tTokenCounter func(ctx context.Context, msg []adk.Message, tools []*schema.ToolInfo) (int64, error)\n\n\t// MaxTokensForClear is the maximum number of tokens allowed in the conversation before clearing is attempted.\n\t// Required. Default is 30000.\n\tMaxTokensForClear int64\n\n\t// ClearRetentionSuffixLimit is the number of most recent messages to retain without clearing.\n\t// This ensures the model has some immediate context.\n\t// Optional. Default is 1.\n\tClearRetentionSuffixLimit int\n\n\t// ClearPostProcess is clear post process handler.\n\t// Optional.\n\tClearPostProcess func(ctx context.Context, state *adk.ChatModelAgentState) context.Context\n\n\t// ToolConfig is the specific configuration that applies to tools by name.\n\t// This configuration takes precedence over GeneralConfig for the specified tools.\n\t// Optional.\n\tToolConfig map[string]*ToolReductionConfig\n}\n\ntype ToolReductionConfig struct {\n\t// Backend is the storage backend where truncated content will be saved.\n\t// Required.\n\tBackend Backend\n\n\t// SkipTruncation skip truncation for this tool.\n\tSkipTruncation bool\n\n\t// TruncHandler is used to process tool call results during truncation.\n\t// Optional. Default using defaultTruncHandler when SkipTruncation is false but TruncHandler is nil.\n\tTruncHandler func(ctx context.Context, detail *ToolDetail) (*TruncResult, error)\n\n\t// SkipClear skip clear for this tool.\n\tSkipClear bool\n\n\t// ClearHandler is used to process tool call arguments and results during clearing.\n\t// Optional. Default using defaultClearHandler when SkipClear is false but ClearHandler is nil.\n\tClearHandler func(ctx context.Context, detail *ToolDetail) (*ClearResult, error)\n}\n\ntype ToolDetail struct {\n\t// ToolContext provides metadata about the tool call (e.g., tool name, call ID).\n\tToolContext *adk.ToolContext\n\n\t// ToolArgument contains the arguments passed to the tool.\n\tToolArgument *schema.ToolArgument\n\n\t// ToolResult contains the output returned by the tool.\n\tToolResult *schema.ToolResult\n}\n\ntype TruncResult struct {\n\t// NeedTrunc indicates whether the tool result should be truncated.\n\tNeedTrunc bool\n\n\t// ToolResult contains the result returned by the tool after trunc\n\t// Required when NeedTrunc is true.\n\tToolResult *schema.ToolResult\n\n\t// NeedOffload indicates whether the tool result should be offloaded.\n\tNeedOffload bool\n\n\t// OffloadFilePath is the path where the offloaded content should be stored.\n\t// This path is typically relative to the backend's root.\n\t// Required when NeedOffload is true.\n\tOffloadFilePath string\n\n\t// OffloadContent is the actual content to be written to the storage backend.\n\t// Required when NeedOffload is true.\n\tOffloadContent string\n}\n\n// ClearResult contains the result of the Handler's decision.\ntype ClearResult struct {\n\t// NeedClear indicates whether the tool argument and result should be cleared.\n\tNeedClear bool\n\n\t// ToolArgument contains the arguments passed to the tool after clear.\n\t// Required when NeedClear is true.\n\tToolArgument *schema.ToolArgument\n\n\t// ToolResult contains the output returned by the tool after clear.\n\t// Required when NeedClear is true\n\tToolResult *schema.ToolResult\n\n\t// NeedOffload indicates whether the tool argument and result should be offloaded.\n\tNeedOffload bool\n\n\t// OffloadFilePath is the path where the offloaded content should be stored.\n\t// This path is typically relative to the backend's root.\n\t// Required when NeedOffload is true.\n\tOffloadFilePath string\n\n\t// OffloadContent is the actual content to be written to the storage backend.\n\t// Required when NeedOffload is true.\n\tOffloadContent string\n}\n\nfunc (t *Config) copyAndFillDefaults() (*Config, error) {\n\tcfg := &Config{\n\t\tBackend:                   t.Backend,\n\t\tSkipTruncation:            t.SkipTruncation,\n\t\tSkipClear:                 t.SkipClear,\n\t\tReadFileToolName:          t.ReadFileToolName,\n\t\tRootDir:                   t.RootDir,\n\t\tMaxLengthForTrunc:         t.MaxLengthForTrunc,\n\t\tTokenCounter:              t.TokenCounter,\n\t\tMaxTokensForClear:         t.MaxTokensForClear,\n\t\tClearRetentionSuffixLimit: t.ClearRetentionSuffixLimit,\n\t\tClearPostProcess:          t.ClearPostProcess,\n\t}\n\tif cfg.TokenCounter == nil {\n\t\tcfg.TokenCounter = defaultTokenCounter\n\t}\n\tif cfg.ClearRetentionSuffixLimit == 0 {\n\t\tcfg.ClearRetentionSuffixLimit = 1\n\t}\n\tif cfg.ReadFileToolName == \"\" {\n\t\tcfg.ReadFileToolName = \"read_file\"\n\t}\n\tif cfg.RootDir == \"\" {\n\t\tcfg.RootDir = \"/tmp\"\n\t}\n\tif cfg.MaxLengthForTrunc == 0 {\n\t\tcfg.MaxLengthForTrunc = 50000\n\t}\n\tif t.ToolConfig != nil {\n\t\tcfg.ToolConfig = make(map[string]*ToolReductionConfig, len(t.ToolConfig))\n\t\tfor toolName, trc := range t.ToolConfig {\n\t\t\tcpConfig := &ToolReductionConfig{\n\t\t\t\tBackend:        trc.Backend,\n\t\t\t\tSkipTruncation: trc.SkipTruncation,\n\t\t\t\tSkipClear:      trc.SkipClear,\n\t\t\t\tTruncHandler:   trc.TruncHandler,\n\t\t\t\tClearHandler:   trc.ClearHandler,\n\t\t\t}\n\t\t\tcfg.ToolConfig[toolName] = cpConfig\n\t\t}\n\t}\n\n\treturn cfg, nil\n}\n\n// New creates tool reduction middleware from config\nfunc New(_ context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) {\n\tvar err error\n\tif config == nil {\n\t\treturn nil, fmt.Errorf(\"config must not be nil\")\n\t}\n\tif config.Backend == nil && !config.SkipTruncation {\n\t\treturn nil, fmt.Errorf(\"backend must be set when not skipping truncation\")\n\t}\n\n\tconfig, err = config.copyAndFillDefaults()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdefaultReductionConfig := &ToolReductionConfig{\n\t\tBackend:        config.Backend,\n\t\tSkipTruncation: config.SkipTruncation,\n\t\tSkipClear:      config.SkipClear,\n\t}\n\tif !defaultReductionConfig.SkipTruncation {\n\t\tdefaultReductionConfig.TruncHandler = defaultTruncHandler(config.RootDir, config.MaxLengthForTrunc)\n\t}\n\tif !defaultReductionConfig.SkipClear {\n\t\tdefaultReductionConfig.ClearHandler = defaultClearHandler(config.RootDir, config.Backend != nil, config.ReadFileToolName)\n\t}\n\n\treturn &toolReductionMiddleware{\n\t\tconfig:        config,\n\t\tdefaultConfig: defaultReductionConfig,\n\t}, nil\n}\n\ntype toolReductionMiddleware struct {\n\tadk.BaseChatModelAgentMiddleware\n\n\tconfig        *Config\n\tdefaultConfig *ToolReductionConfig\n}\n\nfunc (t *toolReductionMiddleware) getToolConfig(toolName string, sc scene) *ToolReductionConfig {\n\tif t.config.ToolConfig != nil {\n\t\tif cfg, ok := t.config.ToolConfig[toolName]; ok {\n\t\t\tif (sc == sceneTruncation && !cfg.SkipTruncation && cfg.TruncHandler == nil) ||\n\t\t\t\t(sc == sceneClear && !cfg.SkipClear && cfg.ClearHandler == nil) {\n\t\t\t\treturn t.defaultConfig\n\t\t\t}\n\t\t\treturn cfg\n\t\t}\n\t}\n\treturn t.defaultConfig\n}\n\nfunc (t *toolReductionMiddleware) WrapInvokableToolCall(_ context.Context, endpoint adk.InvokableToolCallEndpoint, tCtx *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) {\n\tcfg := t.getToolConfig(tCtx.Name, sceneTruncation)\n\tif cfg == nil || cfg.TruncHandler == nil {\n\t\treturn endpoint, nil\n\t}\n\n\treturn func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\t\toutput, err := endpoint(ctx, argumentsInJSON, opts...)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tdetail := &ToolDetail{\n\t\t\tToolContext: tCtx,\n\t\t\tToolArgument: &schema.ToolArgument{\n\t\t\t\tText: argumentsInJSON,\n\t\t\t},\n\t\t\tToolResult: &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t{Type: schema.ToolPartTypeText, Text: output},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\ttruncResult, err := cfg.TruncHandler(ctx, detail)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tif !truncResult.NeedTrunc {\n\t\t\treturn output, nil\n\t\t}\n\t\tif truncResult.NeedOffload {\n\t\t\tif cfg.Backend == nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"truncation: no backend for offload\")\n\t\t\t}\n\t\t\tif err = cfg.Backend.Write(ctx, &filesystem.WriteRequest{\n\t\t\t\tFilePath: truncResult.OffloadFilePath,\n\t\t\t\tContent:  truncResult.OffloadContent,\n\t\t\t}); err != nil {\n\t\t\t\treturn \"\", err\n\t\t\t}\n\t\t}\n\t\treturn truncResult.ToolResult.Parts[0].Text, nil\n\t}, nil\n}\n\nfunc (t *toolReductionMiddleware) WrapStreamableToolCall(_ context.Context, endpoint adk.StreamableToolCallEndpoint, tCtx *adk.ToolContext) (adk.StreamableToolCallEndpoint, error) {\n\tcfg := t.getToolConfig(tCtx.Name, sceneTruncation)\n\tif cfg == nil || cfg.TruncHandler == nil {\n\t\treturn endpoint, nil\n\t}\n\n\treturn func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\t\toutput, err := endpoint(ctx, argumentsInJSON, opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tvar chunks []string\n\t\treaders := output.Copy(2)\n\t\toutput = readers[0]\n\t\torigResp := readers[1]\n\t\tdefer output.Close()\n\n\t\tfor {\n\t\t\tvar recvErr error\n\t\t\tchunk, recvErr := output.Recv()\n\t\t\tif recvErr != nil {\n\t\t\t\tif recvErr != io.EOF {\n\t\t\t\t\treturn origResp, nil\n\t\t\t\t}\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tchunks = append(chunks, chunk)\n\t\t}\n\n\t\tresult := strings.Join(chunks, \"\")\n\t\tdetail := &ToolDetail{\n\t\t\tToolContext: tCtx,\n\t\t\tToolArgument: &schema.ToolArgument{\n\t\t\t\tText: argumentsInJSON,\n\t\t\t},\n\t\t\tToolResult: &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t{Type: schema.ToolPartTypeText, Text: result},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\ttruncResult, err := cfg.TruncHandler(ctx, detail)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tif !truncResult.NeedTrunc {\n\t\t\treturn origResp, nil\n\t\t}\n\t\torigResp.Close() // close err resp when not using it\n\n\t\tif truncResult.NeedOffload {\n\t\t\tif cfg.Backend == nil {\n\t\t\t\treturn nil, fmt.Errorf(\"truncation: no backend for offload\")\n\t\t\t}\n\t\t\tif err = cfg.Backend.Write(ctx, &filesystem.WriteRequest{\n\t\t\t\tFilePath: truncResult.OffloadFilePath,\n\t\t\t\tContent:  truncResult.OffloadContent,\n\t\t\t}); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t\treturn schema.StreamReaderFromArray([]string{truncResult.ToolResult.Parts[0].Text}), nil\n\t}, nil\n}\n\nfunc (t *toolReductionMiddleware) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (\n\tcontext.Context, *adk.ChatModelAgentState, error) {\n\n\tvar (\n\t\terr             error\n\t\testimatedTokens int64\n\t)\n\n\t// init msg tokens\n\testimatedTokens, err = t.config.TokenCounter(ctx, state.Messages, mc.Tools)\n\tif err != nil {\n\t\treturn ctx, state, err\n\t}\n\n\tif estimatedTokens < t.config.MaxTokensForClear {\n\t\treturn ctx, state, nil\n\t}\n\n\t// calc range\n\tvar (\n\t\tstart = 0\n\t\tend   = len(state.Messages)\n\t)\n\tfor ; start < len(state.Messages); start++ {\n\t\tmsg := state.Messages[start]\n\t\tif msg.Role == schema.Assistant && !getMsgOffloadedFlag(msg) {\n\t\t\tbreak\n\t\t}\n\t}\n\tretention := t.config.ClearRetentionSuffixLimit\n\tfor ; retention > 0 && end > 0; end-- {\n\t\tmsg := state.Messages[end-1]\n\t\tif msg.Role == schema.Assistant && len(msg.ToolCalls) > 0 {\n\t\t\tretention--\n\t\t\tif retention == 0 {\n\t\t\t\tend--\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\tif start >= end {\n\t\treturn ctx, state, nil\n\t}\n\n\t// recursively handle\n\ttcMsgIndex := start\n\n\tfor tcMsgIndex < end {\n\t\ttcMsg := state.Messages[tcMsgIndex]\n\t\tif tcMsg.Role == schema.Assistant && len(tcMsg.ToolCalls) > 0 {\n\t\t\ttrMsgEnd := tcMsgIndex + 1 + len(tcMsg.ToolCalls)\n\t\t\tif trMsgEnd > len(state.Messages) {\n\t\t\t\ttrMsgEnd = len(state.Messages)\n\t\t\t}\n\n\t\t\tj := tcMsgIndex\n\t\t\tfor tcIndex, toolCall := range tcMsg.ToolCalls {\n\t\t\t\tj++\n\t\t\t\tif j >= end {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tresultMsg := state.Messages[j]\n\t\t\t\tif resultMsg.Role != schema.Tool { // unexpected\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tcfg := t.getToolConfig(toolCall.Function.Name, sceneClear)\n\t\t\t\tif cfg == nil || cfg.ClearHandler == nil {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\n\t\t\t\ttoolResult, fromContent, toolResultErr := toolResultFromMessage(resultMsg)\n\t\t\t\tif toolResultErr != nil {\n\t\t\t\t\treturn ctx, state, toolResultErr\n\t\t\t\t}\n\n\t\t\t\ttd := &ToolDetail{\n\t\t\t\t\tToolContext: &adk.ToolContext{\n\t\t\t\t\t\tName:   toolCall.Function.Name,\n\t\t\t\t\t\tCallID: toolCall.ID,\n\t\t\t\t\t},\n\t\t\t\t\tToolArgument: &schema.ToolArgument{\n\t\t\t\t\t\tText: toolCall.Function.Arguments,\n\t\t\t\t\t},\n\t\t\t\t\tToolResult: toolResult,\n\t\t\t\t}\n\n\t\t\t\toffloadInfo, offloadErr := cfg.ClearHandler(ctx, td)\n\t\t\t\tif offloadErr != nil {\n\t\t\t\t\treturn ctx, state, offloadErr\n\t\t\t\t}\n\t\t\t\tif !offloadInfo.NeedClear {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\t\t\t\tif offloadInfo.NeedOffload {\n\t\t\t\t\tif cfg.Backend == nil {\n\t\t\t\t\t\treturn ctx, state, fmt.Errorf(\"clear: no backend for offload\")\n\t\t\t\t\t}\n\t\t\t\t\twriteErr := cfg.Backend.Write(ctx, &filesystem.WriteRequest{\n\t\t\t\t\t\tFilePath: offloadInfo.OffloadFilePath,\n\t\t\t\t\t\tContent:  offloadInfo.OffloadContent,\n\t\t\t\t\t})\n\t\t\t\t\tif writeErr != nil {\n\t\t\t\t\t\treturn ctx, state, writeErr\n\t\t\t\t\t}\n\t\t\t\t}\n\n\t\t\t\ttcMsg.ToolCalls[tcIndex].Function.Arguments = offloadInfo.ToolArgument.Text\n\t\t\t\tif fromContent {\n\t\t\t\t\tif len(offloadInfo.ToolResult.Parts) > 0 {\n\t\t\t\t\t\tresultMsg.Content = offloadInfo.ToolResult.Parts[0].Text\n\t\t\t\t\t}\n\t\t\t\t} else {\n\t\t\t\t\tvar convErr error\n\t\t\t\t\tresultMsg.UserInputMultiContent, convErr = offloadInfo.ToolResult.ToMessageInputParts()\n\t\t\t\t\tif convErr != nil {\n\t\t\t\t\t\treturn ctx, state, convErr\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\n\t\t\t// set dedup flag\n\t\t\tsetMsgOffloadedFlag(tcMsg)\n\t\t}\n\t\ttcMsgIndex++\n\t}\n\n\tif t.config.ClearPostProcess != nil {\n\t\tctx = t.config.ClearPostProcess(ctx, state)\n\t}\n\n\treturn ctx, state, nil\n}\n\n// defaultTokenCounter estimates tokens, which treats one token as ~4 characters of text for common English text.\n// github.com/tiktoken-go/tokenizer is highly recommended to replace it.\nfunc defaultTokenCounter(_ context.Context, msgs []*schema.Message, tools []*schema.ToolInfo) (int64, error) {\n\tvar tokens int64\n\tfor _, msg := range msgs {\n\t\tif msg == nil {\n\t\t\tcontinue\n\t\t}\n\t\tif cached, ok := getMsgCachedToken(msg); ok {\n\t\t\ttokens += cached\n\t\t\tcontinue\n\t\t}\n\n\t\tvar sb strings.Builder\n\t\tsb.WriteString(string(msg.Role))\n\t\tsb.WriteString(\"\\n\")\n\t\tsb.WriteString(msg.ReasoningContent)\n\t\tsb.WriteString(\"\\n\")\n\t\tsb.WriteString(msg.Content)\n\t\tsb.WriteString(\"\\n\")\n\t\tif msg.Role == schema.Assistant && len(msg.ToolCalls) > 0 {\n\t\t\tfor _, tc := range msg.ToolCalls {\n\t\t\t\tsb.WriteString(tc.Function.Name)\n\t\t\t\tsb.WriteString(\"\\n\")\n\t\t\t\tsb.WriteString(tc.Function.Arguments)\n\t\t\t}\n\t\t}\n\n\t\tn := int64(len(sb.String()) / 4)\n\t\tsetMsgCachedToken(msg, n)\n\t\ttokens += n\n\t}\n\n\tfor _, tl := range tools {\n\t\ttl_ := *tl\n\t\ttl_.Extra = nil\n\t\ttext, err := sonic.MarshalString(tl_)\n\t\tif err != nil {\n\t\t\treturn 0, fmt.Errorf(\"failed to marshal tool info: %w\", err)\n\t\t}\n\n\t\ttokens += int64(len(text) / 4)\n\t}\n\n\treturn tokens, nil\n}\n\nfunc defaultTruncHandler(rootDir string, truncMaxLength int) func(ctx context.Context, detail *ToolDetail) (truncResult *TruncResult, err error) {\n\treturn func(ctx context.Context, detail *ToolDetail) (offloadInfo *TruncResult, err error) {\n\t\tresultText := detail.ToolResult.Parts[0].Text\n\t\tif len(resultText) <= truncMaxLength {\n\t\t\treturn &TruncResult{NeedTrunc: false}, nil\n\t\t}\n\n\t\tfilePath := filepath.Join(rootDir, \"trunc\", detail.ToolContext.CallID)\n\t\tpreviewSize := truncMaxLength / 2\n\t\ttruncNotify, err := pyfmt.Fmt(getTruncFmt(), map[string]any{\n\t\t\t\"original_size\": len(resultText),\n\t\t\t\"file_path\":     filePath,\n\t\t\t\"preview_size\":  previewSize,\n\t\t\t\"preview_first\": resultText[:previewSize],\n\t\t\t\"preview_last\":  resultText[len(resultText)-previewSize:],\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn &TruncResult{\n\t\t\tToolResult: &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t{Type: schema.ToolPartTypeText, Text: resultText[:truncMaxLength] + truncNotify},\n\t\t\t\t},\n\t\t\t},\n\t\t\tNeedTrunc:       true,\n\t\t\tNeedOffload:     true,\n\t\t\tOffloadFilePath: filePath,\n\t\t\tOffloadContent:  resultText,\n\t\t}, nil\n\t}\n}\n\nfunc defaultClearHandler(rootDir string, needOffload bool, readFileToolName string) func(ctx context.Context, detail *ToolDetail) (*ClearResult, error) {\n\treturn func(ctx context.Context, detail *ToolDetail) (clearResult *ClearResult, err error) {\n\t\tif len(detail.ToolResult.Parts) == 0 {\n\t\t\treturn &ClearResult{NeedClear: false}, nil\n\t\t}\n\t\tfor _, part := range detail.ToolResult.Parts {\n\t\t\tif part.Type != schema.ToolPartTypeText {\n\t\t\t\t// brutal judge\n\t\t\t\treturn nil, fmt.Errorf(\"default offload currently not support multimodal content type=%v\", part.Type)\n\t\t\t}\n\t\t}\n\n\t\tfileName := detail.ToolContext.CallID\n\t\tif fileName == \"\" {\n\t\t\tfileName = uuid.NewString()\n\t\t}\n\n\t\tvar nResult string\n\t\tif needOffload {\n\t\t\tfilePath := filepath.Join(rootDir, \"clear\", fileName)\n\t\t\tnResult, err = pyfmt.Fmt(getClearWithOffloadingFmt(), map[string]any{\n\t\t\t\t\"file_path\":      filePath,\n\t\t\t\t\"read_tool_name\": readFileToolName,\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tclearResult = &ClearResult{\n\t\t\t\tToolArgument:    detail.ToolArgument,\n\t\t\t\tNeedClear:       true,\n\t\t\t\tNeedOffload:     true,\n\t\t\t\tOffloadFilePath: filePath,\n\t\t\t\tOffloadContent:  detail.ToolResult.Parts[0].Text,\n\t\t\t}\n\t\t} else {\n\t\t\tnResult = getClearWithoutOffloadingFmt()\n\t\t\tclearResult = &ClearResult{\n\t\t\t\tToolArgument: detail.ToolArgument,\n\t\t\t\tNeedClear:    true,\n\t\t\t\tNeedOffload:  false,\n\t\t\t}\n\t\t}\n\n\t\tclearResult.ToolResult = &schema.ToolResult{\n\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t{Type: schema.ToolPartTypeText, Text: nResult},\n\t\t\t},\n\t\t}\n\n\t\treturn clearResult, nil\n\t}\n}\n\nfunc getMsgOffloadedFlag(msg *schema.Message) (offloaded bool) {\n\tif msg.Extra == nil {\n\t\treturn false\n\t}\n\tv, ok := msg.Extra[msgReducedFlag].(bool)\n\tif !ok {\n\t\treturn false\n\t}\n\treturn v\n}\n\nfunc setMsgOffloadedFlag(msg *schema.Message) {\n\tif msg.Extra == nil {\n\t\tmsg.Extra = make(map[string]any)\n\t}\n\tmsg.Extra[msgReducedFlag] = true\n}\n\nfunc getMsgCachedToken(msg *schema.Message) (int64, bool) {\n\tif msg.Extra == nil {\n\t\treturn 0, false\n\t}\n\ttokens, ok := msg.Extra[msgReducedTokens].(int64)\n\treturn tokens, ok\n}\n\nfunc setMsgCachedToken(msg *schema.Message, tokens int64) {\n\tif msg.Extra == nil {\n\t\tmsg.Extra = make(map[string]any)\n\t}\n\tmsg.Extra[msgReducedTokens] = tokens\n}\n\nfunc toolResultFromMessage(msg *schema.Message) (result *schema.ToolResult, fromContent bool, err error) {\n\tif msg.Role != schema.Tool {\n\t\treturn nil, false, fmt.Errorf(\"message role %s is not a tool\", msg.Role)\n\t}\n\tif msg.Content != \"\" {\n\t\treturn &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: msg.Content}}}, true, nil\n\t}\n\tresult = &schema.ToolResult{Parts: make([]schema.ToolOutputPart, 0, len(msg.UserInputMultiContent))}\n\tfor _, part := range msg.UserInputMultiContent {\n\t\ttop, convErr := convMessageInputPartToToolOutputPart(part)\n\t\tif convErr != nil {\n\t\t\treturn nil, false, convErr\n\t\t}\n\t\tresult.Parts = append(result.Parts, top)\n\t}\n\treturn result, false, nil\n}\n\nfunc convMessageInputPartToToolOutputPart(msgPart schema.MessageInputPart) (schema.ToolOutputPart, error) {\n\tswitch msgPart.Type {\n\tcase schema.ChatMessagePartTypeText:\n\t\treturn schema.ToolOutputPart{\n\t\t\tType: schema.ToolPartTypeText,\n\t\t\tText: msgPart.Text,\n\t\t}, nil\n\tcase schema.ChatMessagePartTypeImageURL:\n\t\treturn schema.ToolOutputPart{\n\t\t\tType: schema.ToolPartTypeImage,\n\t\t\tImage: &schema.ToolOutputImage{\n\t\t\t\tMessagePartCommon: msgPart.Image.MessagePartCommon,\n\t\t\t},\n\t\t}, nil\n\tcase schema.ChatMessagePartTypeAudioURL:\n\t\treturn schema.ToolOutputPart{\n\t\t\tType: schema.ToolPartTypeAudio,\n\t\t\tAudio: &schema.ToolOutputAudio{\n\t\t\t\tMessagePartCommon: msgPart.Audio.MessagePartCommon,\n\t\t\t},\n\t\t}, nil\n\tcase schema.ChatMessagePartTypeVideoURL:\n\t\treturn schema.ToolOutputPart{\n\t\t\tType: schema.ToolPartTypeVideo,\n\t\t\tVideo: &schema.ToolOutputVideo{\n\t\t\t\tMessagePartCommon: msgPart.Video.MessagePartCommon,\n\t\t\t},\n\t\t}, nil\n\tcase schema.ChatMessagePartTypeFileURL:\n\t\treturn schema.ToolOutputPart{\n\t\t\tType: schema.ToolPartTypeFile,\n\t\t\tFile: &schema.ToolOutputFile{\n\t\t\t\tMessagePartCommon: msgPart.File.MessagePartCommon,\n\t\t\t},\n\t\t}, nil\n\tdefault:\n\t\treturn schema.ToolOutputPart{}, fmt.Errorf(\"unknown msg part type: %v\", msgPart.Type)\n\t}\n}\n"
  },
  {
    "path": "adk/middlewares/reduction/reduction_test.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage reduction\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/components/tool/utils\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestReductionMiddlewareTrunc(t *testing.T) {\n\tctx := context.Background()\n\tit := mockInvokableTool()\n\tst := mockStreamableTool()\n\n\tt.Run(\"test invokable max length trunc\", func(t *testing.T) {\n\t\ttCtx := &adk.ToolContext{\n\t\t\tName:   \"mock_invokable_tool\",\n\t\t\tCallID: \"12345\",\n\t\t}\n\t\tbackend := filesystem.NewInMemoryBackend()\n\t\tconfig := &Config{\n\t\t\tBackend: backend,\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"mock_invokable_tool\": {\n\t\t\t\t\tBackend:        backend,\n\t\t\t\t\tSkipTruncation: false,\n\t\t\t\t\tTruncHandler:   defaultTruncHandler(\"/tmp\", 70),\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\texp := \"hello worldhello worldhello worldhello worldhello worldhello worldhell<persisted-output>\\nOutput too large (199). Full output saved to: /tmp/trunc/12345\\nPreview (first 35):\\nhello worldhello worldhello worldhe\\n\\nPreview (last 35):\\nldhello worldhello worldhello world\\n\\n</persisted-output>\"\n\n\t\tedp, err := mw.WrapInvokableToolCall(ctx, it.InvokableRun, tCtx)\n\t\tassert.NoError(t, err)\n\t\tresp, err := edp(ctx, `{\"value\":\"asd\"}`)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, exp, resp)\n\t\tcontent, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: \"/tmp/trunc/12345\"})\n\t\tassert.NoError(t, err)\n\t\texpOrigContent := `hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world\nhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world`\n\t\tassert.Equal(t, expOrigContent, content.Content)\n\t})\n\n\tt.Run(\"test streamable line and max length trunc\", func(t *testing.T) {\n\t\ttCtx := &adk.ToolContext{\n\t\t\tName:   \"mock_streamable_tool\",\n\t\t\tCallID: \"54321\",\n\t\t}\n\t\tbackend := filesystem.NewInMemoryBackend()\n\t\tconfig := &Config{\n\t\t\tSkipTruncation: true,\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"mock_streamable_tool\": {\n\t\t\t\t\tBackend:        backend,\n\t\t\t\t\tSkipTruncation: false,\n\t\t\t\t\tTruncHandler:   defaultTruncHandler(\"/tmp\", 70),\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\texp := \"hello worldhello worldhello worldhello worldhello worldhello worldhell<persisted-output>\\nOutput too large (199). Full output saved to: /tmp/trunc/54321\\nPreview (first 35):\\nhello worldhello worldhello worldhe\\n\\nPreview (last 35):\\nldhello worldhello worldhello world\\n\\n</persisted-output>\"\n\n\t\tedp, err := mw.WrapStreamableToolCall(ctx, st.StreamableRun, tCtx)\n\t\tassert.NoError(t, err)\n\t\tresp, err := edp(ctx, `{\"value\":\"asd\"}`)\n\t\tassert.NoError(t, err)\n\t\ts, err := resp.Recv()\n\t\tassert.NoError(t, err)\n\t\tresp.Close()\n\t\tassert.Equal(t, exp, s)\n\t\tcontent, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: \"/tmp/trunc/54321\"})\n\t\tassert.NoError(t, err)\n\t\texpOrigContent := `hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world\nhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world`\n\t\tassert.Equal(t, expOrigContent, content.Content)\n\t})\n}\n\nfunc TestReductionMiddlewareClear(t *testing.T) {\n\tctx := context.Background()\n\tit := mockInvokableTool()\n\tst := mockStreamableTool()\n\ttools := []tool.BaseTool{it, st}\n\tvar toolsInfo []*schema.ToolInfo\n\tfor _, bt := range tools {\n\t\tti, _ := bt.Info(ctx)\n\t\ttoolsInfo = append(toolsInfo, ti)\n\t}\n\ttype OffloadContent struct {\n\t\tArguments map[string]string `json:\"arguments\"`\n\t\tResult    string            `json:\"result\"`\n\t}\n\n\tt.Run(\"test default clear\", func(t *testing.T) {\n\t\tbackend := filesystem.NewInMemoryBackend()\n\t\tconfig := &Config{\n\t\t\tSkipTruncation:            true,\n\t\t\tTokenCounter:              defaultTokenCounter,\n\t\t\tMaxTokensForClear:         20,\n\t\t\tClearRetentionSuffixLimit: 0,\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"get_weather\": {\n\t\t\t\t\tBackend:      backend,\n\t\t\t\t\tSkipClear:    false,\n\t\t\t\t\tClearHandler: defaultClearHandler(\"/tmp\", true, \"read_file\"),\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\t_, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.SystemMessage(\"you are a helpful assistant\"),\n\t\t\t\tschema.UserMessage(\"If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C.\"),\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:       \"call_987654321\",\n\t\t\t\t\t\tType:     \"function\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t\tschema.ToolMessage(\"Sunny\", \"call_123456789\"),\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:       \"call_123456789\",\n\t\t\t\t\t\tType:     \"function\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t\tschema.ToolMessage(\"Sunny\", \"call_123456789\"),\n\t\t\t},\n\t\t}, &adk.ModelContext{\n\t\t\tTools: toolsInfo,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_987654321\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[2].ToolCalls)\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_123456789\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[4].ToolCalls)\n\t\tassert.Equal(t, \"<persisted-output>Tool result saved to: /tmp/clear/call_987654321\\nUse read_file to view</persisted-output>\", s.Messages[3].Content)\n\t\tfileContent, err := backend.Read(ctx, &filesystem.ReadRequest{\n\t\t\tFilePath: \"/tmp/clear/call_987654321\",\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfileContentStr := strings.TrimPrefix(strings.TrimSpace(fileContent.Content), \"1\\t\")\n\t\tassert.Equal(t, \"Sunny\", fileContentStr)\n\t})\n\n\tt.Run(\"test default clear without offloading\", func(t *testing.T) {\n\t\tconfig := &Config{\n\t\t\tSkipTruncation:            true,\n\t\t\tTokenCounter:              defaultTokenCounter,\n\t\t\tMaxTokensForClear:         20,\n\t\t\tClearRetentionSuffixLimit: 0,\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"get_weather\": {\n\t\t\t\t\tSkipClear:    false,\n\t\t\t\t\tClearHandler: defaultClearHandler(\"\", false, \"\"),\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\t_, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.SystemMessage(\"you are a helpful assistant\"),\n\t\t\t\tschema.UserMessage(\"If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C.\"),\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:       \"call_987654321\",\n\t\t\t\t\t\tType:     \"function\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t\tschema.ToolMessage(\"Sunny\", \"call_123456789\"),\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:       \"call_123456789\",\n\t\t\t\t\t\tType:     \"function\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t\tschema.ToolMessage(\"Sunny\", \"call_123456789\"),\n\t\t\t},\n\t\t}, &adk.ModelContext{\n\t\t\tTools: toolsInfo,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_987654321\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[2].ToolCalls)\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_123456789\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[4].ToolCalls)\n\t\tassert.Equal(t, \"[Old tool result content cleared]\", s.Messages[3].Content)\n\t})\n\n\tt.Run(\"test clear\", func(t *testing.T) {\n\t\tbackend := filesystem.NewInMemoryBackend()\n\t\thandler := func(ctx context.Context, detail *ToolDetail) (*ClearResult, error) {\n\t\t\targuments := make(map[string]string)\n\t\t\tif err := json.Unmarshal([]byte(detail.ToolArgument.Text), &arguments); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\toffloadContent := &OffloadContent{\n\t\t\t\tArguments: arguments,\n\t\t\t\tResult:    detail.ToolResult.Parts[0].Text,\n\t\t\t}\n\t\t\treplacedArguments := make(map[string]string, len(arguments))\n\t\t\tfilePath := fmt.Sprintf(\"/tmp/%s\", detail.ToolContext.CallID)\n\t\t\tfor k := range arguments {\n\t\t\t\treplacedArguments[k] = \"argument offloaded\"\n\t\t\t}\n\t\t\treturn &ClearResult{\n\t\t\t\tToolArgument: &schema.ToolArgument{Text: toJson(replacedArguments)},\n\t\t\t\tToolResult: &schema.ToolResult{\n\t\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t\t{Type: schema.ToolPartTypeText, Text: \"result offloaded, retrieve it from \" + filePath},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\tNeedClear:       true,\n\t\t\t\tNeedOffload:     true,\n\t\t\t\tOffloadFilePath: filePath,\n\t\t\t\tOffloadContent:  toJson(offloadContent),\n\t\t\t}, nil\n\t\t}\n\t\tconfig := &Config{\n\t\t\tSkipTruncation:            true,\n\t\t\tTokenCounter:              defaultTokenCounter,\n\t\t\tMaxTokensForClear:         20,\n\t\t\tClearRetentionSuffixLimit: 1,\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"get_weather\": {\n\t\t\t\t\tBackend:      backend,\n\t\t\t\t\tSkipClear:    false,\n\t\t\t\t\tClearHandler: handler,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\t_, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.SystemMessage(\"you are a helpful assistant\"),\n\t\t\t\tschema.UserMessage(\"If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C.\"),\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:       \"call_987654321\",\n\t\t\t\t\t\tType:     \"function\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t\tschema.ToolMessage(\"Sunny\", \"call_123456789\"),\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:       \"call_123456789\",\n\t\t\t\t\t\tType:     \"function\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t\tschema.ToolMessage(\"Sunny\", \"call_123456789\"),\n\t\t\t},\n\t\t}, &adk.ModelContext{\n\t\t\tTools: toolsInfo,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_987654321\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\":\"argument offloaded\",\"unit\":\"argument offloaded\"}`},\n\t\t\t},\n\t\t}, s.Messages[2].ToolCalls)\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_123456789\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[4].ToolCalls)\n\t\tassert.Equal(t, \"result offloaded, retrieve it from /tmp/call_987654321\", s.Messages[3].Content)\n\t\tfileContent, err := backend.Read(ctx, &filesystem.ReadRequest{\n\t\t\tFilePath: \"/tmp/call_987654321\",\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfileContentStr := strings.TrimPrefix(strings.TrimSpace(fileContent.Content), \"1\\t\")\n\t\toc := &OffloadContent{}\n\t\terr = json.Unmarshal([]byte(fileContentStr), oc)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, &OffloadContent{\n\t\t\tArguments: map[string]string{\n\t\t\t\t\"location\": \"London, UK\",\n\t\t\t\t\"unit\":     \"c\",\n\t\t\t},\n\t\t\tResult: \"Sunny\",\n\t\t}, oc)\n\t})\n\n\tt.Run(\"test skip handled ones\", func(t *testing.T) {\n\t\tbackend := filesystem.NewInMemoryBackend()\n\t\tconfig := &Config{\n\t\t\tSkipTruncation:            true,\n\t\t\tTokenCounter:              defaultTokenCounter,\n\t\t\tMaxTokensForClear:         20,\n\t\t\tClearRetentionSuffixLimit: 0,\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"get_weather\": {\n\t\t\t\t\tBackend:      backend,\n\t\t\t\t\tSkipClear:    false,\n\t\t\t\t\tClearHandler: defaultClearHandler(\"/tmp\", true, \"read_file\"),\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tmsgs := []adk.Message{\n\t\t\tschema.SystemMessage(\"you are a helpful assistant\"),\n\t\t\tschema.UserMessage(\"If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C.\"),\n\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID:       \"call_987654321\",\n\t\t\t\t\tType:     \"function\",\n\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t},\n\t\t\t}),\n\t\t\tschema.ToolMessage(\"Sunny\", \"call_123456789\"),\n\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID:       \"call_123456789\",\n\t\t\t\t\tType:     \"function\",\n\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t},\n\t\t\t}),\n\t\t\tschema.ToolMessage(\"Sunny\", \"call_123456789\"),\n\t\t}\n\t\t_, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: msgs}, &adk.ModelContext{Tools: toolsInfo})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_987654321\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[2].ToolCalls)\n\t\tassert.NotNil(t, msgs[2].Extra[msgReducedFlag])\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_123456789\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[4].ToolCalls)\n\t\tassert.Equal(t, \"<persisted-output>Tool result saved to: /tmp/clear/call_987654321\\nUse read_file to view</persisted-output>\", s.Messages[3].Content)\n\t\tfileContent, err := backend.Read(ctx, &filesystem.ReadRequest{\n\t\t\tFilePath: \"/tmp/clear/call_987654321\",\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tfileContentStr := strings.TrimPrefix(strings.TrimSpace(fileContent.Content), \"1\\t\")\n\t\tassert.Equal(t, \"Sunny\", fileContentStr)\n\n\t\tmsgs = append(msgs, []*schema.Message{\n\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID:       \"call_8877665544\",\n\t\t\t\t\tType:     \"function\",\n\t\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t\t},\n\t\t\t}),\n\t\t\tschema.ToolMessage(\"Sunny\", \"call_8877665544\"),\n\t\t}...)\n\t\t_, s, err = mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: msgs}, &adk.ModelContext{Tools: toolsInfo})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_987654321\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[2].ToolCalls)\n\t\tassert.NotNil(t, msgs[2].Extra[msgReducedFlag])\n\t\tassert.Equal(t, []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID:       \"call_123456789\",\n\t\t\t\tType:     \"function\",\n\t\t\t\tFunction: schema.FunctionCall{Name: \"get_weather\", Arguments: `{\"location\": \"London, UK\", \"unit\": \"c\"}`},\n\t\t\t},\n\t\t}, s.Messages[4].ToolCalls)\n\t\tassert.NotNil(t, msgs[4].Extra[msgReducedFlag])\n\t\tassert.Equal(t, \"<persisted-output>Tool result saved to: /tmp/clear/call_987654321\\nUse read_file to view</persisted-output>\", s.Messages[3].Content)\n\t\tassert.Equal(t, \"<persisted-output>Tool result saved to: /tmp/clear/call_123456789\\nUse read_file to view</persisted-output>\", s.Messages[5].Content)\n\t})\n}\n\nfunc TestDefaultOffloadHandler(t *testing.T) {\n\tctx := context.Background()\n\tdetail := &ToolDetail{\n\t\tToolContext: &adk.ToolContext{\n\t\t\tName:   \"mock_name\",\n\t\t\tCallID: \"mock_call_id_12345\",\n\t\t},\n\t\tToolArgument: &schema.ToolArgument{Text: \"anything\"},\n\t\tToolResult:   &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: \"hello\"}}},\n\t}\n\n\tfn := defaultClearHandler(\"/tmp\", true, \"read_file\")\n\tinfo, err := fn(ctx, detail)\n\tassert.NoError(t, err)\n\tassert.Equal(t, &ClearResult{\n\t\tToolArgument: &schema.ToolArgument{Text: \"anything\"},\n\t\tToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{\n\t\t\t{\n\t\t\t\tType: schema.ToolPartTypeText,\n\t\t\t\tText: \"<persisted-output>Tool result saved to: /tmp/clear/mock_call_id_12345\\nUse read_file to view</persisted-output>\",\n\t\t\t},\n\t\t}},\n\t\tNeedClear:       true,\n\t\tNeedOffload:     true,\n\t\tOffloadFilePath: \"/tmp/clear/mock_call_id_12345\",\n\t\tOffloadContent:  \"hello\",\n\t}, info)\n\n}\n\nfunc mockInvokableTool() tool.InvokableTool {\n\ttype ContentContainer struct {\n\t\tValue string `json:\"value\"`\n\t}\n\ts1 := strings.Repeat(\"hello world\", 10) + \"\\n\"\n\ts2 := strings.Repeat(\"hello world\", 8)\n\ts3 := s1 + s2\n\tt, _ := utils.InferTool(\"mock_invokable_tool\", \"test desc\", func(ctx context.Context, input *ContentContainer) (output string, err error) {\n\t\treturn s3, nil\n\t})\n\treturn t\n}\n\nfunc mockStreamableTool() tool.StreamableTool {\n\ttype ContentContainer struct {\n\t\tValue string `json:\"value\"`\n\t}\n\ts1 := strings.Repeat(\"hello world\", 10) + \"\\n\"\n\ts2 := strings.Repeat(\"hello world\", 8)\n\ts3 := s1 + s2\n\tt, _ := utils.InferStreamTool(\"mock_streamable_tool\", \"test desc\", func(ctx context.Context, input ContentContainer) (output *schema.StreamReader[string], err error) {\n\t\tsr, sw := schema.Pipe[string](11)\n\t\tfor _, part := range splitStrings(s3, 10) {\n\t\t\tsw.Send(part, nil)\n\t\t}\n\t\tsw.Close()\n\t\treturn sr, nil\n\t})\n\treturn t\n}\n\nfunc splitStrings(s string, n int) []string {\n\tif n <= 0 {\n\t\tn = 1\n\t}\n\tif n == 1 {\n\t\treturn []string{s}\n\t}\n\tif len(s) <= n {\n\t\tparts := make([]string, n)\n\t\tfor i := 0; i < len(s); i++ {\n\t\t\tparts[i] = string(s[i])\n\t\t}\n\t\treturn parts\n\t}\n\tbaseLen := len(s) / n\n\textra := len(s) % n\n\tparts := make([]string, 0, n)\n\tstart := 0\n\tfor i := 0; i < n; i++ {\n\t\tend := start + baseLen\n\t\tif i < extra {\n\t\t\tend++\n\t\t}\n\t\tparts = append(parts, s[start:end])\n\t\tstart = end\n\t}\n\treturn parts\n}\n\nfunc toJson(v any) string {\n\tb, _ := json.Marshal(v)\n\treturn string(b)\n}\n\nfunc TestToolResultFromMessage(t *testing.T) {\n\tt.Run(\"test from content\", func(t *testing.T) {\n\t\tmsg := schema.ToolMessage(\"test content\", \"call_123\")\n\t\tresult, fromContent, err := toolResultFromMessage(msg)\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, fromContent)\n\t\tassert.NotNil(t, result)\n\t\tassert.Len(t, result.Parts, 1)\n\t\tassert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type)\n\t\tassert.Equal(t, \"test content\", result.Parts[0].Text)\n\t})\n\n\tt.Run(\"test from user input multi content\", func(t *testing.T) {\n\t\tmsg := schema.ToolMessage(\"\", \"call_456\")\n\t\tmsg.UserInputMultiContent = []schema.MessageInputPart{\n\t\t\t{\n\t\t\t\tType: schema.ChatMessagePartTypeText,\n\t\t\t\tText: \"test text\",\n\t\t\t},\n\t\t}\n\t\tresult, fromContent, err := toolResultFromMessage(msg)\n\t\tassert.NoError(t, err)\n\t\tassert.False(t, fromContent)\n\t\tassert.NotNil(t, result)\n\t\tassert.Len(t, result.Parts, 1)\n\t\tassert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type)\n\t\tassert.Equal(t, \"test text\", result.Parts[0].Text)\n\t})\n\n\tt.Run(\"test invalid role\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test user message\")\n\t\t_, _, err := toolResultFromMessage(msg)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"message role\")\n\t})\n}\n\nfunc TestConvMessageInputPartToToolOutputPart(t *testing.T) {\n\tt.Run(\"test text type\", func(t *testing.T) {\n\t\tpart := schema.MessageInputPart{\n\t\t\tType: schema.ChatMessagePartTypeText,\n\t\t\tText: \"test text\",\n\t\t}\n\t\tresult, err := convMessageInputPartToToolOutputPart(part)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, schema.ToolPartTypeText, result.Type)\n\t\tassert.Equal(t, \"test text\", result.Text)\n\t})\n\n\tt.Run(\"test image url type\", func(t *testing.T) {\n\t\tpart := schema.MessageInputPart{\n\t\t\tType:  schema.ChatMessagePartTypeImageURL,\n\t\t\tImage: &schema.MessageInputImage{},\n\t\t}\n\t\tresult, err := convMessageInputPartToToolOutputPart(part)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, schema.ToolPartTypeImage, result.Type)\n\t\tassert.NotNil(t, result.Image)\n\t})\n\n\tt.Run(\"test audio url type\", func(t *testing.T) {\n\t\tpart := schema.MessageInputPart{\n\t\t\tType:  schema.ChatMessagePartTypeAudioURL,\n\t\t\tAudio: &schema.MessageInputAudio{},\n\t\t}\n\t\tresult, err := convMessageInputPartToToolOutputPart(part)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, schema.ToolPartTypeAudio, result.Type)\n\t\tassert.NotNil(t, result.Audio)\n\t})\n\n\tt.Run(\"test video url type\", func(t *testing.T) {\n\t\tpart := schema.MessageInputPart{\n\t\t\tType:  schema.ChatMessagePartTypeVideoURL,\n\t\t\tVideo: &schema.MessageInputVideo{},\n\t\t}\n\t\tresult, err := convMessageInputPartToToolOutputPart(part)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, schema.ToolPartTypeVideo, result.Type)\n\t\tassert.NotNil(t, result.Video)\n\t})\n\n\tt.Run(\"test file url type\", func(t *testing.T) {\n\t\tpart := schema.MessageInputPart{\n\t\t\tType: schema.ChatMessagePartTypeFileURL,\n\t\t\tFile: &schema.MessageInputFile{},\n\t\t}\n\t\tresult, err := convMessageInputPartToToolOutputPart(part)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, schema.ToolPartTypeFile, result.Type)\n\t\tassert.NotNil(t, result.File)\n\t})\n\n\tt.Run(\"test unknown type\", func(t *testing.T) {\n\t\tpart := schema.MessageInputPart{\n\t\t\tType: \"unknown_type\",\n\t\t}\n\t\t_, err := convMessageInputPartToToolOutputPart(part)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"unknown msg part type\")\n\t})\n}\n\nfunc TestGetSetMsgOffloadedFlag(t *testing.T) {\n\tt.Run(\"test get offloaded flag - not set\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\tassert.False(t, getMsgOffloadedFlag(msg))\n\t})\n\n\tt.Run(\"test get offloaded flag - set\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\tsetMsgOffloadedFlag(msg)\n\t\tassert.True(t, getMsgOffloadedFlag(msg))\n\t})\n\n\tt.Run(\"test set offloaded flag - nil extra\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\tsetMsgOffloadedFlag(msg)\n\t\tassert.True(t, getMsgOffloadedFlag(msg))\n\t})\n\n\tt.Run(\"test set offloaded flag - existing extra\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\tmsg.Extra = map[string]any{\"existing\": \"value\"}\n\t\tsetMsgOffloadedFlag(msg)\n\t\tassert.True(t, getMsgOffloadedFlag(msg))\n\t\tassert.Equal(t, \"value\", msg.Extra[\"existing\"])\n\t})\n}\n\nfunc TestGetSetMsgCachedToken(t *testing.T) {\n\tt.Run(\"test get cached token - not set\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\ttokens, ok := getMsgCachedToken(msg)\n\t\tassert.False(t, ok)\n\t\tassert.Equal(t, int64(0), tokens)\n\t})\n\n\tt.Run(\"test get cached token - set\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\tsetMsgCachedToken(msg, 100)\n\t\ttokens, ok := getMsgCachedToken(msg)\n\t\tassert.True(t, ok)\n\t\tassert.Equal(t, int64(100), tokens)\n\t})\n\n\tt.Run(\"test set cached token - nil extra\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\tsetMsgCachedToken(msg, 200)\n\t\ttokens, ok := getMsgCachedToken(msg)\n\t\tassert.True(t, ok)\n\t\tassert.Equal(t, int64(200), tokens)\n\t})\n\n\tt.Run(\"test set cached token - existing extra\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\tmsg.Extra = map[string]any{\"existing\": \"value\"}\n\t\tsetMsgCachedToken(msg, 300)\n\t\ttokens, ok := getMsgCachedToken(msg)\n\t\tassert.True(t, ok)\n\t\tassert.Equal(t, int64(300), tokens)\n\t\tassert.Equal(t, \"value\", msg.Extra[\"existing\"])\n\t})\n}\n\nfunc TestNewErrors(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"test nil config\", func(t *testing.T) {\n\t\t_, err := New(ctx, nil)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"config must not be nil\")\n\t})\n\n\tt.Run(\"test no backend when not skipping truncation\", func(t *testing.T) {\n\t\tconfig := &Config{\n\t\t\tBackend:        nil,\n\t\t\tSkipTruncation: false,\n\t\t}\n\t\t_, err := New(ctx, config)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"backend must be set\")\n\t})\n}\n\nfunc TestGetToolConfig(t *testing.T) {\n\tctx := context.Background()\n\tbackend := filesystem.NewInMemoryBackend()\n\n\tt.Run(\"test no tool config\", func(t *testing.T) {\n\t\tconfig := &Config{\n\t\t\tBackend:        backend,\n\t\t\tSkipTruncation: true,\n\t\t\tSkipClear:      true,\n\t\t}\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\ttrmw, ok := mw.(*toolReductionMiddleware)\n\t\tassert.True(t, ok)\n\n\t\tcfg := trmw.getToolConfig(\"non_existent_tool\", sceneTruncation)\n\t\tassert.NotNil(t, cfg)\n\t})\n\n\tt.Run(\"test with tool config\", func(t *testing.T) {\n\t\tconfig := &Config{\n\t\t\tBackend:        backend,\n\t\t\tSkipTruncation: true,\n\t\t\tSkipClear:      true,\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"test_tool\": {\n\t\t\t\t\tSkipTruncation: true,\n\t\t\t\t\tSkipClear:      true,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\ttrmw, ok := mw.(*toolReductionMiddleware)\n\t\tassert.True(t, ok)\n\n\t\tcfg := trmw.getToolConfig(\"test_tool\", sceneTruncation)\n\t\tassert.NotNil(t, cfg)\n\t\tassert.True(t, cfg.SkipTruncation)\n\t})\n\n\tt.Run(\"test with tool config needing default handler\", func(t *testing.T) {\n\t\tconfig := &Config{\n\t\t\tBackend:        backend,\n\t\t\tSkipTruncation: false,\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"test_tool\": {\n\t\t\t\t\tSkipTruncation: false,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tmw, err := New(ctx, config)\n\t\tassert.NoError(t, err)\n\t\ttrmw, ok := mw.(*toolReductionMiddleware)\n\t\tassert.True(t, ok)\n\n\t\tcfg := trmw.getToolConfig(\"test_tool\", sceneTruncation)\n\t\tassert.NotNil(t, cfg)\n\t\tassert.NotNil(t, cfg.TruncHandler)\n\t})\n}\n\nfunc TestCopyAndFillDefaults(t *testing.T) {\n\tt.Run(\"test empty config\", func(t *testing.T) {\n\t\tcfg := &Config{}\n\t\tresult, err := cfg.copyAndFillDefaults()\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, result)\n\t\tassert.Equal(t, \"/tmp\", result.RootDir)\n\t\tassert.Equal(t, \"read_file\", result.ReadFileToolName)\n\t\tassert.Equal(t, 50000, result.MaxLengthForTrunc)\n\t\tassert.Equal(t, 1, result.ClearRetentionSuffixLimit)\n\t\tassert.NotNil(t, result.TokenCounter)\n\t})\n\n\tt.Run(\"test with tool config\", func(t *testing.T) {\n\t\tcfg := &Config{\n\t\t\tToolConfig: map[string]*ToolReductionConfig{\n\t\t\t\t\"test_tool\": {\n\t\t\t\t\tSkipTruncation: true,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult, err := cfg.copyAndFillDefaults()\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, result.ToolConfig)\n\t\tassert.True(t, result.ToolConfig[\"test_tool\"].SkipTruncation)\n\t})\n}\n\nfunc TestDefaultTokenCounter(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"test with nil messages\", func(t *testing.T) {\n\t\tmsgs := []*schema.Message{nil}\n\t\ttokens, err := defaultTokenCounter(ctx, msgs, nil)\n\t\tassert.NoError(t, err)\n\t\tassert.GreaterOrEqual(t, tokens, int64(0))\n\t})\n\n\tt.Run(\"test with tool info\", func(t *testing.T) {\n\t\ttoolInfo := &schema.ToolInfo{\n\t\t\tName: \"test_tool\",\n\t\t\tDesc: \"test description\",\n\t\t}\n\t\ttokens, err := defaultTokenCounter(ctx, nil, []*schema.ToolInfo{toolInfo})\n\t\tassert.NoError(t, err)\n\t\tassert.GreaterOrEqual(t, tokens, int64(0))\n\t})\n}\n\nfunc TestDefaultClearHandler(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"test empty parts\", func(t *testing.T) {\n\t\thandler := defaultClearHandler(\"/tmp\", true, \"read_file\")\n\t\tdetail := &ToolDetail{\n\t\t\tToolContext: &adk.ToolContext{\n\t\t\t\tCallID: \"test_call\",\n\t\t\t},\n\t\t\tToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{}},\n\t\t}\n\t\tresult, err := handler(ctx, detail)\n\t\tassert.NoError(t, err)\n\t\tassert.False(t, result.NeedClear)\n\t})\n\n\tt.Run(\"test multimodal content\", func(t *testing.T) {\n\t\thandler := defaultClearHandler(\"/tmp\", true, \"read_file\")\n\t\tdetail := &ToolDetail{\n\t\t\tToolContext: &adk.ToolContext{\n\t\t\t\tCallID: \"test_call\",\n\t\t\t},\n\t\t\tToolResult: &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeImage}},\n\t\t\t},\n\t\t}\n\t\t_, err := handler(ctx, detail)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"not support multimodal\")\n\t})\n\n\tt.Run(\"test no call id\", func(t *testing.T) {\n\t\thandler := defaultClearHandler(\"/tmp\", true, \"read_file\")\n\t\tdetail := &ToolDetail{\n\t\t\tToolContext: &adk.ToolContext{},\n\t\t\tToolResult: &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: \"test\"}},\n\t\t\t},\n\t\t}\n\t\tresult, err := handler(ctx, detail)\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, result.NeedClear)\n\t\tassert.NotEmpty(t, result.OffloadFilePath)\n\t})\n}\n"
  },
  {
    "path": "adk/middlewares/skill/filesystem_backend.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage skill\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"strings\"\n\n\t\"gopkg.in/yaml.v3\"\n\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n)\n\nconst skillFileName = \"SKILL.md\"\n\ntype filesystemBackend struct {\n\tbackend filesystem.Backend\n\tbaseDir string\n}\n\n// BackendFromFilesystemConfig contains configuration for NewBackendFromFilesystem.\ntype BackendFromFilesystemConfig struct {\n\t// Backend is the filesystem.Backend implementation used for file operations.\n\tBackend filesystem.Backend\n\t// BaseDir is the base directory where skill directories are located.\n\t// Each skill should be in a subdirectory containing a SKILL.md file.\n\tBaseDir string\n}\n\n// NewBackendFromFilesystem creates a new Backend implementation that reads skills from a filesystem.\n// It searches for SKILL.md files in immediate subdirectories of the configured BaseDir.\n// Only first-level subdirectories are scanned; deeply nested SKILL.md files are ignored.\nfunc NewBackendFromFilesystem(_ context.Context, config *BackendFromFilesystemConfig) (Backend, error) {\n\tif config == nil {\n\t\treturn nil, fmt.Errorf(\"config is required\")\n\t}\n\tif config.Backend == nil {\n\t\treturn nil, fmt.Errorf(\"backend is required\")\n\t}\n\tif config.BaseDir == \"\" {\n\t\treturn nil, fmt.Errorf(\"baseDir is required\")\n\t}\n\n\treturn &filesystemBackend{\n\t\tbackend: config.Backend,\n\t\tbaseDir: config.BaseDir,\n\t}, nil\n}\n\nfunc (b *filesystemBackend) List(ctx context.Context) ([]FrontMatter, error) {\n\tskills, err := b.list(ctx)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to list skills: %w\", err)\n\t}\n\n\tmatters := make([]FrontMatter, 0, len(skills))\n\tfor _, skill := range skills {\n\t\tmatters = append(matters, skill.FrontMatter)\n\t}\n\n\treturn matters, nil\n}\n\nfunc (b *filesystemBackend) Get(ctx context.Context, name string) (Skill, error) {\n\tskills, err := b.list(ctx)\n\tif err != nil {\n\t\treturn Skill{}, fmt.Errorf(\"failed to list skills: %w\", err)\n\t}\n\n\tfor _, skill := range skills {\n\t\tif skill.Name == name {\n\t\t\treturn skill, nil\n\t\t}\n\t}\n\n\treturn Skill{}, fmt.Errorf(\"skill not found: %s\", name)\n}\n\nfunc (b *filesystemBackend) list(ctx context.Context) ([]Skill, error) {\n\tvar skills []Skill\n\n\tpattern := \"*/\" + skillFileName\n\tentries, err := b.backend.GlobInfo(ctx, &filesystem.GlobInfoRequest{\n\t\tPattern: pattern,\n\t\tPath:    b.baseDir,\n\t})\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to glob skill files: %w\", err)\n\t}\n\n\tfor _, entry := range entries {\n\t\tfilePath := entry.Path\n\t\tif !filepath.IsAbs(filePath) {\n\t\t\tfilePath = filepath.Join(b.baseDir, filePath)\n\t\t}\n\t\tskill, loadErr := b.loadSkillFromFile(ctx, filePath)\n\t\tif loadErr != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to load skill from %s: %w\", filePath, loadErr)\n\t\t}\n\n\t\tskills = append(skills, skill)\n\t}\n\n\treturn skills, nil\n}\n\nfunc (b *filesystemBackend) loadSkillFromFile(ctx context.Context, path string) (Skill, error) {\n\tfileContent, err := b.backend.Read(ctx, &filesystem.ReadRequest{\n\t\tFilePath: path,\n\t})\n\tif err != nil {\n\t\treturn Skill{}, fmt.Errorf(\"failed to read file: %w\", err)\n\t}\n\n\tdata := stripLineNumbers(fileContent.Content)\n\n\tfrontmatter, content, err := parseFrontmatter(data)\n\tif err != nil {\n\t\treturn Skill{}, fmt.Errorf(\"failed to parse frontmatter: %w\", err)\n\t}\n\n\tvar fm FrontMatter\n\tif err = yaml.Unmarshal([]byte(frontmatter), &fm); err != nil {\n\t\treturn Skill{}, fmt.Errorf(\"failed to unmarshal frontmatter: %w\", err)\n\t}\n\n\tabsDir := filepath.Dir(path)\n\n\treturn Skill{\n\t\tFrontMatter:   fm,\n\t\tContent:       strings.TrimSpace(content),\n\t\tBaseDirectory: absDir,\n\t}, nil\n}\n\nfunc stripLineNumbers(data string) string {\n\tlines := strings.Split(data, \"\\n\")\n\tresult := make([]string, 0, len(lines))\n\tfor _, line := range lines {\n\t\tidx := strings.Index(line, \"\\t\")\n\t\tif idx != -1 {\n\t\t\tline = line[idx+1:]\n\t\t}\n\t\tresult = append(result, line)\n\t}\n\treturn strings.Join(result, \"\\n\")\n}\n\nfunc parseFrontmatter(data string) (frontmatter string, content string, err error) {\n\tconst delimiter = \"---\"\n\n\tdata = strings.TrimSpace(data)\n\n\tif !strings.HasPrefix(data, delimiter) {\n\t\treturn \"\", \"\", fmt.Errorf(\"file does not start with frontmatter delimiter\")\n\t}\n\n\trest := data[len(delimiter):]\n\tendIdx := strings.Index(rest, \"\\n\"+delimiter)\n\tif endIdx == -1 {\n\t\treturn \"\", \"\", fmt.Errorf(\"frontmatter closing delimiter not found\")\n\t}\n\n\tfrontmatter = strings.TrimSpace(rest[:endIdx])\n\tcontent = rest[endIdx+len(\"\\n\"+delimiter):]\n\n\tif strings.HasPrefix(content, \"\\n\") {\n\t\tcontent = content[1:]\n\t}\n\n\treturn frontmatter, content, nil\n}\n"
  },
  {
    "path": "adk/middlewares/skill/filesystem_backend_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage skill\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n)\n\nfunc TestNewBackendFromFilesystem(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"nil config returns error\", func(t *testing.T) {\n\t\tbackend, err := NewBackendFromFilesystem(ctx, nil)\n\t\tassert.Nil(t, backend)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"config is required\")\n\t})\n\n\tt.Run(\"nil backend returns error\", func(t *testing.T) {\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\tassert.Nil(t, backend)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"backend is required\")\n\t})\n\n\tt.Run(\"empty baseDir returns error\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"\",\n\t\t})\n\t\tassert.Nil(t, backend)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"baseDir is required\")\n\t})\n\n\tt.Run(\"valid config succeeds\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, backend)\n\t})\n}\n\nfunc TestFilesystemBackend_List(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"empty directory returns empty list\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/.keep\",\n\t\t\tContent:  \"\",\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskills, err := backend.List(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, skills)\n\t})\n\n\tt.Run(\"directory with no SKILL.md files returns empty list\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/subdir/other.txt\",\n\t\t\tContent:  \"some content\",\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskills, err := backend.List(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, skills)\n\t})\n\n\tt.Run(\"files in root directory are ignored\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/SKILL.md\",\n\t\t\tContent: `---\nname: root-skill\ndescription: Root skill\n---\nContent`,\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskills, err := backend.List(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, skills)\n\t})\n\n\tt.Run(\"valid skill directory returns skill\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/my-skill/SKILL.md\",\n\t\t\tContent: `---\nname: pdf-processing\ndescription: Extract text and tables from PDF files, fill forms, merge documents.\nlicense: Apache-2.0\nmetadata:\n  author: example-org\n  version: \"1.0\"\n---\nThis is the skill content.`,\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskills, err := backend.List(ctx)\n\t\tassert.NoError(t, err)\n\t\trequire.Len(t, skills, 1)\n\t\tassert.Equal(t, \"pdf-processing\", skills[0].Name)\n\t\tassert.Equal(t, \"Extract text and tables from PDF files, fill forms, merge documents.\", skills[0].Description)\n\t})\n\n\tt.Run(\"multiple skill directories returns all skills\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/skill-1/SKILL.md\",\n\t\t\tContent: `---\nname: skill-1\ndescription: First skill\n---\nContent 1`,\n\t\t})\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/skill-2/SKILL.md\",\n\t\t\tContent: `---\nname: skill-2\ndescription: Second skill\n---\nContent 2`,\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskills, err := backend.List(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, skills, 2)\n\n\t\tnames := []string{skills[0].Name, skills[1].Name}\n\t\tassert.Contains(t, names, \"skill-1\")\n\t\tassert.Contains(t, names, \"skill-2\")\n\t})\n\n\tt.Run(\"invalid SKILL.md returns error\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/invalid-skill/SKILL.md\",\n\t\t\tContent:  `No frontmatter here`,\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskills, err := backend.List(ctx)\n\t\tassert.Error(t, err)\n\t\tassert.Nil(t, skills)\n\t\tassert.Contains(t, err.Error(), \"failed to load skill\")\n\t})\n\n\tt.Run(\"non-existent baseDir returns empty list\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/path/that/does/not/exist\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskills, err := backend.List(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, skills)\n\t})\n\n\tt.Run(\"deeply nested SKILL.md is ignored\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/valid-skill/SKILL.md\",\n\t\t\tContent: `---\nname: valid-skill\ndescription: Valid skill\n---\nContent`,\n\t\t})\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/deep/nested/SKILL.md\",\n\t\t\tContent: `---\nname: nested-skill\ndescription: Nested skill\n---\nContent`,\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskills, err := backend.List(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, skills, 1)\n\t\tassert.Equal(t, \"valid-skill\", skills[0].Name)\n\t})\n}\n\nfunc TestFilesystemBackend_Get(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"skill not found returns error\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/.keep\",\n\t\t\tContent:  \"\",\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskill, err := backend.Get(ctx, \"non-existent\")\n\t\tassert.Error(t, err)\n\t\tassert.Empty(t, skill)\n\t\tassert.Contains(t, err.Error(), \"skill not found\")\n\t})\n\n\tt.Run(\"existing skill is returned\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/test-skill/SKILL.md\",\n\t\t\tContent: `---\nname: test-skill\ndescription: Test skill description\n---\nTest content here.`,\n\t\t})\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskill, err := backend.Get(ctx, \"test-skill\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"test-skill\", skill.Name)\n\t\tassert.Equal(t, \"Test skill description\", skill.Description)\n\t\tassert.Equal(t, \"Test content here.\", skill.Content)\n\t})\n\n\tt.Run(\"get specific skill from multiple\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\tfor _, name := range []string{\"alpha\", \"beta\", \"gamma\"} {\n\t\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\t\tFilePath: \"/skills/\" + name + \"/SKILL.md\",\n\t\t\t\tContent: `---\nname: ` + name + `\ndescription: Skill ` + name + `\n---\nContent for ` + name,\n\t\t\t})\n\t\t}\n\n\t\tbackend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{\n\t\t\tBackend: fsBackend,\n\t\t\tBaseDir: \"/skills\",\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\tskill, err := backend.Get(ctx, \"beta\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"beta\", skill.Name)\n\t\tassert.Equal(t, \"Skill beta\", skill.Description)\n\t\tassert.Equal(t, \"Content for beta\", skill.Content)\n\t})\n}\n\nfunc TestParseFrontmatter(t *testing.T) {\n\tt.Run(\"valid frontmatter\", func(t *testing.T) {\n\t\tdata := `---\nname: test\ndescription: test description\n---\nThis is the content.`\n\n\t\tfm, content, err := parseFrontmatter(data)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"name: test\\ndescription: test description\", fm)\n\t\tassert.Equal(t, \"This is the content.\", content)\n\t})\n\n\tt.Run(\"frontmatter with multiline content\", func(t *testing.T) {\n\t\tdata := `---\nname: test\n---\nLine 1\nLine 2\nLine 3`\n\n\t\tfm, content, err := parseFrontmatter(data)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"name: test\", fm)\n\t\tassert.Equal(t, \"Line 1\\nLine 2\\nLine 3\", content)\n\t})\n\n\tt.Run(\"frontmatter with leading/trailing whitespace\", func(t *testing.T) {\n\t\tdata := `  \n---\nname: test\n---\nContent  `\n\n\t\tfm, content, err := parseFrontmatter(data)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"name: test\", fm)\n\t\tassert.Equal(t, \"Content\", content)\n\t})\n\n\tt.Run(\"missing opening delimiter returns error\", func(t *testing.T) {\n\t\tdata := `name: test\n---\nContent`\n\n\t\tfm, content, err := parseFrontmatter(data)\n\t\tassert.Error(t, err)\n\t\tassert.Empty(t, fm)\n\t\tassert.Empty(t, content)\n\t\tassert.Contains(t, err.Error(), \"does not start with frontmatter delimiter\")\n\t})\n\n\tt.Run(\"missing closing delimiter returns error\", func(t *testing.T) {\n\t\tdata := `---\nname: test\nContent without closing`\n\n\t\tfm, content, err := parseFrontmatter(data)\n\t\tassert.Error(t, err)\n\t\tassert.Empty(t, fm)\n\t\tassert.Empty(t, content)\n\t\tassert.Contains(t, err.Error(), \"closing delimiter not found\")\n\t})\n\n\tt.Run(\"empty frontmatter\", func(t *testing.T) {\n\t\tdata := `---\n---\nContent only`\n\n\t\tfm, content, err := parseFrontmatter(data)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, fm)\n\t\tassert.Equal(t, \"Content only\", content)\n\t})\n\n\tt.Run(\"empty content\", func(t *testing.T) {\n\t\tdata := `---\nname: test\n---`\n\n\t\tfm, content, err := parseFrontmatter(data)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"name: test\", fm)\n\t\tassert.Empty(t, content)\n\t})\n\n\tt.Run(\"content with --- inside\", func(t *testing.T) {\n\t\tdata := `---\nname: test\n---\nContent with --- in the middle`\n\n\t\tfm, content, err := parseFrontmatter(data)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"name: test\", fm)\n\t\tassert.Equal(t, \"Content with --- in the middle\", content)\n\t})\n}\n\nfunc TestLoadSkillFromFile(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"valid skill file\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/SKILL.md\",\n\t\t\tContent: `---\nname: file-skill\ndescription: Skill from file\n---\nFile skill content.`,\n\t\t})\n\n\t\tbackend := &filesystemBackend{backend: fsBackend, baseDir: \"/skills\"}\n\t\tskill, err := backend.loadSkillFromFile(ctx, \"/skills/SKILL.md\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"file-skill\", skill.Name)\n\t\tassert.Equal(t, \"Skill from file\", skill.Description)\n\t\tassert.Equal(t, \"File skill content.\", skill.Content)\n\t\tassert.Equal(t, \"/skills\", skill.BaseDirectory)\n\t})\n\n\tt.Run(\"non-existent file returns error\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\tbackend := &filesystemBackend{backend: fsBackend, baseDir: \"/tmp\"}\n\t\tskill, err := backend.loadSkillFromFile(ctx, \"/path/to/nonexistent/SKILL.md\")\n\t\tassert.Error(t, err)\n\t\tassert.Empty(t, skill)\n\t\tassert.Contains(t, err.Error(), \"failed to read file\")\n\t})\n\n\tt.Run(\"invalid yaml in frontmatter returns error\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/SKILL.md\",\n\t\t\tContent: `---\nname: [invalid yaml\n---\nContent`,\n\t\t})\n\n\t\tbackend := &filesystemBackend{backend: fsBackend, baseDir: \"/skills\"}\n\t\tskill, err := backend.loadSkillFromFile(ctx, \"/skills/SKILL.md\")\n\t\tassert.Error(t, err)\n\t\tassert.Empty(t, skill)\n\t\tassert.Contains(t, err.Error(), \"failed to unmarshal frontmatter\")\n\t})\n\n\tt.Run(\"content with extra whitespace is trimmed\", func(t *testing.T) {\n\t\tfsBackend := filesystem.NewInMemoryBackend()\n\t\t_ = fsBackend.Write(ctx, &filesystem.WriteRequest{\n\t\t\tFilePath: \"/skills/SKILL.md\",\n\t\t\tContent: `---\nname: trimmed-skill\ndescription: desc\n---\n\n   Content with whitespace   \n\n`,\n\t\t})\n\n\t\tbackend := &filesystemBackend{backend: fsBackend, baseDir: \"/skills\"}\n\t\tskill, err := backend.loadSkillFromFile(ctx, \"/skills/SKILL.md\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"Content with whitespace\", skill.Content)\n\t})\n}\n"
  },
  {
    "path": "adk/middlewares/skill/prompt.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage skill\n\nconst (\n\tsystemPrompt = `\n# Skills System\n\n**How to Use Skills (Progressive Disclosure):**\n\nSkills follow a **progressive disclosure** pattern - you see their name and description above, but only read full instructions when needed:\n\n1. **Recognize when a skill applies**: Check if the user's task matches a skill's description\n2. **Read the skill's full instructions**: Use the '{tool_name}' tool to load skill\n3. **Follow the skill's instructions**: tool result contains step-by-step workflows, best practices, and examples\n4. **Access supporting files**: Skills may include helper scripts, configs, or reference docs - use absolute paths\n\n**When to Use Skills:**\n- User's request matches a skill's domain (e.g., \"research X\" -> web-research skill)\n- You need specialized knowledge or structured workflows\n- A skill provides proven patterns for complex tasks\n\n**Executing Skill Scripts:**\nSkills may contain Python scripts or other executable files. Always use absolute paths.\n\n**Example Workflow:**\n\nUser: \"Can you research the latest developments in quantum computing?\"\n\n1. Check available skills -> See \"web-research\" skill\n2. Call '{tool_name}' tool to read the full skill instructions\n3. Follow the skill's research workflow (search -> organize -> synthesize)\n4. Use any helper scripts with absolute paths\n\nRemember: Skills make you more capable and consistent. When in doubt, check if a skill exists for the task!\n`\n\n\tsystemPromptChinese = `\n# Skill 系统\n\n**如何使用 Skill（技能）（渐进式展示）：**\n\nSkill 遵循**渐进式展示**模式 - 你可以在上方看到 Skill 的名称和描述，但只在需要时才阅读完整说明：\n\n1. **识别 Skill 适用场景**：检查用户的任务是否匹配某个 Skill 的描述\n2. **阅读 Skill 的完整说明**：使用 '{tool_name}' 工具加载 Skill\n3. **遵循 Skill 说明操作**：工具结果包含逐步工作流程、最佳实践和示例\n4. **访问支持文件**：Skill 可能包含辅助脚本、配置或参考文档 - 使用绝对路径访问\n\n**何时使用 Skill：**\n- 用户请求匹配某个 Skill 的领域（例如\"研究 X\" -> web-research Skill）\n- 你需要专业知识或结构化工作流程\n- 某个 Skill 为复杂任务提供了经过验证的模式\n\n**执行 Skill 脚本：**\nSkill 可能包含 Python 脚本或其他可执行文件。始终使用绝对路径。\n\n**示例工作流程：**\n\n用户：\"你能研究一下量子计算的最新发展吗？\"\n\n1. 检查可用 Skill -> 发现 \"web-research\" Skill\n2. 调用 '{tool_name}' 工具读取完整的 Skill 说明\n3. 遵循 Skill 的研究工作流程（搜索 -> 整理 -> 综合）\n4. 使用绝对路径运行任何辅助脚本\n\n记住：Skill 让你更加强大和稳定。如有疑问，请检查是否存在适用于该任务的 Skill！\n`\n\n\ttoolDescriptionBase = `Execute a skill within the main conversation\n\n<skills_instructions>\nWhen users ask you to perform tasks, check if any of the available skills below can help complete the task more effectively. Skills provide specialized capabilities and domain knowledge.\n\nHow to invoke:\n- Use the exact string inside <name> tag as the skill name (no arguments)\n- Examples:\n  - ` + \"`\" + `skill: \"pdf\"` + \"`\" + ` - invoke the pdf skill\n  - ` + \"`\" + `skill: \"xlsx\"` + \"`\" + ` - invoke the xlsx skill\n  - ` + \"`\" + `skill: \"ms-office-suite:pdf\"` + \"`\" + ` - invoke using fully qualified name\n\nImportant:\n- When a skill is relevant, you must invoke this tool IMMEDIATELY as your first action\n- NEVER just announce or mention a skill in your text response without actually calling this tool\n- This is a BLOCKING REQUIREMENT: invoke the relevant Skill tool BEFORE generating any other response about the task\n- Only use skills listed in <available_skills> below\n- Do not invoke a skill that is already running\n- Skill content may contain relative paths. Convert them to absolute paths using the base directory provided in the tool result\n</skills_instructions>\n\n`\n\ttoolDescriptionBaseChinese = `在主对话中执行 Skill（技能）\n\n<skills_instructions>\n当用户要求你执行任务时，检查下方可用 Skill 列表中是否有 Skill 可以更有效地完成任务。Skill 提供专业能力和领域知识。\n\n如何调用：\n- 使用 <name> 标签内的完整字符串作为 Skill 名称（无需其他参数）\n- 示例：\n  - ` + \"`\" + `skill: \"pdf\"` + \"`\" + ` - 调用 pdf Skill\n  - ` + \"`\" + `skill: \"xlsx\"` + \"`\" + ` - 调用 xlsx Skill\n  - ` + \"`\" + `skill: \"ms-office-suite:pdf\"` + \"`\" + ` - 使用完全限定名称调用\n\n重要说明：\n- 当 Skill 相关时，你必须立即调用此工具作为第一个动作\n- 切勿仅在文本回复中提及 Skill 而不实际调用此工具\n- 这是阻塞性要求：在生成任何关于任务的其他响应之前，先调用相关的 Skill 工具\n- 仅使用 <available_skills> 中列出的 Skill\n- 不要调用已经运行中的 Skill\n- Skill 内容中可能包含相对路径，需使用工具返回的 base directory 将其转换为绝对路径\n</skills_instructions>\n\n`\n\ttoolDescriptionTemplate = `\n<available_skills>\n{{- range .Matters }}\n<skill>\n<name>\n{{ .Name }}\n</name>\n<description>\n{{ .Description }}\n</description>\n</skill>\n{{- end }}\n</available_skills>\n`\n\ttoolResult        = \"Launching skill: %s\\n\"\n\ttoolResultChinese = \"正在启动 Skill：%s\\n\"\n\tuserContent       = `Base directory for this skill: %s\n\n%s`\n\tuserContentChinese = `此 Skill 的目录：%s\n\n%s`\n\ttoolName = \"skill\"\n\n\tsubAgentResultFormat        = \"Skill \\\"%s\\\" completed (sub-agent execution).\\n\\nResult:\\n%s\"\n\tsubAgentResultFormatChinese = \"Skill \\\"%s\\\" 已完成（子 Agent 执行）。\\n\\n结果：\\n%s\"\n)\n"
  },
  {
    "path": "adk/middlewares/skill/skill.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package skill provides the skill middleware, types, and a local filesystem backend.\npackage skill\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"strings\"\n\t\"text/template\"\n\n\t\"github.com/slongfield/pyfmt\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype ContextMode string\n\nconst (\n\t// ContextModeForkWithContext forks a new agent to run the skill,\n\t// carrying over the original message history from the parent agent.\n\tContextModeForkWithContext ContextMode = \"fork_with_context\"\n\t// ContextModeFork forks a new agent to run the skill\n\t// with a clean context, discarding the original message history.\n\tContextModeFork ContextMode = \"fork\"\n)\n\ntype FrontMatter struct {\n\tName        string      `yaml:\"name\"`\n\tDescription string      `yaml:\"description\"`\n\tContext     ContextMode `yaml:\"context\"`\n\tAgent       string      `yaml:\"agent\"`\n\tModel       string      `yaml:\"model\"`\n}\n\ntype Skill struct {\n\tFrontMatter\n\tContent       string\n\tBaseDirectory string\n}\n\ntype Backend interface {\n\tList(ctx context.Context) ([]FrontMatter, error)\n\tGet(ctx context.Context, name string) (Skill, error)\n}\n\n// AgentHubOptions contains options passed to AgentHub.Get when creating an agent for skill execution.\ntype AgentHubOptions struct {\n\t// Model is the resolved model instance when a skill specifies a \"model\" field in frontmatter.\n\t// nil means the skill did not specify a model override; implementations should use their default.\n\tModel model.ToolCallingChatModel\n}\n\n// AgentHub provides agent instances for context mode (fork/fork_with_context) execution.\ntype AgentHub interface {\n\t// Get returns an Agent by name. When name is empty, implementations should return a default agent.\n\t// The opts parameter carries skill-level overrides (e.g., model) resolved by the framework.\n\tGet(ctx context.Context, name string, opts *AgentHubOptions) (adk.Agent, error)\n}\n\n// ModelHub resolves model instances by name for skills that specify a \"model\" field in frontmatter.\ntype ModelHub interface {\n\tGet(ctx context.Context, name string) (model.ToolCallingChatModel, error)\n}\n\n// SystemPromptFunc is a function that returns a custom system prompt.\n// The toolName parameter is the name of the skill tool (default: \"skill\").\ntype SystemPromptFunc func(ctx context.Context, toolName string) string\n\n// ToolDescriptionFunc is a function that returns a custom tool description.\n// The skills parameter contains all available skill front matters.\ntype ToolDescriptionFunc func(ctx context.Context, skills []FrontMatter) string\n\n// Config is the configuration for the skill middleware.\ntype Config struct {\n\t// Backend is the backend for retrieving skills.\n\tBackend Backend\n\t// SkillToolName is the custom name for the skill tool. If nil, the default name \"skill\" is used.\n\tSkillToolName *string\n\t// Deprecated: Use adk.SetLanguage(adk.LanguageChinese) instead to enable Chinese prompts globally.\n\t// This field will be removed in a future version.\n\tUseChinese bool\n\t// AgentHub provides agent factories for context mode (fork/isolate) execution.\n\t// Required when skills use \"context: fork\" or \"context: isolate\" in frontmatter.\n\t// The agent factory is retrieved by agent name (skill.Agent) from this hub.\n\t// When skill.Agent is empty, AgentHub.Get is called with an empty string,\n\t// allowing the hub implementation to return a default agent.\n\tAgentHub AgentHub\n\t// ModelHub provides model instances for skills that specify a \"model\" field in frontmatter.\n\t// Used in two scenarios:\n\t//   - With context mode (fork/isolate): The model is passed to the AgentFactory\n\t//   - Without context mode (inline): The model becomes active for subsequent ChatModel requests\n\t// If nil, skills with model specification will be ignored in inline mode,\n\t// or return an error in context mode.\n\tModelHub ModelHub\n\n\t// CustomSystemPrompt allows customizing the system prompt injected into the agent.\n\t// If nil, the default system prompt is used.\n\t// The function receives the skill tool name as a parameter.\n\tCustomSystemPrompt SystemPromptFunc\n\t// CustomToolDescription allows customizing the tool description for the skill tool.\n\t// If nil, the default tool description is used.\n\t// The function receives all available skill front matters as a parameter.\n\tCustomToolDescription ToolDescriptionFunc\n}\n\n// NewMiddleware creates a new skill middleware handler for ChatModelAgent.\n//\n// The handler provides a skill tool that allows agents to load and execute skills\n// defined in SKILL.md files. Skills can run in different modes based on their\n// frontmatter configuration:\n//\n//   - Inline mode (default): Skill content is returned directly as tool result\n//   - Fork mode (context: fork): Forks a new agent with a clean context, discarding message history\n//   - Fork with context mode (context: fork_with_context): Forks a new agent carrying over message history\n//\n// Example usage:\n//\n//\thandler, err := skill.NewMiddleware(ctx, &skill.Config{\n//\t    Backend:  backend,\n//\t    AgentHub: myAgentHub,\n//\t    ModelHub: myModelHub,\n//\t})\n//\tif err != nil {\n//\t    return err\n//\t}\n//\n//\tagent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n//\t    // ...\n//\t    Middlewares: []adk.ChatModelAgentMiddleware{handler},\n//\t})\nfunc NewMiddleware(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) {\n\tif config == nil {\n\t\treturn nil, fmt.Errorf(\"config is required\")\n\t}\n\tif config.Backend == nil {\n\t\treturn nil, fmt.Errorf(\"backend is required\")\n\t}\n\n\tname := toolName\n\tif config.SkillToolName != nil {\n\t\tname = *config.SkillToolName\n\t}\n\n\tvar instruction string\n\tif config.CustomSystemPrompt != nil {\n\t\tinstruction = config.CustomSystemPrompt(ctx, name)\n\t} else {\n\t\tvar err error\n\t\tinstruction, err = buildSystemPrompt(name, config.UseChinese)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\treturn &skillHandler{\n\t\tinstruction: instruction,\n\t\ttool: &skillTool{\n\t\t\tb:                     config.Backend,\n\t\t\ttoolName:              name,\n\t\t\tuseChinese:            config.UseChinese,\n\t\t\tagentHub:              config.AgentHub,\n\t\t\tmodelHub:              config.ModelHub,\n\t\t\tcustomToolDescription: config.CustomToolDescription,\n\t\t},\n\t}, nil\n}\n\ntype skillHandler struct {\n\t*adk.BaseChatModelAgentMiddleware\n\tinstruction string\n\ttool        *skillTool\n}\n\nfunc (h *skillHandler) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) {\n\trunCtx.Instruction = runCtx.Instruction + \"\\n\" + h.instruction\n\trunCtx.Tools = append(runCtx.Tools, h.tool)\n\treturn ctx, runCtx, nil\n}\n\nfunc (h *skillHandler) WrapModel(ctx context.Context, m model.BaseChatModel, mc *adk.ModelContext) (model.BaseChatModel, error) {\n\tif h.tool.modelHub == nil {\n\t\treturn m, nil\n\t}\n\tmodelName, found, err := adk.GetRunLocalValue(ctx, activeModelKey)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to get active model from run local value: %w\", err)\n\t}\n\tif !found {\n\t\treturn m, nil\n\t}\n\tname, ok := modelName.(string)\n\tif !ok || name == \"\" {\n\t\treturn m, nil\n\t}\n\tnewModel, err := h.tool.modelHub.Get(ctx, name)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to get model '%s' from ModelHub: %w\", name, err)\n\t}\n\treturn newModel, nil\n}\n\nconst activeModelKey = \"__skill_active_model__\"\n\n// New creates a new skill middleware.\n// It provides a tool for the agent to use skills.\n//\n// Deprecated: Use NewChatModelAgentMiddleware instead. New does not support fork mode execution\n// because AgentMiddleware cannot save message history for fork mode.\nfunc New(ctx context.Context, config *Config) (adk.AgentMiddleware, error) {\n\tif config == nil {\n\t\treturn adk.AgentMiddleware{}, fmt.Errorf(\"config is required\")\n\t}\n\tif config.Backend == nil {\n\t\treturn adk.AgentMiddleware{}, fmt.Errorf(\"backend is required\")\n\t}\n\n\tname := toolName\n\tif config.SkillToolName != nil {\n\t\tname = *config.SkillToolName\n\t}\n\n\tvar sp string\n\tif config.CustomSystemPrompt != nil {\n\t\tsp = config.CustomSystemPrompt(ctx, name)\n\t} else {\n\t\tvar err error\n\t\tsp, err = buildSystemPrompt(name, config.UseChinese)\n\t\tif err != nil {\n\t\t\treturn adk.AgentMiddleware{}, err\n\t\t}\n\t}\n\n\treturn adk.AgentMiddleware{\n\t\tAdditionalInstruction: sp,\n\t\tAdditionalTools: []tool.BaseTool{&skillTool{\n\t\t\tb:                     config.Backend,\n\t\t\ttoolName:              name,\n\t\t\tuseChinese:            config.UseChinese,\n\t\t\tcustomToolDescription: config.CustomToolDescription,\n\t\t}},\n\t}, nil\n}\n\nfunc buildSystemPrompt(skillToolName string, useChinese bool) (string, error) {\n\tvar prompt string\n\tif useChinese {\n\t\tprompt = systemPromptChinese\n\t} else {\n\t\tprompt = internal.SelectPrompt(internal.I18nPrompts{\n\t\t\tEnglish: systemPrompt,\n\t\t\tChinese: systemPromptChinese,\n\t\t})\n\t}\n\treturn pyfmt.Fmt(prompt, map[string]string{\n\t\t\"tool_name\": skillToolName,\n\t})\n}\n\ntype skillTool struct {\n\tb                     Backend\n\ttoolName              string\n\tuseChinese            bool\n\tagentHub              AgentHub\n\tmodelHub              ModelHub\n\tcustomToolDescription ToolDescriptionFunc\n}\n\ntype descriptionTemplateHelper struct {\n\tMatters []FrontMatter\n}\n\nfunc (s *skillTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\tskills, err := s.b.List(ctx)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to list skills: %w\", err)\n\t}\n\n\tvar fullDesc string\n\tif s.customToolDescription != nil {\n\t\tfullDesc = s.customToolDescription(ctx, skills)\n\t} else {\n\t\tdesc, err := renderToolDescription(skills)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to render skill tool description: %w\", err)\n\t\t}\n\n\t\tdescBase := internal.SelectPrompt(internal.I18nPrompts{\n\t\t\tEnglish: toolDescriptionBase,\n\t\t\tChinese: toolDescriptionBaseChinese,\n\t\t})\n\t\tfullDesc = descBase + desc\n\t}\n\n\tparamDesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: \"The skill name (no arguments). E.g., \\\"pdf\\\" or \\\"xlsx\\\"\",\n\t\tChinese: \"Skill 名称（无需其他参数）。例如：\\\"pdf\\\" 或 \\\"xlsx\\\"\",\n\t})\n\n\treturn &schema.ToolInfo{\n\t\tName: s.toolName,\n\t\tDesc: fullDesc,\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"skill\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     paramDesc,\n\t\t\t\tRequired: true,\n\t\t\t},\n\t\t}),\n\t}, nil\n}\n\ntype inputArguments struct {\n\tSkill string `json:\"skill\"`\n}\n\nfunc (s *skillTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\targs := &inputArguments{}\n\terr := json.Unmarshal([]byte(argumentsInJSON), args)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to unmarshal arguments: %w\", err)\n\t}\n\tskill, err := s.b.Get(ctx, args.Skill)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to get skill: %w\", err)\n\t}\n\n\tswitch skill.Context {\n\tcase ContextModeForkWithContext:\n\t\treturn s.runAgentMode(ctx, skill, true)\n\tcase ContextModeFork:\n\t\treturn s.runAgentMode(ctx, skill, false)\n\tdefault:\n\t\tif skill.Model != \"\" {\n\t\t\ts.setActiveModel(ctx, skill.Model)\n\t\t}\n\t\treturn s.buildSkillResult(skill)\n\t}\n}\n\nfunc (s *skillTool) setActiveModel(ctx context.Context, modelName string) {\n\t_ = adk.SetRunLocalValue(ctx, activeModelKey, modelName)\n}\n\nfunc (s *skillTool) buildSkillResult(skill Skill) (string, error) {\n\tresultFmt := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: toolResult,\n\t\tChinese: toolResultChinese,\n\t})\n\tcontentFmt := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: userContent,\n\t\tChinese: userContentChinese,\n\t})\n\n\treturn fmt.Sprintf(resultFmt, skill.Name) + fmt.Sprintf(contentFmt, skill.BaseDirectory, skill.Content), nil\n}\n\nfunc (s *skillTool) runAgentMode(ctx context.Context, skill Skill, forkHistory bool) (string, error) {\n\tif s.agentHub == nil {\n\t\treturn \"\", fmt.Errorf(\"skill '%s' requires context:%s but AgentHub is not configured\", skill.Name, skill.Context)\n\t}\n\n\topts := &AgentHubOptions{}\n\tif skill.Model != \"\" {\n\t\tif s.modelHub == nil {\n\t\t\treturn \"\", fmt.Errorf(\"skill '%s' requires model '%s' but ModelHub is not configured\", skill.Name, skill.Model)\n\t\t}\n\t\tm, err := s.modelHub.Get(ctx, skill.Model)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to get model '%s' from ModelHub: %w\", skill.Model, err)\n\t\t}\n\t\topts.Model = m\n\t}\n\n\tagent, err := s.agentHub.Get(ctx, skill.Agent, opts)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to get agent '%s' from AgentHub: %w\", skill.Agent, err)\n\t}\n\n\tvar messages []adk.Message\n\tskillContent, err := s.buildSkillResult(skill)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to build skill result: %w\", err)\n\t}\n\n\tif forkHistory {\n\t\tmessages, err = s.getMessagesFromState(ctx)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to get messages from state: %w\", err)\n\t\t}\n\t\ttoolCallID := compose.GetToolCallID(ctx)\n\t\tmessages = append(messages, schema.ToolMessage(skillContent, toolCallID))\n\t} else {\n\t\tmessages = []adk.Message{schema.UserMessage(skillContent)}\n\t}\n\n\tinput := &adk.AgentInput{\n\t\tMessages:        messages,\n\t\tEnableStreaming: false,\n\t}\n\n\titer := agent.Run(ctx, input)\n\n\tvar results []string\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\n\t\tif event.Err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to run agent event: %w\", event.Err)\n\t\t}\n\n\t\tif event.Output == nil || event.Output.MessageOutput == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tmsg, msgErr := event.Output.MessageOutput.GetMessage()\n\t\tif msgErr != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to get message from event: %w\", msgErr)\n\t\t}\n\n\t\tif msg != nil && msg.Content != \"\" {\n\t\t\tresults = append(results, msg.Content)\n\t\t}\n\t}\n\n\tresultFmt := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: subAgentResultFormat,\n\t\tChinese: subAgentResultFormatChinese,\n\t})\n\n\treturn fmt.Sprintf(resultFmt, skill.Name, strings.Join(results, \"\\n\")), nil\n}\n\nfunc (s *skillTool) getMessagesFromState(ctx context.Context) ([]adk.Message, error) {\n\tvar messages []adk.Message\n\terr := compose.ProcessState(ctx, func(_ context.Context, st *adk.State) error {\n\t\tmessages = make([]adk.Message, len(st.Messages))\n\t\tcopy(messages, st.Messages)\n\t\treturn nil\n\t})\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to process state: %w\", err)\n\t}\n\treturn messages, nil\n}\n\nfunc renderToolDescription(matters []FrontMatter) (string, error) {\n\ttpl, err := template.New(\"skills\").Parse(toolDescriptionTemplate)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tvar buf bytes.Buffer\n\terr = tpl.Execute(&buf, descriptionTemplateHelper{Matters: matters})\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn buf.String(), nil\n}\n"
  },
  {
    "path": "adk/middlewares/skill/skill_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage skill\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype inMemoryBackend struct {\n\tm []Skill\n}\n\nfunc (i *inMemoryBackend) List(ctx context.Context) ([]FrontMatter, error) {\n\tmatters := make([]FrontMatter, 0, len(i.m))\n\tfor _, skill := range i.m {\n\t\tmatters = append(matters, skill.FrontMatter)\n\t}\n\treturn matters, nil\n}\n\nfunc (i *inMemoryBackend) Get(ctx context.Context, name string) (Skill, error) {\n\tfor _, skill := range i.m {\n\t\tif skill.Name == name {\n\t\t\treturn skill, nil\n\t\t}\n\t}\n\treturn Skill{}, errors.New(\"skill not found\")\n}\n\nfunc TestTool(t *testing.T) {\n\tbackend := &inMemoryBackend{m: []Skill{\n\t\t{\n\t\t\tFrontMatter: FrontMatter{\n\t\t\t\tName:        \"name1\",\n\t\t\t\tDescription: \"desc1\",\n\t\t\t},\n\t\t\tContent:       \"content1\",\n\t\t\tBaseDirectory: \"basedir1\",\n\t\t},\n\t\t{\n\t\t\tFrontMatter: FrontMatter{\n\t\t\t\tName:        \"name2\",\n\t\t\t\tDescription: \"desc2\",\n\t\t\t},\n\t\t\tContent:       \"content2\",\n\t\t\tBaseDirectory: \"basedir2\",\n\t\t},\n\t}}\n\n\tctx := context.Background()\n\tm, err := New(ctx, &Config{Backend: backend})\n\tassert.NoError(t, err)\n\tassert.Len(t, m.AdditionalTools, 1)\n\n\tto := m.AdditionalTools[0].(tool.InvokableTool)\n\n\tinfo, err := to.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"skill\", info.Name)\n\tdesc := strings.TrimPrefix(info.Desc, toolDescriptionBase)\n\tassert.Equal(t, `\n<available_skills>\n<skill>\n<name>\nname1\n</name>\n<description>\ndesc1\n</description>\n</skill>\n<skill>\n<name>\nname2\n</name>\n<description>\ndesc2\n</description>\n</skill>\n</available_skills>\n`, desc)\n\n\tresult, err := to.InvokableRun(ctx, `{\"skill\": \"name1\"}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, `Launching skill: name1\nBase directory for this skill: basedir1\n\ncontent1`, result)\n\n\t// chinese\n\tinternal.SetLanguage(internal.LanguageChinese)\n\tdefer internal.SetLanguage(internal.LanguageEnglish)\n\tm, err = New(ctx, &Config{Backend: backend})\n\tassert.NoError(t, err)\n\tassert.Len(t, m.AdditionalTools, 1)\n\n\tto = m.AdditionalTools[0].(tool.InvokableTool)\n\n\tinfo, err = to.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"skill\", info.Name)\n\tdesc = strings.TrimPrefix(info.Desc, toolDescriptionBaseChinese)\n\tassert.Equal(t, `\n<available_skills>\n<skill>\n<name>\nname1\n</name>\n<description>\ndesc1\n</description>\n</skill>\n<skill>\n<name>\nname2\n</name>\n<description>\ndesc2\n</description>\n</skill>\n</available_skills>\n`, desc)\n\n\tresult, err = to.InvokableRun(ctx, `{\"skill\": \"name1\"}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, `正在启动 Skill：name1\n此 Skill 的目录：basedir1\n\ncontent1`, result)\n}\n\nfunc TestSkillToolName(t *testing.T) {\n\tctx := context.Background()\n\n\t// default\n\tm, err := New(ctx, &Config{Backend: &inMemoryBackend{m: []Skill{}}})\n\tassert.NoError(t, err)\n\t// instruction\n\tassert.Contains(t, m.AdditionalInstruction, \"'skill'\")\n\t// tool name\n\tinfo, err := m.AdditionalTools[0].Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"skill\", info.Name)\n\n\t// customized\n\tname := \"load_skill\"\n\tm, err = New(ctx, &Config{Backend: &inMemoryBackend{m: []Skill{}}, SkillToolName: &name})\n\tassert.NoError(t, err)\n\tassert.Contains(t, m.AdditionalInstruction, \"'load_skill'\")\n\tinfo, err = m.AdditionalTools[0].Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"load_skill\", info.Name)\n}\n\n// --- Mock types for NewMiddleware tests ---\n\ntype mockModel struct {\n\tmodel.ToolCallingChatModel\n\tname string\n}\n\ntype mockModelHub struct {\n\tmodels map[string]model.ToolCallingChatModel\n}\n\nfunc (h *mockModelHub) Get(_ context.Context, name string) (model.ToolCallingChatModel, error) {\n\tm, ok := h.models[name]\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"model not found: %s\", name)\n\t}\n\treturn m, nil\n}\n\ntype mockAgent struct {\n\tevents []*adk.AgentEvent\n}\n\nfunc (a *mockAgent) Name(_ context.Context) string        { return \"mock-agent\" }\nfunc (a *mockAgent) Description(_ context.Context) string  { return \"mock agent for testing\" }\nfunc (a *mockAgent) Run(_ context.Context, _ *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\titer, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tgo func() {\n\t\tdefer gen.Close()\n\t\tfor _, e := range a.events {\n\t\t\tgen.Send(e)\n\t\t}\n\t}()\n\treturn iter\n}\n\ntype mockAgentHub struct {\n\tagents     map[string]adk.Agent\n\tlastOpts   *AgentHubOptions\n\tdefaultAgent adk.Agent\n}\n\nfunc (h *mockAgentHub) Get(_ context.Context, name string, opts *AgentHubOptions) (adk.Agent, error) {\n\th.lastOpts = opts\n\tif name == \"\" && h.defaultAgent != nil {\n\t\treturn h.defaultAgent, nil\n\t}\n\ta, ok := h.agents[name]\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"agent not found: %s\", name)\n\t}\n\treturn a, nil\n}\n\ntype errorBackend struct {\n\tlistErr error\n\tgetErr  error\n}\n\nfunc (b *errorBackend) List(_ context.Context) ([]FrontMatter, error) {\n\treturn nil, b.listErr\n}\nfunc (b *errorBackend) Get(_ context.Context, _ string) (Skill, error) {\n\treturn Skill{}, b.getErr\n}\n\n// --- NewMiddleware tests ---\n\nfunc TestNewMiddleware(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"nil config returns error\", func(t *testing.T) {\n\t\thandler, err := NewMiddleware(ctx, nil)\n\t\tassert.Nil(t, handler)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"config is required\")\n\t})\n\n\tt.Run(\"nil backend returns error\", func(t *testing.T) {\n\t\thandler, err := NewMiddleware(ctx, &Config{})\n\t\tassert.Nil(t, handler)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"backend is required\")\n\t})\n\n\tt.Run(\"valid config succeeds\", func(t *testing.T) {\n\t\tbackend := &inMemoryBackend{m: []Skill{}}\n\t\thandler, err := NewMiddleware(ctx, &Config{Backend: backend})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, handler)\n\t})\n\n\tt.Run(\"custom tool name\", func(t *testing.T) {\n\t\tbackend := &inMemoryBackend{m: []Skill{\n\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Description: \"d1\"}, Content: \"c1\"},\n\t\t}}\n\t\tname := \"load_skill\"\n\t\thandler, err := NewMiddleware(ctx, &Config{Backend: backend, SkillToolName: &name})\n\t\trequire.NoError(t, err)\n\n\t\th := handler.(*skillHandler)\n\t\tassert.Contains(t, h.instruction, \"'load_skill'\")\n\t\tassert.Equal(t, \"load_skill\", h.tool.toolName)\n\t})\n\n\tt.Run(\"custom system prompt\", func(t *testing.T) {\n\t\tbackend := &inMemoryBackend{m: []Skill{}}\n\t\thandler, err := NewMiddleware(ctx, &Config{\n\t\t\tBackend: backend,\n\t\t\tCustomSystemPrompt: func(_ context.Context, toolName string) string {\n\t\t\t\treturn \"custom prompt for \" + toolName\n\t\t\t},\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\th := handler.(*skillHandler)\n\t\tassert.Equal(t, \"custom prompt for skill\", h.instruction)\n\t})\n\n\tt.Run(\"custom tool description\", func(t *testing.T) {\n\t\tbackend := &inMemoryBackend{m: []Skill{\n\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Description: \"d1\"}, Content: \"c1\"},\n\t\t}}\n\t\thandler, err := NewMiddleware(ctx, &Config{\n\t\t\tBackend: backend,\n\t\t\tCustomToolDescription: func(_ context.Context, skills []FrontMatter) string {\n\t\t\t\treturn fmt.Sprintf(\"custom desc with %d skills\", len(skills))\n\t\t\t},\n\t\t})\n\t\trequire.NoError(t, err)\n\n\t\th := handler.(*skillHandler)\n\t\tinfo, err := h.tool.Info(ctx)\n\t\trequire.NoError(t, err)\n\t\tassert.Equal(t, \"custom desc with 1 skills\", info.Desc)\n\t})\n}\n\nfunc TestBeforeAgent(t *testing.T) {\n\tctx := context.Background()\n\tbackend := &inMemoryBackend{m: []Skill{\n\t\t{FrontMatter: FrontMatter{Name: \"s1\", Description: \"d1\"}, Content: \"c1\"},\n\t}}\n\thandler, err := NewMiddleware(ctx, &Config{Backend: backend})\n\trequire.NoError(t, err)\n\n\trunCtx := &adk.ChatModelAgentContext{\n\t\tInstruction: \"base instruction\",\n\t\tTools:       []tool.BaseTool{},\n\t}\n\t_, newRunCtx, err := handler.BeforeAgent(ctx, runCtx)\n\tassert.NoError(t, err)\n\tassert.Contains(t, newRunCtx.Instruction, \"base instruction\")\n\tassert.Contains(t, newRunCtx.Instruction, \"Skills System\")\n\tassert.Len(t, newRunCtx.Tools, 1)\n\n\t// verify the added tool is the skill tool\n\tinfo, err := newRunCtx.Tools[0].Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"skill\", info.Name)\n}\n\nfunc TestSkillToolInfo(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"list error propagates\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb:        &errorBackend{listErr: errors.New(\"list failed\")},\n\t\t\ttoolName: \"skill\",\n\t\t}\n\t\tinfo, err := st.Info(ctx)\n\t\tassert.Nil(t, info)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"list failed\")\n\t})\n\n\tt.Run(\"description contains all skills\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"alpha\", Description: \"desc-alpha\"}},\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"beta\", Description: \"desc-beta\"}},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t}\n\t\tinfo, err := st.Info(ctx)\n\t\trequire.NoError(t, err)\n\t\tassert.Contains(t, info.Desc, \"alpha\")\n\t\tassert.Contains(t, info.Desc, \"desc-alpha\")\n\t\tassert.Contains(t, info.Desc, \"beta\")\n\t\tassert.Contains(t, info.Desc, \"desc-beta\")\n\t})\n}\n\nfunc TestInvokableRun_InlineMode(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"invalid json returns error\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb:        &inMemoryBackend{m: []Skill{}},\n\t\t\ttoolName: \"skill\",\n\t\t}\n\t\t_, err := st.InvokableRun(ctx, \"not json\")\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"failed to unmarshal\")\n\t})\n\n\tt.Run(\"skill not found returns error\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb:        &inMemoryBackend{m: []Skill{}},\n\t\t\ttoolName: \"skill\",\n\t\t}\n\t\t_, err := st.InvokableRun(ctx, `{\"skill\": \"nonexistent\"}`)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"failed to get skill\")\n\t})\n\n\tt.Run(\"inline mode returns skill content\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{\n\t\t\t\t\tFrontMatter:   FrontMatter{Name: \"pdf\", Description: \"PDF processing\"},\n\t\t\t\t\tContent:       \"Process PDF files here\",\n\t\t\t\t\tBaseDirectory: \"/skills/pdf\",\n\t\t\t\t},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t}\n\t\tresult, err := st.InvokableRun(ctx, `{\"skill\": \"pdf\"}`)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, result, \"pdf\")\n\t\tassert.Contains(t, result, \"/skills/pdf\")\n\t\tassert.Contains(t, result, \"Process PDF files here\")\n\t})\n}\n\nfunc TestInvokableRun_AgentMode(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"fork mode without AgentHub returns error\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Context: ContextModeFork}, Content: \"c1\"},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t}\n\t\t_, err := st.InvokableRun(ctx, `{\"skill\": \"s1\"}`)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"AgentHub is not configured\")\n\t})\n\n\tt.Run(\"fork_with_context mode without AgentHub returns error\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Context: ContextModeForkWithContext}, Content: \"c1\"},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t}\n\t\t_, err := st.InvokableRun(ctx, `{\"skill\": \"s1\"}`)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"AgentHub is not configured\")\n\t})\n\n\tt.Run(\"model specified without ModelHub returns error\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Context: ContextModeFork, Model: \"gpt-4\"}, Content: \"c1\"},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t\tagentHub: &mockAgentHub{},\n\t\t}\n\t\t_, err := st.InvokableRun(ctx, `{\"skill\": \"s1\"}`)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"ModelHub is not configured\")\n\t})\n\n\tt.Run(\"model not found in ModelHub returns error\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Context: ContextModeFork, Model: \"gpt-4\"}, Content: \"c1\"},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t\tagentHub: &mockAgentHub{},\n\t\t\tmodelHub: &mockModelHub{models: map[string]model.ToolCallingChatModel{}},\n\t\t}\n\t\t_, err := st.InvokableRun(ctx, `{\"skill\": \"s1\"}`)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"failed to get model\")\n\t})\n\n\tt.Run(\"agent not found in AgentHub returns error\", func(t *testing.T) {\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Context: ContextModeFork, Agent: \"nonexistent\"}, Content: \"c1\"},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t\tagentHub: &mockAgentHub{agents: map[string]adk.Agent{}},\n\t\t}\n\t\t_, err := st.InvokableRun(ctx, `{\"skill\": \"s1\"}`)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"failed to get agent\")\n\t})\n\n\tt.Run(\"fork mode runs agent and returns result\", func(t *testing.T) {\n\t\tagent := &mockAgent{\n\t\t\tevents: []*adk.AgentEvent{\n\t\t\t\t{\n\t\t\t\t\tOutput: &adk.AgentOutput{\n\t\t\t\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(\"agent response\", nil),\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\thub := &mockAgentHub{defaultAgent: agent}\n\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{\n\t\t\t\t\tFrontMatter:   FrontMatter{Name: \"test-skill\", Context: ContextModeFork},\n\t\t\t\t\tContent:       \"skill content\",\n\t\t\t\t\tBaseDirectory: \"/skills/test\",\n\t\t\t\t},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t\tagentHub: hub,\n\t\t}\n\n\t\tresult, err := st.InvokableRun(ctx, `{\"skill\": \"test-skill\"}`)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, result, \"test-skill\")\n\t\tassert.Contains(t, result, \"agent response\")\n\t\tassert.Contains(t, result, \"completed\")\n\t\t// verify no model was passed in opts\n\t\tassert.NotNil(t, hub.lastOpts)\n\t\tassert.Nil(t, hub.lastOpts.Model)\n\t})\n\n\tt.Run(\"fork mode with model passes model to AgentHub\", func(t *testing.T) {\n\t\tm := &mockModel{name: \"test-model\"}\n\t\tagent := &mockAgent{\n\t\t\tevents: []*adk.AgentEvent{\n\t\t\t\t{\n\t\t\t\t\tOutput: &adk.AgentOutput{\n\t\t\t\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(\"response\", nil),\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\thub := &mockAgentHub{defaultAgent: agent}\n\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{\n\t\t\t\t\tFrontMatter:   FrontMatter{Name: \"s1\", Context: ContextModeFork, Model: \"test-model\"},\n\t\t\t\t\tContent:       \"c1\",\n\t\t\t\t\tBaseDirectory: \"/skills/s1\",\n\t\t\t\t},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t\tagentHub: hub,\n\t\t\tmodelHub: &mockModelHub{models: map[string]model.ToolCallingChatModel{\"test-model\": m}},\n\t\t}\n\n\t\tresult, err := st.InvokableRun(ctx, `{\"skill\": \"s1\"}`)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, result, \"s1\")\n\t\t// verify model was passed\n\t\tassert.NotNil(t, hub.lastOpts)\n\t\tassert.Equal(t, m, hub.lastOpts.Model)\n\t})\n\n\tt.Run(\"agent returns multiple events\", func(t *testing.T) {\n\t\tagent := &mockAgent{\n\t\t\tevents: []*adk.AgentEvent{\n\t\t\t\t{\n\t\t\t\t\tOutput: &adk.AgentOutput{\n\t\t\t\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(\"part1\", nil),\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{Output: nil}, // nil output should be skipped\n\t\t\t\t{\n\t\t\t\t\tOutput: &adk.AgentOutput{\n\t\t\t\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(\"part2\", nil),\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\thub := &mockAgentHub{defaultAgent: agent}\n\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Context: ContextModeFork}, Content: \"c1\", BaseDirectory: \"/d\"},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t\tagentHub: hub,\n\t\t}\n\n\t\tresult, err := st.InvokableRun(ctx, `{\"skill\": \"s1\"}`)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, result, \"part1\")\n\t\tassert.Contains(t, result, \"part2\")\n\t})\n\n\tt.Run(\"agent returns empty content events\", func(t *testing.T) {\n\t\tagent := &mockAgent{\n\t\t\tevents: []*adk.AgentEvent{\n\t\t\t\t{\n\t\t\t\t\tOutput: &adk.AgentOutput{\n\t\t\t\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(\"\", nil),\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\thub := &mockAgentHub{defaultAgent: agent}\n\n\t\tst := &skillTool{\n\t\t\tb: &inMemoryBackend{m: []Skill{\n\t\t\t\t{FrontMatter: FrontMatter{Name: \"s1\", Context: ContextModeFork}, Content: \"c1\", BaseDirectory: \"/d\"},\n\t\t\t}},\n\t\t\ttoolName: \"skill\",\n\t\t\tagentHub: hub,\n\t\t}\n\n\t\tresult, err := st.InvokableRun(ctx, `{\"skill\": \"s1\"}`)\n\t\tassert.NoError(t, err)\n\t\t// result should contain skill name but no extra content\n\t\tassert.Contains(t, result, \"s1\")\n\t})\n}\n"
  },
  {
    "path": "adk/middlewares/summarization/consts.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage summarization\n\nconst (\n\textraKeyContentType = \"_eino_summarization_content_type\"\n)\n\ntype summarizationContentType string\n\nconst (\n\tcontentTypeSummary summarizationContentType = \"summary\"\n)\n\ntype ActionType string\n\nconst (\n\tActionTypeBeforeSummarize ActionType = \"before_summarize\"\n\tActionTypeAfterSummarize  ActionType = \"after_summarize\"\n)\n"
  },
  {
    "path": "adk/middlewares/summarization/customized_action.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage summarization\n\nimport (\n\t\"github.com/cloudwego/eino/adk\"\n)\n\ntype CustomizedAction struct {\n\t// Type is the action type.\n\tType ActionType `json:\"type\"`\n\n\t// Before is set when Type is ActionTypeBeforeSummarize.\n\t// Emitted after trigger condition is met, before calling model to generate summary.\n\tBefore *BeforeSummarizeAction `json:\"before,omitempty\"`\n\n\t// After is set when Type is ActionTypeAfterSummarize.\n\t// Emitted after summarization.\n\tAfter *AfterSummarizeAction `json:\"after,omitempty\"`\n}\n\ntype BeforeSummarizeAction struct {\n\t// Messages is the original state messages before summarization.\n\tMessages []adk.Message `json:\"messages,omitempty\"`\n}\n\ntype AfterSummarizeAction struct {\n\t// Messages is the final state messages after summarization.\n\tMessages []adk.Message `json:\"messages,omitempty\"`\n}\n"
  },
  {
    "path": "adk/middlewares/summarization/prompt.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage summarization\n\nimport (\n\t\"regexp\"\n\n\t\"github.com/cloudwego/eino/adk/internal\"\n)\n\nvar allUserMessagesTagRegex = regexp.MustCompile(`(?s)<all_user_messages>.*</all_user_messages>`)\n\nfunc getSystemInstruction() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: systemInstruction,\n\t\tChinese: systemInstructionZh,\n\t})\n}\n\nfunc getUserSummaryInstruction() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: userSummaryInstruction,\n\t\tChinese: userSummaryInstructionZh,\n\t})\n}\n\nfunc getSummaryPreamble() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: summaryPreamble,\n\t\tChinese: summaryPreambleZh,\n\t})\n}\n\nfunc getContinueInstruction() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: continueInstruction,\n\t\tChinese: continueInstructionZh,\n\t})\n}\n\nfunc getTranscriptPathInstruction() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: transcriptPathInstruction,\n\t\tChinese: transcriptPathInstructionZh,\n\t})\n}\n\nfunc getTruncatedMarkerFormat() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: truncatedMarkerFormat,\n\t\tChinese: truncatedMarkerFormatZh,\n\t})\n}\n\nfunc getUserMessagesReplacedNote() string {\n\treturn internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: userMessagesReplacedNote,\n\t\tChinese: userMessagesReplacedNoteZh,\n\t})\n}\n\nconst systemInstruction = `You are a helpful AI assistant tasked with summarizing conversations.`\n\nconst systemInstructionZh = `你是一个负责总结对话的 AI 助手。`\n\nconst userSummaryInstruction = `Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions.\nThis summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context.\n\nBefore providing your final summary, wrap your analysis in <analysis> tags to organize your thoughts and ensure you've covered all necessary points. In your analysis process:\n\n1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify:\n   - The user's explicit requests and intents\n   - Your approach to addressing the user's requests\n   - Key decisions, technical concepts and code patterns\n   - Specific details like:\n     - file names\n     - full code snippets\n     - function signatures\n     - file edits\n  - Errors that you ran into and how you fixed them\n  - Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.\n2. Double-check for technical accuracy and completeness, addressing each required element thoroughly.\n\nYour summary should include the following sections:\n\n1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail\n2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.\n3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.\n4. Errors and fixes: List all errors that you ran into, and how you fixed them. Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.\n5. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.\n6. All user messages: List ALL user messages that are not tool results, and wrap them in the <all_user_messages>...</all_user_messages> block. These are critical for understanding the users' feedback and changing intent.\n7. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.\n8. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.\n9. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's most recent explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests or really old requests that were already completed without confirming with the user first.\n                       If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation.\n\nHere's an example of how your output should be structured:\n\n<example>\n<analysis>\n[Your thought process, ensuring all points are covered thoroughly and accurately]\n</analysis>\n\n<summary>\n1. Primary Request and Intent:\n   [Detailed description]\n\n2. Key Technical Concepts:\n   - [Concept 1]\n   - [Concept 2]\n   - [...]\n\n3. Files and Code Sections:\n   - [File Name 1]\n      - [Summary of why this file is important]\n      - [Summary of the changes made to this file, if any]\n      - [Important Code Snippet]\n   - [File Name 2]\n      - [Important Code Snippet]\n   - [...]\n\n4. Errors and fixes:\n    - [Detailed description of error 1]:\n      - [How you fixed the error]\n      - [User feedback on the error if any]\n    - [...]\n\n5. Problem Solving:\n   [Description of solved problems and ongoing troubleshooting]\n\n6. All user messages: \n<all_user_messages>\n    - [Detailed non tool use user message]\n    - [...]\n</all_user_messages>\n\n7. Pending Tasks:\n   - [Task 1]\n   - [Task 2]\n   - [...]\n\n8. Current Work:\n   [Precise description of current work]\n\n9. Optional Next Step:\n   [Optional Next step to take]\n\n</summary>\n</example>\n\nPlease provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response. \n\nThere may be additional summarization instructions provided in the included context. If so, remember to follow these instructions when creating the above summary. Examples of instructions include:\n<example>\n## Compact Instructions\nWhen summarizing the conversation focus on typescript code changes and also remember the mistakes you made and how you fixed them.\n</example>\n\n<example>\n# Summary instructions\nWhen you are using compact - please focus on test output and code changes. Include file reads verbatim.\n</example>\n\n\nIMPORTANT: Do NOT use any tools. You MUST respond with ONLY the <summary>...</summary> block as your text output.\n`\n\nconst userSummaryInstructionZh = `你的任务是对目前为止的对话创建一份详细的总结，需要密切关注用户的明确请求和你之前的操作。\n这份总结应该全面捕捉技术细节、代码模式和架构决策，以确保继续开发工作时不丢失上下文。\n\n在提供最终总结之前，请将你的分析过程包裹在 <analysis> 标签中，以组织思路并确保涵盖所有必要的要点。在分析过程中：\n\n1. 按时间顺序分析对话中的每条消息和每个部分。对于每个部分，需要全面识别：\n   - 用户的明确请求和意图\n   - 你处理用户请求的方法\n   - 关键决策、技术概念和代码模式\n   - 具体细节，例如：\n     - 文件名\n     - 完整代码片段\n     - 函数签名\n     - 文件编辑\n   - 你遇到的错误以及如何修复它们\n   - 特别注意你收到的具体用户反馈，尤其是用户要求你以不同方式处理的情况\n2. 仔细检查技术准确性和完整性，彻底处理每个必需的元素。\n\n你的总结应包含以下部分：\n\n1. 主要请求和意图：详细捕捉用户所有的明确请求和意图\n2. 关键技术概念：列出讨论过的所有重要技术概念、技术和框架\n3. 文件和代码部分：列举检查、修改或创建的具体文件和代码部分。特别注意最近的消息，在适用的地方包含完整的代码片段，并总结为什么这个文件的读取或编辑很重要\n4. 错误和修复：列出你遇到的所有错误以及如何修复它们。特别注意你收到的具体用户反馈，尤其是用户要求你以不同方式处理的情况\n5. 问题解决：记录已解决的问题和任何正在进行的故障排除工作\n6. 所有用户消息：列出所有非工具结果的用户消息，并将它们包裹在 <all_user_messages>...</all_user_messages> 块中。这些对于理解用户的反馈和变化的意图至关重要\n7. 待处理任务：列出明确要求你处理的任何待处理任务\n8. 当前工作：详细描述在此总结请求之前正在进行的工作，特别注意用户和助手的最近消息。在适用的地方包含文件名和代码片段\n9. 可选的下一步：列出与你最近工作相关的下一步操作。重要提示：确保这一步与用户最近的明确请求以及你在此总结请求之前正在处理的任务直接相关。如果你的上一个任务已经完成，则只有在与用户请求明确相关时才列出下一步。不要在未与用户确认的情况下开始处理无关的请求或已经完成的旧请求。\n   如果有下一步，请包含最近对话中的直接引用，准确显示你正在处理的任务以及你停止的位置。这应该是逐字引用，以确保任务理解不会偏离。\n\n以下是输出结构的示例：\n\n<example>\n<analysis>\n[你的思考过程，确保全面准确地涵盖所有要点]\n</analysis>\n\n<summary>\n1. 主要请求和意图：\n   [详细描述]\n\n2. 关键技术概念：\n   - [概念 1]\n   - [概念 2]\n   - [...]\n\n3. 文件和代码部分：\n   - [文件名 1]\n      - [为什么这个文件重要的总结]\n      - [对这个文件所做更改的总结（如有）]\n      - [重要代码片段]\n   - [文件名 2]\n      - [重要代码片段]\n   - [...]\n\n4. 错误和修复：\n    - [错误 1 的详细描述]：\n      - [如何修复该错误]\n      - [用户对该错误的反馈（如有）]\n    - [...]\n\n5. 问题解决：\n   [已解决问题和正在进行的故障排除的描述]\n\n6. 所有用户消息：\n<all_user_messages>\n    - [详细的非工具使用用户消息]\n    - [...]\n</all_user_messages>\n\n7. 待处理任务：\n   - [任务 1]\n   - [任务 2]\n   - [...]\n\n8. 当前工作：\n   [当前工作的精确描述]\n\n9. 可选的下一步：\n   [可选的下一步操作]\n\n</summary>\n</example>\n\n请根据目前为止的对话提供你的总结，遵循此结构并确保回复的精确性和全面性。\n\n上下文中可能包含额外的总结指令。如果有，请在创建上述总结时记得遵循这些指令。指令示例包括：\n<example>\n## 压缩指令\n在总结对话时，重点关注 typescript 代码更改，并记住你犯的错误以及如何修复它们。\n</example>\n\n<example>\n# 总结指令\n当你使用压缩时，请重点关注测试输出和代码更改。逐字包含文件读取内容。\n</example>\n\n\n重要提示：不要使用任何工具。你必须只以 <summary>...</summary> 块作为文本输出进行回复。\n`\n\nconst summaryPreamble = `This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.`\n\nconst summaryPreambleZh = `此会话延续自此前一段因上下文耗尽而终止的对话。以下总结概述了此前对话的内容。`\n\nconst continueInstruction = `Please continue the conversation from where we left it off without asking the user any further questions. Continue with the last task that you were asked to work on.`\n\nconst continueInstructionZh = `请从我们中断的地方继续对话，无需向用户提出任何进一步的问题。继续完成先前指令中未完成的任务。`\n\nconst transcriptPathInstruction = `If you need specific details from before compaction (like exact code snippets, error messages, or content you generated), read the full transcript at: %s`\n\nconst transcriptPathInstructionZh = `如果你需要压缩之前的具体细节（如精确的代码片段、错误消息或你生成的内容），完整的对话记录位于：%s`\n\nconst truncatedMarkerFormat = \"…%d characters truncated…\"\n\nconst truncatedMarkerFormatZh = \"…已截断 %d 个字符…\"\n\nconst userMessagesReplacedNote = \"Some earlier user messages have been cleared. Below are the most recent user messages:\"\n\nconst userMessagesReplacedNoteZh = \"部分较早的用户消息已被清除，以下是保留的最近用户消息：\"\n"
  },
  {
    "path": "adk/middlewares/summarization/summarization.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package summarization provides a middleware that automatically summarizes\n// conversation history when token count exceeds the configured threshold.\npackage summarization\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"regexp\"\n\t\"strings\"\n\t\"unicode/utf8\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc init() {\n\tschema.RegisterName[*CustomizedAction](\"_eino_adk_summarization_mw_customized_action\")\n}\n\ntype (\n\tTokenCounterFunc  func(ctx context.Context, input *TokenCounterInput) (int, error)\n\tGenModelInputFunc func(ctx context.Context, defaultSystemInstruction, userInstruction adk.Message, originalMsgs []adk.Message) ([]adk.Message, error)\n\tFinalizeFunc      func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error)\n\tCallbackFunc      func(ctx context.Context, before, after adk.ChatModelAgentState) error\n)\n\n// Config defines the configuration for the summarization middleware.\ntype Config struct {\n\t// Model is the chat model used to generate summaries.\n\tModel model.BaseChatModel\n\n\t// ModelOptions specifies options passed to the model when generating summaries.\n\t// Optional.\n\tModelOptions []model.Option\n\n\t// TokenCounter calculates the token count for a message.\n\t// Optional. Defaults to a simple estimator (~4 chars/token).\n\tTokenCounter TokenCounterFunc\n\n\t// Trigger specifies the conditions that activate summarization.\n\t// Optional. Defaults to triggering when total tokens exceed 190k.\n\tTrigger *TriggerCondition\n\n\t// EmitInternalEvents indicates whether internal events should be emitted during summarization,\n\t// allowing external observers to track the summarization process.\n\t//\n\t// Event Scoping:\n\t//   - ActionTypeBeforeSummarize: emitted before calling model to generate summary\n\t//   - ActionTypeAfterSummarize: emitted after summary generation completes\n\t// Optional. Defaults to false.\n\tEmitInternalEvents bool\n\n\t// UserInstruction serves as the user-level instruction to guide the model on how to summarize the context.\n\t// It is appended to the message history as a User message.\n\t// If provided, it overrides the default user summarization instruction.\n\t// Optional.\n\tUserInstruction string\n\n\t// TranscriptFilePath is the path to the file containing the full conversation history.\n\t// It is appended to the summary to remind the model where to read the original context.\n\t// Optional but strongly recommended.\n\tTranscriptFilePath string\n\n\t// GenModelInput allows full control over the summarization model input construction.\n\t//\n\t// Parameters:\n\t//   - defaultSystemInstruction: System message defining the model's role\n\t//   - userInstruction: User message with the task instruction\n\t//   - originalMsgs: original complete message list\n\t//\n\t// Typical model input order: systemInstruction -> contextMessages -> userInstruction.\n\t//\n\t// Optional.\n\tGenModelInput GenModelInputFunc\n\n\t// Finalize is called after summary generation. The returned messages are used as the final output.\n\t// Optional.\n\tFinalize FinalizeFunc\n\n\t// Callback is called after Finalize, before exiting the middleware.\n\t// Read-only, do not modify state.\n\t// Optional.\n\tCallback CallbackFunc\n\n\t// PreserveUserMessages controls whether to preserve original user messages in the summary.\n\t// When enabled, replaces the <all_user_messages> section in the model-generated summary\n\t// with recent original user messages from the conversation.\n\t// When disabled, the model-generated content is kept unchanged.\n\t// Optional. Enabled by default.\n\tPreserveUserMessages *PreserveUserMessages\n}\n\ntype TokenCounterInput struct {\n\tMessages []adk.Message\n\tTools    []*schema.ToolInfo\n}\n\n// TriggerCondition specifies when summarization should be activated.\n// Summarization triggers if ANY of the set conditions is met.\ntype TriggerCondition struct {\n\t// ContextTokens triggers summarization when total token count exceeds this threshold.\n\tContextTokens int\n\t// ContextMessages triggers summarization when total messages count exceeds this threshold.\n\tContextMessages int\n}\n\n// PreserveUserMessages controls whether to preserve original user messages in the summary.\ntype PreserveUserMessages struct {\n\tEnabled bool\n\n\t// MaxTokens limits the maximum token count for preserved user messages.\n\t// When set, only the most recent user messages within this limit are preserved.\n\t// Optional. Defaults to 1/3 of TriggerCondition.ContextTokens if not specified.\n\tMaxTokens int\n}\n\n// New creates a summarization middleware that automatically summarizes conversation history\n// when trigger conditions are met.\nfunc New(ctx context.Context, cfg *Config) (adk.ChatModelAgentMiddleware, error) {\n\tif err := cfg.check(); err != nil {\n\t\treturn nil, err\n\t}\n\treturn &middleware{\n\t\tcfg:                          cfg,\n\t\tBaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},\n\t}, nil\n}\n\ntype middleware struct {\n\t*adk.BaseChatModelAgentMiddleware\n\tcfg *Config\n}\n\nfunc (m *middleware) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState,\n\tmtx *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {\n\n\tvar tools []*schema.ToolInfo\n\tif mtx != nil {\n\t\ttools = mtx.Tools\n\t}\n\n\ttriggered, err := m.shouldSummarize(ctx, &TokenCounterInput{\n\t\tMessages: state.Messages,\n\t\tTools:    tools,\n\t})\n\tif err != nil {\n\t\treturn nil, nil, err\n\t}\n\tif !triggered {\n\t\treturn ctx, state, nil\n\t}\n\n\tbeforeState := *state\n\n\tif m.cfg.EmitInternalEvents {\n\t\terr = m.emitEvent(ctx, &CustomizedAction{\n\t\t\tType:   ActionTypeBeforeSummarize,\n\t\t\tBefore: &BeforeSummarizeAction{Messages: state.Messages},\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, nil, err\n\t\t}\n\t}\n\n\tvar (\n\t\tsystemMsgs  []adk.Message\n\t\tcontextMsgs []adk.Message\n\t)\n\n\tfor _, msg := range state.Messages {\n\t\tif msg.Role == schema.System {\n\t\t\tsystemMsgs = append(systemMsgs, msg)\n\t\t} else {\n\t\t\tcontextMsgs = append(contextMsgs, msg)\n\t\t}\n\t}\n\n\tsummary, err := m.summarize(ctx, state.Messages, contextMsgs)\n\tif err != nil {\n\t\treturn nil, nil, err\n\t}\n\n\tsummary, err = m.postProcessSummary(ctx, contextMsgs, summary)\n\tif err != nil {\n\t\treturn nil, nil, err\n\t}\n\n\tif m.cfg.Finalize != nil {\n\t\tstate.Messages, err = m.cfg.Finalize(ctx, state.Messages, summary)\n\t\tif err != nil {\n\t\t\treturn nil, nil, err\n\t\t}\n\t} else {\n\t\tstate.Messages = append(systemMsgs, summary)\n\t}\n\n\tif m.cfg.Callback != nil {\n\t\terr = m.cfg.Callback(ctx, beforeState, *state)\n\t\tif err != nil {\n\t\t\treturn nil, nil, err\n\t\t}\n\t}\n\n\tif m.cfg.EmitInternalEvents {\n\t\terr = m.emitEvent(ctx, &CustomizedAction{\n\t\t\tType:  ActionTypeAfterSummarize,\n\t\t\tAfter: &AfterSummarizeAction{Messages: state.Messages},\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, nil, err\n\t\t}\n\t}\n\n\treturn ctx, state, nil\n}\n\nfunc (m *middleware) shouldSummarize(ctx context.Context, input *TokenCounterInput) (bool, error) {\n\tif m.cfg.Trigger != nil && m.cfg.Trigger.ContextMessages > 0 {\n\t\tif len(input.Messages) > m.cfg.Trigger.ContextMessages {\n\t\t\treturn true, nil\n\t\t}\n\t}\n\ttokens, err := m.countTokens(ctx, input)\n\tif err != nil {\n\t\treturn false, fmt.Errorf(\"failed to count tokens: %w\", err)\n\t}\n\treturn tokens > m.getTriggerContextTokens(), nil\n}\n\nfunc (m *middleware) getTriggerContextTokens() int {\n\tconst defaultTriggerContextTokens = 190000\n\tif m.cfg.Trigger != nil {\n\t\treturn m.cfg.Trigger.ContextTokens\n\t}\n\treturn defaultTriggerContextTokens\n}\n\nfunc (m *middleware) getUserMessageContextTokens() int {\n\tif m.cfg.PreserveUserMessages != nil && m.cfg.PreserveUserMessages.MaxTokens > 0 {\n\t\treturn m.cfg.PreserveUserMessages.MaxTokens\n\t}\n\treturn m.getTriggerContextTokens() / 3\n}\n\nfunc (m *middleware) emitEvent(ctx context.Context, action *CustomizedAction) error {\n\terr := adk.SendEvent(ctx, &adk.AgentEvent{\n\t\tAction: &adk.AgentAction{\n\t\t\tCustomizedAction: action,\n\t\t},\n\t})\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to send internal event: %w\", err)\n\t}\n\treturn nil\n}\n\nfunc (m *middleware) countTokens(ctx context.Context, input *TokenCounterInput) (int, error) {\n\tif m.cfg.TokenCounter != nil {\n\t\treturn m.cfg.TokenCounter(ctx, input)\n\t}\n\treturn defaultTokenCounter(ctx, input)\n}\n\nfunc defaultTokenCounter(ctx context.Context, input *TokenCounterInput) (int, error) {\n\tvar totalTokens int\n\tfor _, msg := range input.Messages {\n\t\ttext := extractTextContent(msg)\n\t\ttotalTokens += estimateTokenCount(text)\n\t}\n\n\tfor _, tl := range input.Tools {\n\t\ttl_ := *tl\n\t\ttl_.Extra = nil\n\t\ttext, err := sonic.MarshalString(tl_)\n\t\tif err != nil {\n\t\t\treturn 0, fmt.Errorf(\"failed to marshal tool info: %w\", err)\n\t\t}\n\n\t\ttotalTokens += estimateTokenCount(text)\n\t}\n\n\treturn totalTokens, nil\n}\n\nfunc estimateTokenCount(text string) int {\n\treturn (len(text) + 3) / 4\n}\n\nfunc (m *middleware) summarize(ctx context.Context, originMsgs, contextMsgs []adk.Message) (adk.Message, error) {\n\tinput, err := m.buildSummarizationModelInput(ctx, originMsgs, contextMsgs)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tresp, err := m.cfg.Model.Generate(ctx, input, m.cfg.ModelOptions...)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to generate summary: %w\", err)\n\t}\n\n\treturn newSummaryMessage(resp.Content), nil\n}\n\nfunc (m *middleware) buildSummarizationModelInput(ctx context.Context, originMsgs, contextMsgs []adk.Message) ([]adk.Message, error) {\n\tuserInstruction := m.cfg.UserInstruction\n\tif userInstruction == \"\" {\n\t\tuserInstruction = getUserSummaryInstruction()\n\t}\n\n\tuserInstructionMsg := &schema.Message{\n\t\tRole:    schema.User,\n\t\tContent: userInstruction,\n\t}\n\n\tsysInstructionMsg := &schema.Message{\n\t\tRole:    schema.System,\n\t\tContent: getSystemInstruction(),\n\t}\n\n\tif m.cfg.GenModelInput != nil {\n\t\tinput, err := m.cfg.GenModelInput(ctx, sysInstructionMsg, userInstructionMsg, originMsgs)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to generate model input: %w\", err)\n\t\t}\n\t\treturn input, nil\n\t}\n\n\tinput := make([]adk.Message, 0, len(contextMsgs)+2)\n\tinput = append(input, sysInstructionMsg)\n\tinput = append(input, contextMsgs...)\n\tinput = append(input, userInstructionMsg)\n\n\treturn input, nil\n}\n\nfunc newSummaryMessage(content string) *schema.Message {\n\tsummary := &schema.Message{\n\t\tRole:    schema.User,\n\t\tContent: content,\n\t}\n\tsetContentType(summary, contentTypeSummary)\n\treturn summary\n}\n\nfunc (m *middleware) postProcessSummary(ctx context.Context, messages []adk.Message, summary adk.Message) (adk.Message, error) {\n\tif m.cfg.PreserveUserMessages == nil || m.cfg.PreserveUserMessages.Enabled {\n\t\tmaxUserMsgTokens := m.getUserMessageContextTokens()\n\t\tcontent, err := m.replaceUserMessagesInSummary(ctx, messages, summary.Content, maxUserMsgTokens)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to replace user messages in summary: %w\", err)\n\t\t}\n\t\tsummary.Content = content\n\t}\n\n\tif path := m.cfg.TranscriptFilePath; path != \"\" {\n\t\tsummary.Content = appendSection(summary.Content, fmt.Sprintf(getTranscriptPathInstruction(), path))\n\t}\n\n\tsummary.Content = appendSection(getSummaryPreamble(), summary.Content)\n\n\tsummary.UserInputMultiContent = []schema.MessageInputPart{\n\t\t{\n\t\t\tType: schema.ChatMessagePartTypeText,\n\t\t\tText: summary.Content,\n\t\t},\n\t\t{\n\t\t\tType: schema.ChatMessagePartTypeText,\n\t\t\tText: getContinueInstruction(),\n\t\t},\n\t}\n\n\tsummary.Content = \"\"\n\n\treturn summary, nil\n}\n\nfunc (m *middleware) replaceUserMessagesInSummary(ctx context.Context, messages []adk.Message, summary string, contextTokens int) (string, error) {\n\tvar userMsgs []adk.Message\n\tfor _, msg := range messages {\n\t\tif typ, ok := getContentType(msg); ok && typ == contentTypeSummary {\n\t\t\tcontinue\n\t\t}\n\t\tif msg.Role == schema.User {\n\t\t\tuserMsgs = append(userMsgs, msg)\n\t\t}\n\t}\n\n\tif len(userMsgs) == 0 {\n\t\treturn summary, nil\n\t}\n\n\tvar selected []adk.Message\n\tif len(userMsgs) == 1 {\n\t\tselected = userMsgs\n\t} else {\n\t\tvar totalTokens int\n\t\tfor i := len(userMsgs) - 1; i >= 0; i-- {\n\t\t\tmsg := userMsgs[i]\n\n\t\t\ttokens, err := m.countTokens(ctx, &TokenCounterInput{\n\t\t\t\tMessages: []adk.Message{msg},\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"failed to count tokens: %w\", err)\n\t\t\t}\n\n\t\t\tremaining := contextTokens - totalTokens\n\t\t\tif tokens <= remaining {\n\t\t\t\ttotalTokens += tokens\n\t\t\t\tselected = append(selected, msg)\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\ttrimmedMsg := defaultTrimUserMessage(msg, remaining)\n\t\t\tif trimmedMsg != nil {\n\t\t\t\tselected = append(selected, trimmedMsg)\n\t\t\t}\n\n\t\t\tbreak\n\t\t}\n\n\t\tfor i, j := 0, len(selected)-1; i < j; i, j = i+1, j-1 {\n\t\t\tselected[i], selected[j] = selected[j], selected[i]\n\t\t}\n\t}\n\n\tvar msgLines []string\n\tfor _, msg := range selected {\n\t\ttext := extractTextContent(msg)\n\t\tif text != \"\" {\n\t\t\tmsgLines = append(msgLines, \"    - \"+text)\n\t\t}\n\t}\n\tuserMsgsText := strings.Join(msgLines, \"\\n\")\n\n\tif userMsgsText == \"\" {\n\t\treturn summary, nil\n\t}\n\n\tlastMatch := findLastMatch(allUserMessagesTagRegex, summary)\n\tif lastMatch == nil {\n\t\treturn summary, nil\n\t}\n\n\tvar replacement string\n\tif len(selected) < len(userMsgs) {\n\t\treplacement = \"<all_user_messages>\\n\" + getUserMessagesReplacedNote() + \"\\n\" + userMsgsText + \"\\n</all_user_messages>\"\n\t} else {\n\t\treplacement = \"<all_user_messages>\\n\" + userMsgsText + \"\\n</all_user_messages>\"\n\t}\n\n\tcontent := summary[:lastMatch[0]] + replacement + summary[lastMatch[1]:]\n\n\treturn content, nil\n}\n\nfunc findLastMatch(re *regexp.Regexp, s string) []int {\n\tmatches := re.FindAllStringIndex(s, -1)\n\tif len(matches) == 0 {\n\t\treturn nil\n\t}\n\treturn matches[len(matches)-1]\n}\n\nfunc appendSection(base, section string) string {\n\tif base == \"\" {\n\t\treturn section\n\t}\n\tif section == \"\" {\n\t\treturn base\n\t}\n\treturn base + \"\\n\\n\" + section\n}\n\nfunc defaultTrimUserMessage(msg adk.Message, remainingTokens int) adk.Message {\n\tif remainingTokens <= 0 {\n\t\treturn nil\n\t}\n\n\ttextContent := extractTextContent(msg)\n\tif len(textContent) == 0 {\n\t\treturn nil\n\t}\n\n\ttrimmed := truncateTextByChars(textContent)\n\tif trimmed == \"\" {\n\t\treturn nil\n\t}\n\n\treturn &schema.Message{\n\t\tRole:    schema.User,\n\t\tContent: trimmed,\n\t}\n}\n\nfunc truncateTextByChars(text string) string {\n\tconst maxRunes = 2000\n\n\tif text == \"\" {\n\t\treturn \"\"\n\t}\n\n\tif utf8.RuneCountInString(text) <= maxRunes {\n\t\treturn text\n\t}\n\n\thalfRunes := maxRunes / 2\n\trunes := []rune(text)\n\ttotalRunes := len(runes)\n\n\tprefix := string(runes[:halfRunes])\n\tsuffix := string(runes[totalRunes-halfRunes:])\n\tremovedChars := totalRunes - maxRunes\n\n\tmarker := fmt.Sprintf(getTruncatedMarkerFormat(), removedChars)\n\n\treturn prefix + marker + suffix\n}\n\nfunc extractTextContent(msg adk.Message) string {\n\tif msg == nil {\n\t\treturn \"\"\n\t}\n\tif msg.Content != \"\" {\n\t\treturn msg.Content\n\t}\n\n\tvar sb strings.Builder\n\tfor _, part := range msg.UserInputMultiContent {\n\t\tif part.Type == schema.ChatMessagePartTypeText && part.Text != \"\" {\n\t\t\tif sb.Len() > 0 {\n\t\t\t\tsb.WriteString(\"\\n\")\n\t\t\t}\n\t\t\tsb.WriteString(part.Text)\n\t\t}\n\t}\n\n\treturn sb.String()\n}\n\nfunc (c *Config) check() error {\n\tif c == nil {\n\t\treturn fmt.Errorf(\"config is required\")\n\t}\n\tif c.Model == nil {\n\t\treturn fmt.Errorf(\"model is required\")\n\t}\n\tif c.Trigger != nil {\n\t\tif err := c.Trigger.check(); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\treturn nil\n}\n\nfunc (c *TriggerCondition) check() error {\n\tif c.ContextTokens < 0 {\n\t\treturn fmt.Errorf(\"trigger.ContextTokens must be non-negative\")\n\t}\n\tif c.ContextMessages < 0 {\n\t\treturn fmt.Errorf(\"trigger.ContextMessages must be non-negative\")\n\t}\n\tif c.ContextTokens == 0 && c.ContextMessages == 0 {\n\t\treturn fmt.Errorf(\"at least one of trigger.ContextTokens or trigger.ContextMessages must be non-negative\")\n\t}\n\treturn nil\n}\n\nfunc setContentType(msg adk.Message, ct summarizationContentType) {\n\tsetExtra(msg, extraKeyContentType, string(ct))\n}\n\nfunc getContentType(msg adk.Message) (summarizationContentType, bool) {\n\tct, ok := getExtra[string](msg, extraKeyContentType)\n\tif !ok {\n\t\treturn \"\", false\n\t}\n\treturn summarizationContentType(ct), true\n}\n\nfunc setExtra(msg adk.Message, key string, value any) {\n\tif msg.Extra == nil {\n\t\tmsg.Extra = make(map[string]any)\n\t}\n\tmsg.Extra[key] = value\n}\n\nfunc getExtra[T any](msg adk.Message, key string) (T, bool) {\n\tvar zero T\n\tif msg == nil || msg.Extra == nil {\n\t\treturn zero, false\n\t}\n\tv, ok := msg.Extra[key].(T)\n\tif !ok {\n\t\treturn zero, false\n\t}\n\treturn v, true\n}\n"
  },
  {
    "path": "adk/middlewares/summarization/summarization_test.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage summarization\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestNew(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"valid config\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\n\t\tcfg := &Config{\n\t\t\tModel: cm,\n\t\t}\n\n\t\tmw, err := New(ctx, cfg)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, mw)\n\t})\n\n\tt.Run(\"nil config returns error\", func(t *testing.T) {\n\t\tmw, err := New(ctx, nil)\n\t\tassert.Error(t, err)\n\t\tassert.Nil(t, mw)\n\t})\n\n\tt.Run(\"nil model returns error\", func(t *testing.T) {\n\t\tmw, err := New(ctx, &Config{})\n\t\tassert.Error(t, err)\n\t\tassert.Nil(t, mw)\n\t})\n}\n\nfunc TestMiddlewareBeforeModelRewriteState(t *testing.T) {\n\tctx := context.Background()\n\tmtx := &adk.ModelContext{}\n\n\tt.Run(\"no summarization when under threshold\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel:   cm,\n\t\t\t\tTrigger: &TriggerCondition{ContextTokens: 1000},\n\t\t\t},\n\t\t\tBaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},\n\t\t}\n\n\t\tstate := &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t\tschema.AssistantMessage(\"hi\", nil),\n\t\t\t},\n\t\t}\n\n\t\t_, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, newState.Messages, 2)\n\t\tassert.Equal(t, \"hello\", newState.Messages[0].Content)\n\t})\n\n\tt.Run(\"summarization triggered when over threshold\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(&schema.Message{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tContent: \"Summary content\",\n\t\t\t}, nil).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel:   cm,\n\t\t\t\tTrigger: &TriggerCondition{ContextTokens: 10},\n\t\t\t},\n\t\t\tBaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},\n\t\t}\n\n\t\tstate := &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(strings.Repeat(\"a\", 100)),\n\t\t\t\tschema.AssistantMessage(strings.Repeat(\"b\", 100), nil),\n\t\t\t},\n\t\t}\n\n\t\t_, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, newState.Messages, 1)\n\t\tassert.Equal(t, schema.User, newState.Messages[0].Role)\n\t})\n\n\tt.Run(\"preserves system messages after summarization\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) {\n\t\t\t\tfor i, msg := range msgs {\n\t\t\t\t\tif i == 0 {\n\t\t\t\t\t\tassert.Equal(t, schema.System, msg.Role)\n\t\t\t\t\t} else {\n\t\t\t\t\t\tassert.NotEqual(t, schema.System, msg.Role)\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn &schema.Message{\n\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\tContent: \"Summary content\",\n\t\t\t\t}, nil\n\t\t\t}).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel:   cm,\n\t\t\t\tTrigger: &TriggerCondition{ContextTokens: 10},\n\t\t\t},\n\t\t\tBaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},\n\t\t}\n\n\t\tstate := &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.SystemMessage(\"You are a helpful assistant\"),\n\t\t\t\tschema.UserMessage(strings.Repeat(\"a\", 100)),\n\t\t\t\tschema.AssistantMessage(strings.Repeat(\"b\", 100), nil),\n\t\t\t},\n\t\t}\n\n\t\t_, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, newState.Messages, 2)\n\t\tassert.Equal(t, schema.System, newState.Messages[0].Role)\n\t\tassert.Equal(t, \"You are a helpful assistant\", newState.Messages[0].Content)\n\t\tassert.Equal(t, schema.User, newState.Messages[1].Role)\n\t})\n\n\tt.Run(\"preserves multiple system messages\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(&schema.Message{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tContent: \"Summary\",\n\t\t\t}, nil).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel:   cm,\n\t\t\t\tTrigger: &TriggerCondition{ContextTokens: 10},\n\t\t\t},\n\t\t\tBaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},\n\t\t}\n\n\t\tstate := &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.SystemMessage(\"System 1\"),\n\t\t\t\tschema.SystemMessage(\"System 2\"),\n\t\t\t\tschema.UserMessage(strings.Repeat(\"a\", 100)),\n\t\t\t},\n\t\t}\n\n\t\t_, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, newState.Messages, 3)\n\t\tassert.Equal(t, schema.System, newState.Messages[0].Role)\n\t\tassert.Equal(t, \"System 1\", newState.Messages[0].Content)\n\t\tassert.Equal(t, schema.System, newState.Messages[1].Role)\n\t\tassert.Equal(t, \"System 2\", newState.Messages[1].Content)\n\t\tassert.Equal(t, schema.User, newState.Messages[2].Role)\n\t})\n\n\tt.Run(\"custom finalize function\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(&schema.Message{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tContent: \"Summary\",\n\t\t\t}, nil).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel:   cm,\n\t\t\t\tTrigger: &TriggerCondition{ContextTokens: 10},\n\t\t\t\tFinalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) {\n\t\t\t\t\treturn []adk.Message{\n\t\t\t\t\t\tschema.SystemMessage(\"system prompt\"),\n\t\t\t\t\t\tsummary,\n\t\t\t\t\t}, nil\n\t\t\t\t},\n\t\t\t},\n\t\t\tBaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},\n\t\t}\n\n\t\tstate := &adk.ChatModelAgentState{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(strings.Repeat(\"a\", 100)),\n\t\t\t},\n\t\t}\n\n\t\t_, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, newState.Messages, 2)\n\t\tassert.Equal(t, schema.System, newState.Messages[0].Role)\n\t\tassert.Equal(t, \"system prompt\", newState.Messages[0].Content)\n\t})\n\n}\n\nfunc TestMiddlewareShouldSummarize(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"returns true when over messages threshold\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tTrigger: &TriggerCondition{ContextMessages: 1},\n\t\t\t},\n\t\t}\n\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(\"msg1\"),\n\t\t\t\tschema.UserMessage(\"msg2\"),\n\t\t\t},\n\t\t}\n\n\t\ttriggered, err := mw.shouldSummarize(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, triggered)\n\t})\n\n\tt.Run(\"returns false when under messages threshold\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tTrigger: &TriggerCondition{\n\t\t\t\t\tContextMessages: 3,\n\t\t\t\t\tContextTokens:   1000,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(\"msg1\"),\n\t\t\t\tschema.UserMessage(\"msg2\"),\n\t\t\t},\n\t\t}\n\n\t\ttriggered, err := mw.shouldSummarize(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.False(t, triggered)\n\t})\n\n\tt.Run(\"returns true when over threshold\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tTrigger: &TriggerCondition{ContextTokens: 10},\n\t\t\t},\n\t\t}\n\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(strings.Repeat(\"a\", 100)),\n\t\t\t},\n\t\t}\n\n\t\ttriggered, err := mw.shouldSummarize(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, triggered)\n\t})\n\n\tt.Run(\"returns false when under threshold\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tTrigger: &TriggerCondition{ContextTokens: 1000},\n\t\t\t},\n\t\t}\n\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(\"short message\"),\n\t\t\t},\n\t\t}\n\n\t\ttriggered, err := mw.shouldSummarize(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.False(t, triggered)\n\t})\n\n\tt.Run(\"uses default threshold when trigger is nil\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{},\n\t\t}\n\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{\n\t\t\t\tschema.UserMessage(\"short message\"),\n\t\t\t},\n\t\t}\n\n\t\ttriggered, err := mw.shouldSummarize(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.False(t, triggered)\n\t})\n}\n\nfunc TestMiddlewareCountTokens(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"uses custom token counter\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tTokenCounter: func(ctx context.Context, input *TokenCounterInput) (int, error) {\n\t\t\t\t\treturn 42, nil\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{schema.UserMessage(\"test\")},\n\t\t}\n\t\ttokens, err := mw.countTokens(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, 42, tokens)\n\t})\n\n\tt.Run(\"uses default token counter when nil\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{},\n\t\t}\n\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{schema.UserMessage(\"test\")},\n\t\t}\n\t\ttokens, err := mw.countTokens(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, 1, tokens)\n\t})\n\n\tt.Run(\"custom token counter error\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tTokenCounter: func(ctx context.Context, input *TokenCounterInput) (int, error) {\n\t\t\t\t\treturn 0, errors.New(\"token count error\")\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{schema.UserMessage(\"test\")},\n\t\t}\n\t\t_, err := mw.countTokens(ctx, input)\n\t\tassert.Error(t, err)\n\t})\n}\n\nfunc TestExtractTextContent(t *testing.T) {\n\tt.Run(\"extracts from Content field\", func(t *testing.T) {\n\t\tmsg := &schema.Message{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"hello world\",\n\t\t}\n\t\tassert.Equal(t, \"hello world\", extractTextContent(msg))\n\t})\n\n\tt.Run(\"extracts from UserInputMultiContent\", func(t *testing.T) {\n\t\tmsg := &schema.Message{\n\t\t\tRole: schema.User,\n\t\t\tUserInputMultiContent: []schema.MessageInputPart{\n\t\t\t\t{Type: schema.ChatMessagePartTypeText, Text: \"part1\"},\n\t\t\t\t{Type: schema.ChatMessagePartTypeText, Text: \"part2\"},\n\t\t\t},\n\t\t}\n\t\tassert.Equal(t, \"part1\\npart2\", extractTextContent(msg))\n\t})\n\n\tt.Run(\"prefers Content over UserInputMultiContent\", func(t *testing.T) {\n\t\tmsg := &schema.Message{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"content field\",\n\t\t\tUserInputMultiContent: []schema.MessageInputPart{\n\t\t\t\t{Type: schema.ChatMessagePartTypeText, Text: \"multi content\"},\n\t\t\t},\n\t\t}\n\t\tassert.Equal(t, \"content field\", extractTextContent(msg))\n\t})\n}\n\nfunc TestTruncateTextByChars(t *testing.T) {\n\tt.Run(\"returns empty for empty string\", func(t *testing.T) {\n\t\tresult := truncateTextByChars(\"\")\n\t\tassert.Equal(t, \"\", result)\n\t})\n\n\tt.Run(\"returns original if under limit\", func(t *testing.T) {\n\t\tresult := truncateTextByChars(\"short\")\n\t\tassert.Equal(t, \"short\", result)\n\t})\n\n\tt.Run(\"truncates long text\", func(t *testing.T) {\n\t\tlongText := strings.Repeat(\"a\", 3000)\n\t\tresult := truncateTextByChars(longText)\n\t\tassert.Less(t, len(result), len(longText))\n\t\tassert.Contains(t, result, \"truncated\")\n\t})\n\n\tt.Run(\"preserves prefix and suffix\", func(t *testing.T) {\n\t\tlongText := strings.Repeat(\"a\", 1000) + strings.Repeat(\"b\", 1000) + strings.Repeat(\"c\", 1000)\n\t\tresult := truncateTextByChars(longText)\n\t\tassert.True(t, strings.HasPrefix(result, strings.Repeat(\"a\", 1000)))\n\t\tassert.True(t, strings.HasSuffix(result, strings.Repeat(\"c\", 1000)))\n\t})\n}\n\nfunc TestAppendSection(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\tbase     string\n\t\tsection  string\n\t\texpected string\n\t}{\n\t\t{\n\t\t\tname:     \"both empty\",\n\t\t\tbase:     \"\",\n\t\t\tsection:  \"\",\n\t\t\texpected: \"\",\n\t\t},\n\t\t{\n\t\t\tname:     \"base empty\",\n\t\t\tbase:     \"\",\n\t\t\tsection:  \"section\",\n\t\t\texpected: \"section\",\n\t\t},\n\t\t{\n\t\t\tname:     \"section empty\",\n\t\t\tbase:     \"base\",\n\t\t\tsection:  \"\",\n\t\t\texpected: \"base\",\n\t\t},\n\t\t{\n\t\t\tname:     \"both non-empty\",\n\t\t\tbase:     \"base\",\n\t\t\tsection:  \"section\",\n\t\t\texpected: \"base\\n\\nsection\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := appendSection(tt.base, tt.section)\n\t\t\tassert.Equal(t, tt.expected, result)\n\t\t})\n\t}\n}\n\nfunc TestAllUserMessagesTagRegex(t *testing.T) {\n\tt.Run(\"matches tag\", func(t *testing.T) {\n\t\ttext := `<all_user_messages>\n    - msg1\n    - msg2\n</all_user_messages>`\n\t\tassert.True(t, allUserMessagesTagRegex.MatchString(text))\n\t})\n\n\tt.Run(\"replaces tag content\", func(t *testing.T) {\n\t\ttext := `before\n<all_user_messages>\n    - old msg\n</all_user_messages>\nafter`\n\t\treplacement := \"<all_user_messages>\\n    - new msg\\n</all_user_messages>\"\n\t\tresult := allUserMessagesTagRegex.ReplaceAllString(text, replacement)\n\t\tassert.Contains(t, result, \"new msg\")\n\t\tassert.NotContains(t, result, \"old msg\")\n\t\tassert.Contains(t, result, \"before\")\n\t\tassert.Contains(t, result, \"after\")\n\t})\n}\n\nfunc TestConfigCheck(t *testing.T) {\n\tt.Run(\"nil config\", func(t *testing.T) {\n\t\tvar c *Config\n\t\terr := c.check()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"config is required\")\n\t})\n\n\tt.Run(\"nil model\", func(t *testing.T) {\n\t\tc := &Config{}\n\t\terr := c.check()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"model is required\")\n\t})\n\n\tt.Run(\"valid config\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\n\t\tc := &Config{\n\t\t\tModel: cm,\n\t\t}\n\t\terr := c.check()\n\t\tassert.NoError(t, err)\n\t})\n\n\tt.Run(\"invalid trigger max tokens\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\n\t\tc := &Config{\n\t\t\tModel:   cm,\n\t\t\tTrigger: &TriggerCondition{ContextTokens: -1},\n\t\t}\n\t\terr := c.check()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"must be non-negative\")\n\t})\n\n\tt.Run(\"invalid trigger max messages\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\n\t\tc := &Config{\n\t\t\tModel:   cm,\n\t\t\tTrigger: &TriggerCondition{ContextMessages: -1},\n\t\t}\n\t\terr := c.check()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"must be non-negative\")\n\t})\n\n\tt.Run(\"both trigger conditions are zero\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\n\t\tc := &Config{\n\t\t\tModel:   cm,\n\t\t\tTrigger: &TriggerCondition{ContextTokens: 0, ContextMessages: 0},\n\t\t}\n\t\terr := c.check()\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"must be non-negative\")\n\t})\n}\n\nfunc TestSetGetContentType(t *testing.T) {\n\tmsg := &schema.Message{\n\t\tRole:    schema.User,\n\t\tContent: \"test\",\n\t}\n\n\tsetContentType(msg, contentTypeSummary)\n\n\tct, ok := getContentType(msg)\n\tassert.True(t, ok)\n\tassert.Equal(t, contentTypeSummary, ct)\n}\n\nfunc TestSetGetExtra(t *testing.T) {\n\tt.Run(\"set and get\", func(t *testing.T) {\n\t\tmsg := &schema.Message{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"test\",\n\t\t}\n\n\t\tsetExtra(msg, \"key\", \"value\")\n\n\t\tv, ok := getExtra[string](msg, \"key\")\n\t\tassert.True(t, ok)\n\t\tassert.Equal(t, \"value\", v)\n\t})\n\n\tt.Run(\"get from nil message\", func(t *testing.T) {\n\t\tv, ok := getExtra[string](nil, \"key\")\n\t\tassert.False(t, ok)\n\t\tassert.Equal(t, \"\", v)\n\t})\n\n\tt.Run(\"get non-existent key\", func(t *testing.T) {\n\t\tmsg := &schema.Message{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"test\",\n\t\t}\n\n\t\tv, ok := getExtra[string](msg, \"non-existent\")\n\t\tassert.False(t, ok)\n\t\tassert.Equal(t, \"\", v)\n\t})\n}\n\nfunc TestMiddlewareSummarize(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"message structure\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) {\n\t\t\t\tassert.GreaterOrEqual(t, len(msgs), 3)\n\t\t\t\tassert.Equal(t, schema.System, msgs[0].Role)\n\t\t\t\tassert.Equal(t, schema.User, msgs[len(msgs)-1].Role)\n\t\t\t\treturn &schema.Message{\n\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\tContent: \"summary\",\n\t\t\t\t}, nil\n\t\t\t}).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel: cm,\n\t\t\t},\n\t\t}\n\n\t\ttestMsg := []adk.Message{schema.UserMessage(\"test\")}\n\t\t_, err := mw.summarize(ctx, testMsg, testMsg)\n\t\tassert.NoError(t, err)\n\t})\n\n\tt.Run(\"uses context messages\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) {\n\t\t\t\t// Verify the context messages are included\n\t\t\t\tfound := false\n\t\t\t\tfor _, msg := range msgs {\n\t\t\t\t\tif msg.Content == \"context message\" {\n\t\t\t\t\t\tfound = true\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\tassert.True(t, found, \"should contain context message\")\n\n\t\t\t\treturn &schema.Message{\n\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\tContent: \"summary\",\n\t\t\t\t}, nil\n\t\t\t}).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel: cm,\n\t\t\t},\n\t\t}\n\n\t\tcontextMsgs := []adk.Message{\n\t\t\tschema.UserMessage(\"context message\"),\n\t\t}\n\t\t_, err := mw.summarize(ctx, contextMsgs, contextMsgs)\n\t\tassert.NoError(t, err)\n\t})\n\n\tt.Run(\"uses GenModelInput\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\n\t\texpectedInput := []adk.Message{\n\t\t\tschema.UserMessage(\"custom input\"),\n\t\t}\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) {\n\t\t\t\tassert.Len(t, msgs, 1)\n\t\t\t\tassert.Equal(t, \"custom input\", msgs[0].Content)\n\t\t\t\treturn &schema.Message{\n\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\tContent: \"summary\",\n\t\t\t\t}, nil\n\t\t\t}).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel: cm,\n\t\t\t\tGenModelInput: func(ctx context.Context, defaultSystemInstruction, userInstruction adk.Message, originalMsgs []adk.Message) ([]adk.Message, error) {\n\t\t\t\t\treturn expectedInput, nil\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\ttestMsg := []adk.Message{schema.UserMessage(\"test\")}\n\t\t_, err := mw.summarize(ctx, testMsg, testMsg)\n\t\tassert.NoError(t, err)\n\t})\n\n\tt.Run(\"GenModelInput error\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel: cm,\n\t\t\t\tGenModelInput: func(ctx context.Context, defaultSystemInstruction, userInstruction adk.Message, originalMsgs []adk.Message) ([]adk.Message, error) {\n\t\t\t\t\treturn nil, errors.New(\"gen input error\")\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\ttestMsg := []adk.Message{schema.UserMessage(\"test\")}\n\t\t_, err := mw.summarize(ctx, testMsg, testMsg)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"gen input error\")\n\t})\n\n\tt.Run(\"uses custom instruction\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) {\n\t\t\t\tlastMsg := msgs[len(msgs)-1]\n\t\t\t\tassert.Equal(t, schema.User, lastMsg.Role)\n\t\t\t\tassert.Contains(t, lastMsg.Content, \"custom instruction\")\n\t\t\t\treturn &schema.Message{\n\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\tContent: \"summary\",\n\t\t\t\t}, nil\n\t\t\t}).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel:           cm,\n\t\t\t\tUserInstruction: \"custom instruction\",\n\t\t\t},\n\t\t}\n\n\t\ttestMsg := []adk.Message{schema.UserMessage(\"test\")}\n\t\t_, err := mw.summarize(ctx, testMsg, testMsg)\n\t\tassert.NoError(t, err)\n\t})\n\n\tt.Run(\"model generate error\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockBaseChatModel(ctrl)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(nil, errors.New(\"generate error\")).Times(1)\n\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tModel: cm,\n\t\t\t},\n\t\t}\n\n\t\ttestMsg := []adk.Message{schema.UserMessage(\"test\")}\n\t\t_, err := mw.summarize(ctx, testMsg, testMsg)\n\t\tassert.Error(t, err)\n\t})\n}\n\nfunc TestReplaceUserMessagesInSummary(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"replaces user messages section\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{},\n\t\t}\n\n\t\tmsgs := []adk.Message{\n\t\t\tschema.UserMessage(\"msg1\"),\n\t\t\tschema.AssistantMessage(\"response1\", nil),\n\t\t\tschema.UserMessage(\"msg2\"),\n\t\t}\n\n\t\tsummary := `1. Primary Request:\n   test\n\n6. All user messages:\n<all_user_messages>\n    - [old message]\n</all_user_messages>\n\n7. Pending Tasks:\n   - task1`\n\n\t\tresult, err := mw.replaceUserMessagesInSummary(ctx, msgs, summary, 1000)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, result, \"msg1\")\n\t\tassert.Contains(t, result, \"msg2\")\n\t\tassert.NotContains(t, result, \"old message\")\n\t\tassert.Contains(t, result, \"7. Pending Tasks:\")\n\t})\n\n\tt.Run(\"returns original if no matching sections\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{},\n\t\t}\n\n\t\tmsgs := []adk.Message{\n\t\t\tschema.UserMessage(\"test\"),\n\t\t}\n\n\t\tsummary := \"summary without sections\"\n\t\tresult, err := mw.replaceUserMessagesInSummary(ctx, msgs, summary, 1000)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, summary, result)\n\t})\n\n\tt.Run(\"skips summary messages\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{},\n\t\t}\n\n\t\tsummaryMsg := &schema.Message{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"summary\",\n\t\t}\n\t\tsetContentType(summaryMsg, contentTypeSummary)\n\n\t\tmsgs := []adk.Message{\n\t\t\tsummaryMsg,\n\t\t\tschema.UserMessage(\"regular message\"),\n\t\t}\n\n\t\tsummary := `6. All user messages:\n<all_user_messages>\n    - [old]\n</all_user_messages>\n\n7. Pending Tasks:\n   - task`\n\n\t\tresult, err := mw.replaceUserMessagesInSummary(ctx, msgs, summary, 1000)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, result, \"regular message\")\n\t\tassert.NotContains(t, result, \"    - summary\")\n\t})\n\n\tt.Run(\"token counter error\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tTokenCounter: func(ctx context.Context, input *TokenCounterInput) (int, error) {\n\t\t\t\t\treturn 0, errors.New(\"count error\")\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmsgs := []adk.Message{\n\t\t\tschema.UserMessage(\"test1\"),\n\t\t\tschema.UserMessage(\"test2\"),\n\t\t}\n\n\t\t_, err := mw.replaceUserMessagesInSummary(ctx, msgs, \"summary\", 1000)\n\t\tassert.Error(t, err)\n\t})\n\n\tt.Run(\"returns original if empty user messages\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{},\n\t\t}\n\n\t\tmsgs := []adk.Message{\n\t\t\tschema.AssistantMessage(\"response\", nil),\n\t\t}\n\n\t\tsummary := `6. All user messages:\n    - [old]\n\n7. Pending Tasks:\n   - task`\n\n\t\tresult, err := mw.replaceUserMessagesInSummary(ctx, msgs, summary, 1000)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, summary, result)\n\t})\n}\n\nfunc TestAllUserMessagesTagRegexMatch(t *testing.T) {\n\tt.Run(\"matches xml tag\", func(t *testing.T) {\n\t\ttext := \"<all_user_messages>\\n    - msg\\n</all_user_messages>\"\n\t\tassert.True(t, allUserMessagesTagRegex.MatchString(text))\n\t})\n\n\tt.Run(\"does not match without tag\", func(t *testing.T) {\n\t\ttext := \"6. All user messages:\\n    - msg\"\n\t\tassert.False(t, allUserMessagesTagRegex.MatchString(text))\n\t})\n}\n\nfunc TestDefaultTrimUserMessage(t *testing.T) {\n\tt.Run(\"returns nil for zero remaining tokens\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"test\")\n\t\tresult := defaultTrimUserMessage(msg, 0)\n\t\tassert.Nil(t, result)\n\t})\n\n\tt.Run(\"returns nil for empty content\", func(t *testing.T) {\n\t\tmsg := schema.UserMessage(\"\")\n\t\tresult := defaultTrimUserMessage(msg, 100)\n\t\tassert.Nil(t, result)\n\t})\n\n\tt.Run(\"trims long message\", func(t *testing.T) {\n\t\tlongText := strings.Repeat(\"a\", 3000)\n\t\tmsg := schema.UserMessage(longText)\n\t\tresult := defaultTrimUserMessage(msg, 100)\n\t\tassert.NotNil(t, result)\n\t\tassert.Less(t, len(result.Content), len(longText))\n\t})\n}\n\nfunc TestDefaultTokenCounter(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"counts tool tokens\", func(t *testing.T) {\n\t\tinput := &TokenCounterInput{\n\t\t\tMessages: []adk.Message{},\n\t\t\tTools: []*schema.ToolInfo{\n\t\t\t\t{Name: \"test_tool\", Desc: \"description\"},\n\t\t\t},\n\t\t}\n\t\tcount, err := defaultTokenCounter(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.Greater(t, count, 0)\n\t})\n}\n\nfunc TestPostProcessSummary(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"with transcript path\", func(t *testing.T) {\n\t\tmw := &middleware{\n\t\t\tcfg: &Config{\n\t\t\t\tTranscriptFilePath: \"/path/to/transcript.txt\",\n\t\t\t},\n\t\t}\n\n\t\tsummary := &schema.Message{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"summary content\",\n\t\t}\n\n\t\tresult, err := mw.postProcessSummary(ctx, []adk.Message{}, summary)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.UserInputMultiContent, 2)\n\t\tassert.Contains(t, result.UserInputMultiContent[0].Text, \"/path/to/transcript.txt\")\n\t})\n}\n"
  },
  {
    "path": "adk/prebuilt/deep/checkpoint_compat_resume_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage deep\n\nimport (\n\t\"context\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"runtime\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype compatCheckpointStore struct {\n\tdata map[string][]byte\n}\n\nfunc newCompatCheckpointStore() *compatCheckpointStore {\n\treturn &compatCheckpointStore{data: make(map[string][]byte)}\n}\n\nfunc (s *compatCheckpointStore) Set(_ context.Context, key string, value []byte) error {\n\ts.data[key] = append([]byte(nil), value...)\n\treturn nil\n}\n\nfunc (s *compatCheckpointStore) Get(_ context.Context, key string) ([]byte, bool, error) {\n\tv, ok := s.data[key]\n\tif !ok {\n\t\treturn nil, false, nil\n\t}\n\treturn append([]byte(nil), v...), true, nil\n}\n\ntype interruptingSubAgentTool struct {\n\tname string\n}\n\nfunc (t *interruptingSubAgentTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: t.name,\n\t\tDesc: \"interrupts on first call and resumes from stored state\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"action\": {Type: schema.String},\n\t\t}),\n\t}, nil\n}\n\nfunc (t *interruptingSubAgentTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\t_, _, _ = tool.GetInterruptState[string](ctx)\n\treturn \"resumed\", nil\n}\n\nfunc readTestdataBytes(t *testing.T, filename string) []byte {\n\tt.Helper()\n\t_, file, _, ok := runtime.Caller(0)\n\tassert.True(t, ok)\n\tp := filepath.Join(filepath.Dir(file), \"testdata\", filename)\n\tb, err := os.ReadFile(p)\n\tassert.NoError(t, err)\n\tassert.NotEmpty(t, b)\n\treturn b\n}\n\nfunc runDeepAgentCheckpointCompat(t *testing.T, checkpointID string, filename string) {\n\tt.Helper()\n\tctx := context.Background()\n\n\tdata := readTestdataBytes(t, filename)\n\n\tstore := newCompatCheckpointStore()\n\tassert.NoError(t, store.Set(ctx, checkpointID, data))\n\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tinterruptToolName := \"interrupt_in_subagent_tool\"\n\tsubTool := &interruptingSubAgentTool{name: interruptToolName}\n\n\tdeepModel := mockModel.NewMockBaseChatModel(ctrl)\n\tsubModel := mockModel.NewMockBaseChatModel(ctrl)\n\n\tdeepModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\treturn schema.AssistantMessage(\"deep done\", nil), nil\n\t\t}).AnyTimes()\n\n\tsubModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\treturn schema.AssistantMessage(\"sub done\", nil), nil\n\t\t}).AnyTimes()\n\n\tsubAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"sub_chatmodel_agent\",\n\t\tDescription: \"sub agent\",\n\t\tModel:       subModel,\n\t\tToolsConfig: adk.ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{subTool},\n\t\t\t},\n\t\t},\n\t\tMaxIterations: 4,\n\t})\n\tassert.NoError(t, err)\n\n\tdeepAgent, err := New(ctx, &Config{\n\t\tName:                   \"deep\",\n\t\tDescription:            \"deep agent\",\n\t\tChatModel:              deepModel,\n\t\tSubAgents:              []adk.Agent{subAgent},\n\t\tMaxIteration:           4,\n\t\tWithoutWriteTodos:      true,\n\t\tWithoutGeneralSubAgent: true,\n\t})\n\tassert.NoError(t, err)\n\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{\n\t\tAgent:           deepAgent,\n\t\tCheckPointStore: store,\n\t})\n\n\tit, err := runner.Resume(ctx, checkpointID)\n\tassert.NoError(t, err)\n\n\tvar sawDeepDone bool\n\tvar sawAnyOutput bool\n\tfor {\n\t\tev, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, ev.Err)\n\t\tif ev.Output != nil && ev.Output.MessageOutput != nil && ev.Output.MessageOutput.Message != nil {\n\t\t\tsawAnyOutput = true\n\t\t\tmsg := ev.Output.MessageOutput.Message\n\t\t\tif msg.Role == schema.Assistant && strings.Contains(msg.Content, \"deep done\") {\n\t\t\t\tsawDeepDone = true\n\t\t\t}\n\t\t}\n\t}\n\n\tassert.True(t, sawAnyOutput)\n\tassert.True(t, sawDeepDone)\n}\n\nfunc TestDeepAgentCheckpointCompat_V0_8_Resume(t *testing.T) {\n\ttests := []struct {\n\t\tname         string\n\t\tcheckpointID string\n\t\tfilename     string\n\t}{\n\t\t{\n\t\t\tname:         \"v0.7.37\",\n\t\t\tcheckpointID: \"checkpoint_compat_v0_7_37\",\n\t\t\tfilename:     \"checkpoint_data_v0.7.37.bin\",\n\t\t},\n\t\t{\n\t\t\tname:         \"v0.8.2\",\n\t\t\tcheckpointID: \"checkpoint_compat_v0_8_2\",\n\t\t\tfilename:     \"checkpoint_data_v0.8.2.bin\",\n\t\t},\n\t\t{\n\t\t\tname:         \"v0.8.3\",\n\t\t\tcheckpointID: \"checkpoint_compat_v0_8_3\",\n\t\t\tfilename:     \"checkpoint_data_v0.8.3.bin\",\n\t\t},\n\t\t{\n\t\t\tname:         \"v0.8.4\",\n\t\t\tcheckpointID: \"checkpoint_compat_v0_8_4\",\n\t\t\tfilename:     \"checkpoint_data_v0.8.4.bin\",\n\t\t},\n\t}\n\n\tfor _, tc := range tests {\n\t\tt.Run(tc.name, func(t *testing.T) {\n\t\t\trunDeepAgentCheckpointCompat(t, tc.checkpointID, tc.filename)\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "adk/prebuilt/deep/deep.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package deep provides a prebuilt agent with deep task orchestration.\npackage deep\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/filesystem\"\n\t\"github.com/cloudwego/eino/adk/internal\"\n\tfilesystem2 \"github.com/cloudwego/eino/adk/middlewares/filesystem\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool/utils\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc init() {\n\tschema.RegisterName[TODO](\"_eino_adk_prebuilt_deep_todo\")\n\tschema.RegisterName[[]TODO](\"_eino_adk_prebuilt_deep_todo_slice\")\n}\n\n// Config defines the configuration for creating a DeepAgent.\ntype Config struct {\n\t// Name is the identifier for the Deep agent.\n\tName string\n\t// Description provides a brief explanation of the agent's purpose.\n\tDescription string\n\n\t// ChatModel is the model used by DeepAgent for reasoning and task execution.\n\t// If the agent uses any tools, this model must support the model.WithTools call option,\n\t// as that's how the agent configures the model with tool information.\n\tChatModel model.BaseChatModel\n\t// Instruction contains the system prompt that guides the agent's behavior.\n\t// When empty, a built-in default system prompt will be used, which includes general assistant\n\t// behavior guidelines, security policies, coding style guidelines, and tool usage policies.\n\tInstruction string\n\t// SubAgents are specialized agents that can be invoked by the agent.\n\tSubAgents []adk.Agent\n\t// ToolsConfig provides the tools and tool-calling configurations available for the agent to invoke.\n\tToolsConfig adk.ToolsConfig\n\t// MaxIteration limits the maximum number of reasoning iterations the agent can perform.\n\tMaxIteration int\n\n\t// Backend provides filesystem operations used by tools and offloading.\n\t// If set, filesystem tools (read_file, write_file, edit_file, glob, grep) will be registered.\n\t// Optional.\n\tBackend filesystem.Backend\n\t// Shell provides shell command execution capability.\n\t// If set, an execute tool will be registered to support shell command execution.\n\t// Optional. Mutually exclusive with StreamingShell.\n\tShell filesystem.Shell\n\t// StreamingShell provides streaming shell command execution capability.\n\t// If set, a streaming execute tool will be registered to support streaming shell command execution.\n\t// Optional. Mutually exclusive with Shell.\n\tStreamingShell filesystem.StreamingShell\n\n\t// WithoutWriteTodos disables the built-in write_todos tool when set to true.\n\tWithoutWriteTodos bool\n\t// WithoutGeneralSubAgent disables the general-purpose subagent when set to true.\n\tWithoutGeneralSubAgent bool\n\t// TaskToolDescriptionGenerator allows customizing the description for the task tool.\n\t// If provided, this function generates the tool description based on available subagents.\n\tTaskToolDescriptionGenerator func(ctx context.Context, availableAgents []adk.Agent) (string, error)\n\n\tMiddlewares []adk.AgentMiddleware\n\n\t// Handlers configures interface-based handlers for extending agent behavior.\n\t// Unlike Middlewares (struct-based), Handlers allow users to:\n\t//   - Add custom methods to their handler implementations\n\t//   - Return modified context from handler methods\n\t//   - Centralize configuration in struct fields instead of closures\n\t//\n\t// Handlers are processed after Middlewares, in registration order.\n\t// See adk.ChatModelAgentMiddleware documentation for when to use Handlers vs Middlewares.\n\tHandlers []adk.ChatModelAgentMiddleware\n\n\tModelRetryConfig *adk.ModelRetryConfig\n\n\t// OutputKey stores the agent's response in the session.\n\t// Optional. When set, stores output via AddSessionValue(ctx, outputKey, msg.Content).\n\tOutputKey string\n}\n\n// New creates a new Deep agent instance with the provided configuration.\n// This function initializes built-in tools, creates a task tool for subagent orchestration,\n// and returns a fully configured ChatModelAgent ready for execution.\nfunc New(ctx context.Context, cfg *Config) (adk.ResumableAgent, error) {\n\thandlers, err := buildBuiltinAgentMiddlewares(ctx, cfg)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tinstruction := cfg.Instruction\n\tif len(instruction) == 0 {\n\t\tinstruction = internal.SelectPrompt(internal.I18nPrompts{\n\t\t\tEnglish: baseAgentInstruction,\n\t\t\tChinese: baseAgentInstructionChinese,\n\t\t})\n\t}\n\n\tif !cfg.WithoutGeneralSubAgent || len(cfg.SubAgents) > 0 {\n\t\ttt, err := newTaskToolMiddleware(\n\t\t\tctx,\n\t\t\tcfg.TaskToolDescriptionGenerator,\n\t\t\tcfg.SubAgents,\n\n\t\t\tcfg.WithoutGeneralSubAgent,\n\t\t\tcfg.ChatModel,\n\t\t\tinstruction,\n\t\t\tcfg.ToolsConfig,\n\t\t\tcfg.MaxIteration,\n\t\t\tcfg.Middlewares,\n\t\t\tappend(handlers, cfg.Handlers...),\n\t\t)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to new task tool: %w\", err)\n\t\t}\n\t\thandlers = append(handlers, tt)\n\t}\n\n\treturn adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:          cfg.Name,\n\t\tDescription:   cfg.Description,\n\t\tInstruction:   instruction,\n\t\tModel:         cfg.ChatModel,\n\t\tToolsConfig:   cfg.ToolsConfig,\n\t\tMaxIterations: cfg.MaxIteration,\n\t\tMiddlewares:   cfg.Middlewares,\n\t\tHandlers:      append(handlers, cfg.Handlers...),\n\n\t\tGenModelInput:    genModelInput,\n\t\tModelRetryConfig: cfg.ModelRetryConfig,\n\t\tOutputKey:        cfg.OutputKey,\n\t})\n}\n\nfunc genModelInput(ctx context.Context, instruction string, input *adk.AgentInput) ([]*schema.Message, error) {\n\tmsgs := make([]*schema.Message, 0, len(input.Messages)+1)\n\n\tif instruction != \"\" {\n\t\tmsgs = append(msgs, schema.SystemMessage(instruction))\n\t}\n\n\tmsgs = append(msgs, input.Messages...)\n\n\treturn msgs, nil\n}\n\nfunc buildBuiltinAgentMiddlewares(ctx context.Context, cfg *Config) ([]adk.ChatModelAgentMiddleware, error) {\n\tvar ms []adk.ChatModelAgentMiddleware\n\tif !cfg.WithoutWriteTodos {\n\t\tt, err := newWriteTodos()\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tms = append(ms, t)\n\t}\n\n\tif cfg.Backend != nil || cfg.Shell != nil || cfg.StreamingShell != nil {\n\t\tfm, err := filesystem2.New(ctx, &filesystem2.MiddlewareConfig{\n\t\t\tBackend:        cfg.Backend,\n\t\t\tShell:          cfg.Shell,\n\t\t\tStreamingShell: cfg.StreamingShell,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tms = append(ms, fm)\n\t}\n\n\treturn ms, nil\n}\n\ntype TODO struct {\n\tContent    string `json:\"content\"`\n\tActiveForm string `json:\"activeForm\"`\n\tStatus     string `json:\"status\" jsonschema:\"enum=pending,enum=in_progress,enum=completed\"`\n}\n\ntype writeTodosArguments struct {\n\tTodos []TODO `json:\"todos\"`\n}\n\nfunc newWriteTodos() (adk.ChatModelAgentMiddleware, error) {\n\ttoolDesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: writeTodosToolDescription,\n\t\tChinese: writeTodosToolDescriptionChinese,\n\t})\n\tresultMsg := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: \"Updated todo list to %s\",\n\t\tChinese: \"已更新待办列表为 %s\",\n\t})\n\n\tt, err := utils.InferTool(\"write_todos\", toolDesc, func(ctx context.Context, input writeTodosArguments) (output string, err error) {\n\t\tadk.AddSessionValue(ctx, SessionKeyTodos, input.Todos)\n\t\ttodos, err := sonic.MarshalString(input.Todos)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn fmt.Sprintf(resultMsg, todos), nil\n\t})\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn buildAppendPromptTool(\"\", t), nil\n}\n"
  },
  {
    "path": "adk/prebuilt/deep/deep_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage deep\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/prebuilt/planexecute\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestGenModelInput(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"WithInstruction\", func(t *testing.T) {\n\t\tinput := &adk.AgentInput{\n\t\t\tMessages: []*schema.Message{\n\t\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t},\n\t\t}\n\n\t\tmsgs, err := genModelInput(ctx, \"You are a helpful assistant\", input)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, msgs, 2)\n\t\tassert.Equal(t, schema.System, msgs[0].Role)\n\t\tassert.Equal(t, \"You are a helpful assistant\", msgs[0].Content)\n\t\tassert.Equal(t, schema.User, msgs[1].Role)\n\t\tassert.Equal(t, \"hello\", msgs[1].Content)\n\t})\n\n\tt.Run(\"WithoutInstruction\", func(t *testing.T) {\n\t\tinput := &adk.AgentInput{\n\t\t\tMessages: []*schema.Message{\n\t\t\t\tschema.UserMessage(\"hello\"),\n\t\t\t},\n\t\t}\n\n\t\tmsgs, err := genModelInput(ctx, \"\", input)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, msgs, 1)\n\t\tassert.Equal(t, schema.User, msgs[0].Role)\n\t\tassert.Equal(t, \"hello\", msgs[0].Content)\n\t})\n}\n\nfunc TestWriteTodos(t *testing.T) {\n\tm, err := buildBuiltinAgentMiddlewares(context.Background(), &Config{WithoutWriteTodos: false})\n\tassert.NoError(t, err)\n\n\twt := m[0].(*appendPromptTool).t.(tool.InvokableTool)\n\n\ttodos := `[{\"content\":\"content1\",\"activeForm\":\"\",\"status\":\"pending\"},{\"content\":\"content2\",\"activeForm\":\"\",\"status\":\"pending\"}]`\n\targs := fmt.Sprintf(`{\"todos\": %s}`, todos)\n\n\tresult, err := wt.InvokableRun(context.Background(), args)\n\tassert.NoError(t, err)\n\tassert.Equal(t, fmt.Sprintf(\"Updated todo list to %s\", todos), result)\n}\n\nfunc TestDeepSubAgentSharesSessionValues(t *testing.T) {\n\tctx := context.Background()\n\tspy := &spySubAgent{}\n\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tcalls := 0\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\tcalls++\n\t\t\tif calls == 1 {\n\t\t\t\tc := schema.ToolCall{ID: \"id-1\", Type: \"function\"}\n\t\t\t\tc.Function.Name = taskToolName\n\t\t\t\tc.Function.Arguments = fmt.Sprintf(`{\"subagent_type\":\"%s\",\"description\":\"from_parent\"}`, spy.Name(ctx))\n\t\t\t\treturn schema.AssistantMessage(\"\", []schema.ToolCall{c}), nil\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"done\", nil), nil\n\t\t}).AnyTimes()\n\n\tagent, err := New(ctx, &Config{\n\t\tName:                   \"deep\",\n\t\tDescription:            \"deep agent\",\n\t\tChatModel:              cm,\n\t\tInstruction:            \"you are deep agent\",\n\t\tSubAgents:              []adk.Agent{spy},\n\t\tToolsConfig:            adk.ToolsConfig{},\n\t\tMaxIteration:           2,\n\t\tWithoutWriteTodos:      true,\n\t\tWithoutGeneralSubAgent: true,\n\t})\n\tassert.NoError(t, err)\n\n\tr := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent})\n\tit := r.Run(ctx, []adk.Message{schema.UserMessage(\"hi\")}, adk.WithSessionValues(map[string]any{\"parent_key\": \"parent_val\"}))\n\tfor {\n\t\tif _, ok := it.Next(); !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassert.Equal(t, \"parent_val\", spy.seenParentValue)\n}\n\nfunc TestDeepSubAgentFollowsStreamingMode(t *testing.T) {\n\tctx := context.Background()\n\tspy := &spyStreamingSubAgent{}\n\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tsubName := spy.Name(ctx)\n\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID:   \"id-1\",\n\t\t\t\t\tType: \"function\",\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      taskToolName,\n\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"subagent_type\":\"%s\",\"description\":\"from_parent\"}`, subName),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}),\n\t\t}), nil).\n\t\tTimes(1)\n\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\tschema.AssistantMessage(\"done\", nil),\n\t\t}), nil).\n\t\tTimes(1)\n\n\tagent, err := New(ctx, &Config{\n\t\tName:                   \"deep\",\n\t\tDescription:            \"deep agent\",\n\t\tChatModel:              cm,\n\t\tInstruction:            \"you are deep agent\",\n\t\tSubAgents:              []adk.Agent{spy},\n\t\tToolsConfig:            adk.ToolsConfig{},\n\t\tMaxIteration:           2,\n\t\tWithoutWriteTodos:      true,\n\t\tWithoutGeneralSubAgent: true,\n\t})\n\tassert.NoError(t, err)\n\n\tr := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent, EnableStreaming: true})\n\tit := r.Run(ctx, []adk.Message{schema.UserMessage(\"hi\")})\n\tfor {\n\t\tif _, ok := it.Next(); !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassert.True(t, spy.seenEnableStreaming)\n}\n\ntype spySubAgent struct {\n\tseenParentValue any\n}\n\nfunc (s *spySubAgent) Name(context.Context) string        { return \"spy-subagent\" }\nfunc (s *spySubAgent) Description(context.Context) string { return \"spy\" }\nfunc (s *spySubAgent) Run(ctx context.Context, _ *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\ts.seenParentValue, _ = adk.GetSessionValue(ctx, \"parent_key\")\n\tit, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tgen.Send(adk.EventFromMessage(schema.AssistantMessage(\"ok\", nil), nil, schema.Assistant, \"\"))\n\tgen.Close()\n\treturn it\n}\n\ntype spyStreamingSubAgent struct {\n\tseenEnableStreaming bool\n}\n\nfunc (s *spyStreamingSubAgent) Name(context.Context) string        { return \"spy-streaming-subagent\" }\nfunc (s *spyStreamingSubAgent) Description(context.Context) string { return \"spy\" }\nfunc (s *spyStreamingSubAgent) Run(ctx context.Context, input *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\tif input != nil {\n\t\ts.seenEnableStreaming = input.EnableStreaming\n\t}\n\tit, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tgen.Send(adk.EventFromMessage(schema.AssistantMessage(\"ok\", nil), nil, schema.Assistant, \"\"))\n\tgen.Close()\n\treturn it\n}\n\nfunc TestDeepAgentWithPlanExecuteSubAgent_InternalEventsEmitted(t *testing.T) {\n\tctx := context.Background()\n\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tdeepModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tplannerModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\texecutorModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\treplannerModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tdeepModel.EXPECT().WithTools(gomock.Any()).Return(deepModel, nil).AnyTimes()\n\tplannerModel.EXPECT().WithTools(gomock.Any()).Return(plannerModel, nil).AnyTimes()\n\texecutorModel.EXPECT().WithTools(gomock.Any()).Return(executorModel, nil).AnyTimes()\n\treplannerModel.EXPECT().WithTools(gomock.Any()).Return(replannerModel, nil).AnyTimes()\n\n\tplannerModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input []*schema.Message, opts ...interface{}) (*schema.StreamReader[*schema.Message], error) {\n\t\t\tsr, sw := schema.Pipe[*schema.Message](1)\n\t\t\tgo func() {\n\t\t\t\tdefer sw.Close()\n\t\t\t\tplanJSON := `{\"steps\":[\"step1\"]}`\n\t\t\t\tmsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:   \"plan_call_1\",\n\t\t\t\t\t\tType: \"function\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      \"plan\",\n\t\t\t\t\t\t\tArguments: planJSON,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t\tsw.Send(msg, nil)\n\t\t\t}()\n\t\t\treturn sr, nil\n\t\t},\n\t).Times(1)\n\n\texecutorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\treturn schema.AssistantMessage(\"executed step1\", nil), nil\n\t\t},\n\t).Times(1)\n\n\treplannerModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input []*schema.Message, opts ...interface{}) (*schema.StreamReader[*schema.Message], error) {\n\t\t\tsr, sw := schema.Pipe[*schema.Message](1)\n\t\t\tgo func() {\n\t\t\t\tdefer sw.Close()\n\t\t\t\tresponseJSON := `{\"response\":\"final response\"}`\n\t\t\t\tmsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID:   \"respond_call_1\",\n\t\t\t\t\t\tType: \"function\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      \"respond\",\n\t\t\t\t\t\t\tArguments: responseJSON,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t\tsw.Send(msg, nil)\n\t\t\t}()\n\t\t\treturn sr, nil\n\t\t},\n\t).Times(1)\n\n\tplanner, err := planexecute.NewPlanner(ctx, &planexecute.PlannerConfig{\n\t\tToolCallingChatModel: plannerModel,\n\t})\n\tassert.NoError(t, err)\n\n\texecutor, err := planexecute.NewExecutor(ctx, &planexecute.ExecutorConfig{\n\t\tModel: executorModel,\n\t})\n\tassert.NoError(t, err)\n\n\treplanner, err := planexecute.NewReplanner(ctx, &planexecute.ReplannerConfig{\n\t\tChatModel: replannerModel,\n\t})\n\tassert.NoError(t, err)\n\n\tplanExecuteAgent, err := planexecute.New(ctx, &planexecute.Config{\n\t\tPlanner:   planner,\n\t\tExecutor:  executor,\n\t\tReplanner: replanner,\n\t})\n\tassert.NoError(t, err)\n\n\tnamedPlanExecuteAgent := &namedPlanExecuteAgent{\n\t\tResumableAgent: planExecuteAgent,\n\t\tname:           \"plan_execute_subagent\",\n\t\tdescription:    \"a plan execute subagent\",\n\t}\n\n\tdeepModelCalls := 0\n\tdeepModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\tdeepModelCalls++\n\t\t\tif deepModelCalls == 1 {\n\t\t\t\tc := schema.ToolCall{ID: \"id-1\", Type: \"function\"}\n\t\t\t\tc.Function.Name = taskToolName\n\t\t\t\tc.Function.Arguments = fmt.Sprintf(`{\"subagent_type\":\"%s\",\"description\":\"execute the plan\"}`, namedPlanExecuteAgent.name)\n\t\t\t\treturn schema.AssistantMessage(\"\", []schema.ToolCall{c}), nil\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"done\", nil), nil\n\t\t}).AnyTimes()\n\n\tdeepAgent, err := New(ctx, &Config{\n\t\tName:                   \"deep\",\n\t\tDescription:            \"deep agent\",\n\t\tChatModel:              deepModel,\n\t\tInstruction:            \"you are deep agent\",\n\t\tSubAgents:              []adk.Agent{namedPlanExecuteAgent},\n\t\tToolsConfig:            adk.ToolsConfig{EmitInternalEvents: true},\n\t\tMaxIteration:           5,\n\t\tWithoutWriteTodos:      true,\n\t\tWithoutGeneralSubAgent: true,\n\t})\n\tassert.NoError(t, err)\n\n\tr := adk.NewRunner(ctx, adk.RunnerConfig{Agent: deepAgent})\n\tit := r.Run(ctx, []adk.Message{schema.UserMessage(\"hi\")})\n\n\tvar events []*adk.AgentEvent\n\tfor {\n\t\tevent, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\tassert.Greater(t, len(events), 0, \"should have at least one event\")\n\n\tvar deepAgentEvents []*adk.AgentEvent\n\tvar plannerEvents []*adk.AgentEvent\n\tvar executorEvents []*adk.AgentEvent\n\tvar replannerEvents []*adk.AgentEvent\n\tvar planExecuteEvents []*adk.AgentEvent\n\n\tfor _, event := range events {\n\t\tswitch event.AgentName {\n\t\tcase \"deep\":\n\t\t\tdeepAgentEvents = append(deepAgentEvents, event)\n\t\tcase \"planner\":\n\t\t\tplannerEvents = append(plannerEvents, event)\n\t\tcase \"executor\":\n\t\t\texecutorEvents = append(executorEvents, event)\n\t\tcase \"replanner\":\n\t\t\treplannerEvents = append(replannerEvents, event)\n\t\tcase \"plan_execute_replan\", \"execute_replan\":\n\t\t\tplanExecuteEvents = append(planExecuteEvents, event)\n\t\t}\n\t}\n\n\tassert.Greater(t, len(deepAgentEvents), 0, \"should have events from deep agent\")\n\n\tassert.Greater(t, len(plannerEvents), 0, \"planner internal events should be emitted when EmitInternalEvents is true\")\n\tassert.Greater(t, len(executorEvents), 0, \"executor internal events should be emitted when EmitInternalEvents is true\")\n\tassert.Greater(t, len(replannerEvents), 0, \"replanner internal events should be emitted when EmitInternalEvents is true\")\n\n\tt.Logf(\"Total events: %d\", len(events))\n\tt.Logf(\"Deep agent events: %d\", len(deepAgentEvents))\n\tt.Logf(\"Planner events: %d\", len(plannerEvents))\n\tt.Logf(\"Executor events: %d\", len(executorEvents))\n\tt.Logf(\"Replanner events: %d\", len(replannerEvents))\n\tt.Logf(\"PlanExecute events: %d\", len(planExecuteEvents))\n}\n\ntype namedPlanExecuteAgent struct {\n\tadk.ResumableAgent\n\tname        string\n\tdescription string\n}\n\nfunc (n *namedPlanExecuteAgent) Name(_ context.Context) string {\n\treturn n.name\n}\n\nfunc (n *namedPlanExecuteAgent) Description(_ context.Context) string {\n\treturn n.description\n}\n\nfunc TestDeepAgentOutputKey(t *testing.T) {\n\tt.Run(\"OutputKeyStoresInSession\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tctrl := gomock.NewController(t)\n\t\tdefer ctrl.Finish()\n\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Hello from DeepAgent\", nil), nil).\n\t\t\tTimes(1)\n\n\t\tagent, err := New(ctx, &Config{\n\t\t\tName:                   \"deep\",\n\t\t\tDescription:            \"deep agent\",\n\t\t\tChatModel:              cm,\n\t\t\tInstruction:            \"you are deep agent\",\n\t\t\tMaxIteration:           2,\n\t\t\tWithoutWriteTodos:      true,\n\t\t\tWithoutGeneralSubAgent: true,\n\t\t\tOutputKey:              \"deep_output\",\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tvar capturedSessionValues map[string]any\n\t\twrappedAgent := &sessionCaptureAgent{\n\t\t\tAgent:          agent,\n\t\t\tcaptureSession: func(values map[string]any) { capturedSessionValues = values },\n\t\t}\n\n\t\tr := adk.NewRunner(ctx, adk.RunnerConfig{Agent: wrappedAgent})\n\t\tit := r.Run(ctx, []adk.Message{schema.UserMessage(\"hi\")})\n\t\tfor {\n\t\t\tif _, ok := it.Next(); !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Contains(t, capturedSessionValues, \"deep_output\")\n\t\tassert.Equal(t, \"Hello from DeepAgent\", capturedSessionValues[\"deep_output\"])\n\t})\n\n\tt.Run(\"OutputKeyWithStreamingStoresInSession\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tctrl := gomock.NewController(t)\n\t\tdefer ctrl.Finish()\n\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"Hello\", nil),\n\t\t\t\tschema.AssistantMessage(\" from\", nil),\n\t\t\t\tschema.AssistantMessage(\" DeepAgent\", nil),\n\t\t\t}), nil).\n\t\t\tTimes(1)\n\n\t\tagent, err := New(ctx, &Config{\n\t\t\tName:                   \"deep\",\n\t\t\tDescription:            \"deep agent\",\n\t\t\tChatModel:              cm,\n\t\t\tInstruction:            \"you are deep agent\",\n\t\t\tMaxIteration:           2,\n\t\t\tWithoutWriteTodos:      true,\n\t\t\tWithoutGeneralSubAgent: true,\n\t\t\tOutputKey:              \"deep_output\",\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tvar capturedSessionValues map[string]any\n\t\twrappedAgent := &sessionCaptureAgent{\n\t\t\tAgent:          agent,\n\t\t\tcaptureSession: func(values map[string]any) { capturedSessionValues = values },\n\t\t}\n\n\t\tr := adk.NewRunner(ctx, adk.RunnerConfig{Agent: wrappedAgent, EnableStreaming: true})\n\t\tit := r.Run(ctx, []adk.Message{schema.UserMessage(\"hi\")})\n\t\tfor {\n\t\t\tif _, ok := it.Next(); !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.Contains(t, capturedSessionValues, \"deep_output\")\n\t\tassert.Equal(t, \"Hello from DeepAgent\", capturedSessionValues[\"deep_output\"])\n\t})\n\n\tt.Run(\"OutputKeyNotSetWhenEmpty\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tctrl := gomock.NewController(t)\n\t\tdefer ctrl.Finish()\n\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"Hello from DeepAgent\", nil), nil).\n\t\t\tTimes(1)\n\n\t\tagent, err := New(ctx, &Config{\n\t\t\tName:                   \"deep\",\n\t\t\tDescription:            \"deep agent\",\n\t\t\tChatModel:              cm,\n\t\t\tInstruction:            \"you are deep agent\",\n\t\t\tMaxIteration:           2,\n\t\t\tWithoutWriteTodos:      true,\n\t\t\tWithoutGeneralSubAgent: true,\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tvar capturedSessionValues map[string]any\n\t\twrappedAgent := &sessionCaptureAgent{\n\t\t\tAgent:          agent,\n\t\t\tcaptureSession: func(values map[string]any) { capturedSessionValues = values },\n\t\t}\n\n\t\tr := adk.NewRunner(ctx, adk.RunnerConfig{Agent: wrappedAgent})\n\t\tit := r.Run(ctx, []adk.Message{schema.UserMessage(\"hi\")})\n\t\tfor {\n\t\t\tif _, ok := it.Next(); !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tassert.NotContains(t, capturedSessionValues, \"deep_output\")\n\t})\n}\n\ntype sessionCaptureAgent struct {\n\tadk.Agent\n\tcaptureSession func(map[string]any)\n}\n\nfunc (s *sessionCaptureAgent) Run(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\tinnerIt := s.Agent.Run(ctx, input, opts...)\n\tit, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tgo func() {\n\t\tdefer gen.Close()\n\t\tfor {\n\t\t\tevent, ok := innerIt.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tgen.Send(event)\n\t\t}\n\t\ts.captureSession(adk.GetSessionValues(ctx))\n\t}()\n\treturn it\n}\n"
  },
  {
    "path": "adk/prebuilt/deep/prompt.go",
    "content": "/*\n * Copyright (c) 2025 Harrison Chase\n * Copyright (c) 2025 CloudWeGo Authors\n * SPDX-License-Identifier: MIT\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage deep\n\n// This file contains prompt templates and tool descriptions adapted from the DeepAgents project and ClaudeCode.\n// Original source: https://github.com/langchain-ai/deepagents  https://claude.com/product/claude-code\n//\n// These prompts are used under the terms of the original project's open source license.\n// When using this code in your own open source project, ensure compliance with the original license requirements.\n\nconst (\n\ttaskPrompt = `\n# 'task' (subagent spawner)\n\nYou have access to a 'task' tool to launch short-lived subagents that handle isolated tasks. These agents are ephemeral — they live only for the duration of the task and return a single result.\nYou should proactively use the 'task' tool with specialized agents when the task at hand matches the agent's description.\n\nWhen to use the task tool:\n- When a task is complex and multi-step, and can be fully delegated in isolation\n- When a task is independent of other tasks and can run in parallel\n- When a task requires focused reasoning or heavy token/context usage that would bloat the orchestrator thread\n- When sandboxing improves reliability (e.g. code execution, structured searches, data formatting)\n- When you only care about the output of the subagent, and not the intermediate steps (ex. performing a lot of research and then returned a synthesized report, performing a series of computations or lookups to achieve a concise, relevant answer.)\n\nSubagent lifecycle:\n1. **Spawn** → Provide clear role, instructions, and expected output\n2. **Run** → The subagent completes the task autonomously\n3. **Return** → The subagent provides a single structured result\n4. **Reconcile** → Incorporate or synthesize the result into the main thread\n\nWhen NOT to use the task tool:\n- If you need to see the intermediate reasoning or steps after the subagent has completed (the task tool hides them)\n- If the task is trivial (a few tool calls or simple lookup)\n- If delegating does not reduce token usage, complexity, or context switching\n- If splitting would add latency without benefit\n\n## Important Task Tool Usage Notes to Remember\n- Whenever possible, parallelize the work that you do. This is true for both tool_calls, and for tasks. Whenever you have independent steps to complete - make tool_calls, or kick off tasks (subagents) in parallel to accomplish them faster. This saves time for the user, which is incredibly important.\n- Remember to use the 'task' tool to silo independent tasks within a multi-part objective.\n- You should use the 'task' tool whenever you have a complex task that will take multiple steps, and is independent from other tasks that the agent needs to complete. These agents are highly competent and efficient.\n`\n\n\tbaseAgentInstruction = `\nYou are a helpful assistant. Use the instructions below and the tools available to you to assist the user.\n\nIMPORTANT: Assist with authorized security testing, defensive security, CTF challenges, and educational contexts. Refuse requests for destructive techniques, DoS attacks, mass targeting, supply chain compromise, or detection evasion for malicious purposes. Dual-use security tools (C2 frameworks, credential testing, exploit development) require clear authorization context: pentesting engagements, CTF competitions, security research, or defensive use cases.\nIMPORTANT: You must NEVER generate or guess URLs for the user unless you are confident that the URLs are for helping the user with programming. You may use URLs provided by the user in their messages or local files.\n\n# Tone and style\n- Only use emojis if the user explicitly requests it. Avoid using emojis in all communication unless asked.\n- Your output will be displayed on a command line interface. Your responses should be short and concise. You can use Github-flavored markdown for formatting, and will be rendered in a monospace font using the CommonMark specification.\n- Output text to communicate with the user; all text you output outside of tool use is displayed to the user. Only use tools to complete tasks. Never use tools like Bash or code comments as means to communicate with the user during the session.\n- NEVER create files unless they're absolutely necessary for achieving your goal. ALWAYS prefer editing an existing file to creating a new one. This includes markdown files.\n- Do not use a colon before tool calls. Your tool calls may not be shown directly in the output, so text like \"Let me read the file:\" followed by a read tool call should just be \"Let me read the file.\" with a period.\n\n# Professional objectivity\nPrioritize technical accuracy and truthfulness over validating the user's beliefs. Focus on facts and problem-solving, providing direct, objective technical info without any unnecessary superlatives, praise, or emotional validation. It is best for the user if agent honestly applies the same rigorous standards to all ideas and disagrees when necessary, even if it may not be what the user wants to hear. Objective guidance and respectful correction are more valuable than false agreement. Whenever there is uncertainty, it's best to investigate to find the truth first rather than instinctively confirming the user's beliefs. Avoid using over-the-top validation or excessive praise when responding to users such as \"You're absolutely right\" or similar phrases.\n\n# Planning without timelines\nWhen planning tasks, provide concrete implementation steps without time estimates. Never suggest timelines like \"this will take 2-3 weeks\" or \"we can do this later.\" Focus on what needs to be done, not when. Break work into actionable steps and let users decide scheduling.\n\n# Doing tasks\nThe user will primarily request you perform software engineering tasks. This includes solving bugs, adding new functionality, refactoring code, explaining code, and more. For these tasks the following steps are recommended:\n- NEVER propose changes to code you haven't read. If a user asks about or wants you to modify a file, read it first. Understand existing code before suggesting modifications.\n- Be careful not to introduce security vulnerabilities such as command injection, XSS, SQL injection, and other OWASP top 10 vulnerabilities. If you notice that you wrote insecure code, immediately fix it.\n- Avoid over-engineering. Only make changes that are directly requested or clearly necessary. Keep solutions simple and focused.\n  - Don't add features, refactor code, or make \"improvements\" beyond what was asked. A bug fix doesn't need surrounding code cleaned up. A simple feature doesn't need extra configurability. Don't add docstrings, comments, or type annotations to code you didn't change. Only add comments where the logic isn't self-evident.\n  - Don't add error handling, fallbacks, or validation for scenarios that can't happen. Trust internal code and framework guarantees. Only validate at system boundaries (user input, external APIs). Don't use feature flags or backwards-compatibility shims when you can just change the code.\n  - Don't create helpers, utilities, or abstractions for one-time operations. Don't design for hypothetical future requirements. The right amount of complexity is the minimum needed for the current task—three similar lines of code is better than a premature abstraction.\n- Avoid backwards-compatibility hacks like renaming unused '_vars', re-exporting types, adding '// removed' comments for removed code, etc. If something is unused, delete it completely.\n\n- Tool results and user messages may include <system-reminder> tags. <system-reminder> tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear.\n- The conversation has unlimited context through automatic summarization.\n\nIMPORTANT: Complete tasks fully. Do not stop mid-task or leave work incomplete. Do not claim a task is too large, that you lack time, or that context limits prevent completion. You have unlimited context through summarization. Continue working until the task is done or the user stops you.\n\n# Tool usage policy\n- You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls.\n- If the user specifies that they want you to run tools \"in parallel\", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls.\n- Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead.\n\n\nIMPORTANT: Assist with authorized security testing, defensive security, CTF challenges, and educational contexts. Refuse requests for destructive techniques, DoS attacks, mass targeting, supply chain compromise, or detection evasion for malicious purposes. Dual-use security tools (C2 frameworks, credential testing, exploit development) require clear authorization context: pentesting engagements, CTF competitions, security research, or defensive use cases.\n\n\n# Code References\n\nWhen referencing specific functions or pieces of code include the pattern 'file_path:line_number' to allow the user to easily navigate to the source code location.\n\n<example>\nuser: Where are errors from the client handled?\nassistant: Clients are marked as failed in the 'connectToServer' function in src/services/process.ts:712.\n</example>\n`\n\tgeneralAgentDescription = `general-purpose agent for researching complex questions, searching for code, and executing multi-step tasks. When you are searching for a keyword or file and are not confident that you will find the right match in the first few tries use this agent to perform the search for you. (Tools: *)`\n\ttaskToolDescription     = `Launch a new agent to handle complex, multi-step tasks autonomously. \n\nThe Task tool launches specialized agents (subprocesses) that autonomously handle complex tasks. Each agent type has specific capabilities and tools available to it.\n\nAvailable agent types and the tools they have access to:\n{other_agents}\n\nWhen using the Task tool, you must specify a subagent_type parameter to select which agent type to use.\n\nWhen NOT to use the Task tool:\n- If you want to read a specific file path, use the Read or Glob tool instead of the Task tool, to find the match more quickly\n- If you are searching for a specific class definition like \"class Foo\", use the Glob tool instead, to find the match more quickly\n- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of the Task tool, to find the match more quickly\n- Other tasks that are not related to the agent descriptions above\n\n\nUsage notes:\n- Launch multiple agents concurrently whenever possible, to maximize performance; to do that, use a single message with multiple tool uses\n- When the agent is done, it will return a single message back to you. The result returned by the agent is not visible to the user. To show the user the result, you should send a text message back to the user with a concise summary of the result.\n- Provide clear, detailed prompts so the agent can work autonomously and return exactly the information you need.\n- The agent's outputs should generally be trusted\n- Clearly tell the agent whether you expect it to write code or just to do research (search, file reads, web fetches, etc.), since it is not aware of the user's intent\n- If the agent description mentions that it should be used proactively, then you should try your best to use it without the user having to ask for it first. Use your judgement.\n- If the user specifies that they want you to run agents \"in parallel\", you MUST send a single message with multiple Task tool use content blocks. For example, if you need to launch both a code-reviewer agent and a test-runner agent in parallel, send a single message with both tool calls.\n\nExample usage:\n\n<example_agent_descriptions>\n\"code-reviewer\": use this agent after you are done writing a significant piece of code\n\"greeting-responder\": use this agent when to respond to user greetings with a friendly joke\n</example_agent_description>\n\n<example>\nuser: \"Please write a function that checks if a number is prime\"\nassistant: Sure let me write a function that checks if a number is prime\nassistant: First let me use the Write tool to write a function that checks if a number is prime\nassistant: I'm going to use the Write tool to write the following code:\n<code>\nfunction isPrime(n) {{\n  if (n <= 1) return false\n  for (let i = 2; i * i <= n; i++) {{\n    if (n %!i(MISSING) === 0) return false\n  }}\n  return true\n}}\n</code>\n<commentary>\nSince a significant piece of code was written and the task was completed, now use the code-reviewer agent to review the code\n</commentary>\nassistant: Now let me use the code-reviewer agent to review the code\nassistant: Uses the Task tool to launch the code-reviewer agent \n</example>\n\n<example>\nuser: \"Hello\"\n<commentary>\nSince the user is greeting, use the greeting-responder agent to respond with a friendly joke\n</commentary>\nassistant: \"I'm going to use the Task tool to launch the greeting-responder agent\"\n</example>\n`\n\twriteTodosToolDescription = `Use this tool to create and manage a structured task list for your current coding session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user.\nIt also helps the user understand the progress of the task and overall progress of their requests.\n\n## When to Use This Tool\nUse this tool proactively in these scenarios:\n\n1. Complex multi-step tasks - When a task requires 3 or more distinct steps or actions\n2. Non-trivial and complex tasks - Tasks that require careful planning or multiple operations\n3. User explicitly requests todo list - When the user directly asks you to use the todo list\n4. User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated)\n5. After receiving new instructions - Immediately capture user requirements as todos\n6. When you start working on a task - Mark it as in_progress BEFORE beginning work. Ideally you should only have one todo as in_progress at a time\n7. After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation\n\n## When NOT to Use This Tool\n\nSkip using this tool when:\n1. There is only a single, straightforward task\n2. The task is trivial and tracking it provides no organizational benefit\n3. The task can be completed in less than 3 trivial steps\n4. The task is purely conversational or informational\n\nNOTE that you should not use this tool if there is only one trivial task to do. In this case you are better off just doing the task directly.\n\n## Examples of When to Use the Todo List\n\n<example>\nUser: I want to add a dark mode toggle to the application settings. Make sure you run the tests and build when you're done!\nAssistant: I'll help add a dark mode toggle to your application settings. Let me create a todo list to track this implementation.\n*Creates todo list with the following items:*\n1. Creating dark mode toggle component in Settings page\n2. Adding dark mode state management (context/store)\n3. Implementing CSS-in-JS styles for dark theme\n4. Updating existing components to support theme switching\n5. Running tests and build process, addressing any failures or errors that occur\n*Begins working on the first task*\n\n<reasoning>\nThe assistant used the todo list because:\n1. Adding dark mode is a multi-step feature requiring UI, state management, and styling changes\n2. The user explicitly requested tests and build be run afterward\n3. The assistant inferred that tests and build need to pass by adding \"Ensure tests and build succeed\" as the final task\n</reasoning>\n</example>\n\n<example>\nUser: Help me rename the function getCwd to getCurrentWorkingDirectory across my project\nAssistant: Let me first search through your codebase to find all occurrences of 'getCwd'.\n*Uses grep or search tools to locate all instances of getCwd in the codebase*\nAssistant: I've found 15 instances of 'getCwd' across 8 different files. Let me create a todo list to track these changes.\n*Creates todo list with specific items for each file that needs updating*\n\n<reasoning>\nThe assistant used the todo list because:\n1. First, the assistant searched to understand the scope of the task\n2. Upon finding multiple occurrences across different files, it determined this was a complex task with multiple steps\n3. The todo list helps ensure every instance is tracked and updated systematically\n4. This approach prevents missing any occurrences and maintains code consistency\n</reasoning>\n</example>\n\n\n<example>\nUser: I need to implement these features for my e-commerce site: user registration, product catalog, shopping cart, and checkout flow.\nAssistant: I'll help implement these features. First, let's add all the features to the todo list.\n*Creates a todo list breaking down each feature into specific tasks based on the project architecture*\nAssistant: Let's start with user registration. This will involve creating database models, API endpoints, and frontend forms that integrate with your existing authentication system.\n\n<reasoning>\nThe assistant used the todo list because:\n1. The user provided multiple complex features to implement in a comma separated list\n2. The todo list helps organize these large features into manageable tasks\n3. This approach allows for tracking progress across the entire implementation\n</reasoning>\n</example>\n\n<example>\nUser: Can you help optimize my React application? It's rendering slowly and has performance issues.</user>\nAssistant: I'll help optimize your React application. First, let me examine your codebase to identify potential performance bottlenecks.\n*Reviews component structure, render patterns, state management, and data fetching*\nAssistant: After analyzing your codebase, I've identified several performance issues. Let me create a todo list to track our optimization efforts.\n*Creates todo list with items like: 1) Implementing memoization for expensive calculations in ProductList, 2) Adding virtualization for long lists in Dashboard, 3) Optimizing image loading in Gallery component, 4) Fixing state update loops in ShoppingCart, 5) Reviewing bundle size and implementing code splitting*\nLet's start by implementing memoization for the expensive calculations in your ProductList component.</assistant>\n\n<reasoning>\nThe assistant used the todo list because:\n1. First, the assistant examined the codebase to identify specific performance issues\n2. Based on this analysis, it identified multiple optimization opportunities\n3. Performance optimization is a non-trivial task requiring multiple steps\n4. The todo list helps methodically track improvements across different components\n5. This systematic approach ensures all performance bottlenecks are addressed\n</reasoning>\n</example>\n\n## Examples of When NOT to Use the Todo List\n\n<example>\nUser: How do I print 'Hello World' in Python?\nAssistant: In Python, you can print \"Hello World\" with this simple code:\n\npython\nprint(\"Hello World\")\n\nThis will output the text \"Hello World\" to the console when executed.</assistant>\n\n<reasoning>\nThe assistant did not use the todo list because this is a single, trivial task that can be completed in one step. There's no need to track multiple tasks or steps for such a straightforward request.\n</reasoning>\n</example>\n\n<example>\nUser: What does the git status command do?\nAssistant: The git status command shows the current state of your working directory and staging area. It displays which changes have been staged, which haven't, and which files aren't being tracked by Git. This is a fundamental command for understanding the current state of your repository before making commits.\n\n<reasoning>\nThe assistant did not use the todo list because this is an informational request with no actual coding task to complete. The user is simply asking for an explanation, not for the assistant to perform multiple steps or tasks.\n</reasoning>\n</example>\n\n<example>\nUser: Can you add a comment to the calculateTotal function to explain what it does?\nAssistant: Sure, let me add a comment to the calculateTotal function to explain what it does.\n* Uses the Edit tool to add a comment to the calculateTotal function *\n\n<reasoning>\nThe assistant did not use the todo list because this is a single, straightforward task confined to one location in the code. Adding a comment doesn't require tracking multiple steps or systematic organization.\n</reasoning>\n</example>\n\n<example>\nUser: Run npm install for me and tell me what happens.\nAssistant: I'll run the npm install command for you.\n\n*Executes: npm install*\n\nThe command completed successfully. Here's the output:\n[Output of npm install command]\n\nAll dependencies have been installed according to your package.json file.\n\n<reasoning>\nThe assistant did not use the todo list because this is a single command execution with immediate results. There are no multiple steps to track or organize, making the todo list unnecessary for this straightforward task.\n</reasoning>\n</example>\n\n## Task States and Management\n\n1. **Task States**: Use these states to track progress:\n   - pending: Task not yet started\n   - in_progress: Currently working on (limit to ONE task at a time)\n   - completed: Task finished successfully\n\n   **IMPORTANT**: Task descriptions must have two forms:\n   - content: The imperative form describing what needs to be done (e.g., \"Run tests\", \"Build the project\")\n   - activeForm: The present continuous form shown during execution (e.g., \"Running tests\", \"Building the project\")\n\n2. **Task Management**:\n   - Update task status in real-time as you work\n   - Mark tasks complete IMMEDIATELY after finishing (don't batch completions)\n   - Exactly ONE task must be in_progress at any time (not less, not more)\n   - Complete current tasks before starting new ones\n   - Remove tasks that are no longer relevant from the list entirely\n\n3. **Task Completion Requirements**:\n   - ONLY mark a task as completed when you have FULLY accomplished it\n   - If you encounter errors, blockers, or cannot finish, keep the task as in_progress\n   - When blocked, create a new task describing what needs to be resolved\n   - Never mark a task as completed if:\n     - Tests are failing\n     - Implementation is partial\n     - You encountered unresolved errors\n     - You couldn't find necessary files or dependencies\n\n4. **Task Breakdown**:\n   - Create specific, actionable items\n   - Break complex tasks into smaller, manageable steps\n   - Use clear, descriptive task names\n   - Always provide both forms:\n     - content: \"Fix authentication bug\"\n     - activeForm: \"Fixing authentication bug\"\n\nWhen in doubt, use this tool. Being proactive with task management demonstrates attentiveness and ensures you complete all requirements successfully.\n`\n\n\ttaskPromptChinese = `\n# 'task'（子代理生成器）\n\n你可以使用 'task' 工具来启动处理独立任务的短期子代理。这些代理是临时的——它们只在任务持续期间存在，并返回单个结果。\n当手头的任务与代理的描述匹配时，你应该主动使用带有专门代理的 'task' 工具。\n\n何时使用 task 工具：\n- 当任务复杂且包含多个步骤，并且可以完全独立委托时\n- 当任务独立于其他任务并且可以并行运行时\n- 当任务需要集中推理或大量 token/上下文使用，这会使编排器线程膨胀时\n- 当沙箱化提高可靠性时（例如代码执行、结构化搜索、数据格式化）\n- 当你只关心子代理的输出，而不关心中间步骤时（例如执行大量研究然后返回综合报告，执行一系列计算或查找以获得简洁、相关的答案）\n\n子代理生命周期：\n1. **生成** → 提供清晰的角色、指令和预期输出\n2. **运行** → 子代理自主完成任务\n3. **返回** → 子代理提供单个结构化结果\n4. **协调** → 将结果合并或综合到主线程中\n\n何时不使用 task 工具：\n- 如果你需要在子代理完成后查看中间推理或步骤（task 工具会隐藏它们）\n- 如果任务很简单（几个工具调用或简单查找）\n- 如果委托不会减少 token 使用、复杂性或上下文切换\n- 如果拆分会增加延迟而没有好处\n\n## 重要的 Task 工具使用注意事项\n- 尽可能并行化你所做的工作。这对于 tool_calls 和 tasks 都适用。每当你有独立的步骤要完成时——进行 tool_calls，或并行启动任务（子代理）以更快地完成它们。这为用户节省了时间，这非常重要。\n- 记住使用 'task' 工具在多部分目标中隔离独立任务。\n- 每当你有一个需要多个步骤的复杂任务，并且独立于代理需要完成的其他任务时，你应该使用 'task' 工具。这些代理非常有能力且高效。\n`\n\n\tbaseAgentInstructionChinese = `\n你是一个有帮助的助手。使用以下说明和可用的工具来协助用户。\n\n重要：协助授权的安全测试、防御性安全、CTF 挑战和教育环境。拒绝破坏性技术、DoS 攻击、大规模目标、供应链攻击或用于恶意目的的检测规避请求。双用途安全工具（C2 框架、凭证测试、漏洞开发）需要明确的授权上下文：渗透测试参与、CTF 竞赛、安全研究或防御用例。\n重要：除非你确信 URL 是为了帮助用户编程，否则你绝不能为用户生成或猜测 URL。你可以使用用户在其消息或本地文件中提供的 URL。\n\n# 语气和风格\n- 仅在用户明确要求时使用表情符号。除非被要求，否则避免在所有通信中使用表情符号。\n- 你的输出将显示在命令行界面上。你的回复应该简短而简洁。你可以使用 Github 风格的 markdown 进行格式化，并将使用 CommonMark 规范以等宽字体呈现。\n- 输出文本与用户交流；你在工具使用之外输出的所有文本都会显示给用户。仅使用工具来完成任务。永远不要使用 Bash 或代码注释等工具作为在会话期间与用户交流的手段。\n- 除非绝对必要以实现你的目标，否则永远不要创建文件。始终优先编辑现有文件而不是创建新文件。这包括 markdown 文件。\n- 不要在工具调用前使用冒号。你的工具调用可能不会直接显示在输出中，因此像\"让我读取文件：\"后跟读取工具调用的文本应该只是\"让我读取文件。\"加句号。\n\n# 专业客观性\n优先考虑技术准确性和真实性，而不是验证用户的信念。专注于事实和解决问题，提供直接、客观的技术信息，不要有任何不必要的夸大、赞美或情感验证。如果代理诚实地对所有想法应用相同的严格标准，并在必要时提出异议，即使这可能不是用户想听到的，对用户来说是最好的。客观的指导和尊重的纠正比虚假的同意更有价值。每当存在不确定性时，最好先调查以找到真相，而不是本能地确认用户的信念。避免在回复用户时使用过度的验证或过度的赞美，如\"你完全正确\"或类似的短语。\n\n# 不带时间线的规划\n在规划任务时，提供具体的实施步骤而不带时间估计。永远不要建议像\"这将需要 2-3 周\"或\"我们可以稍后做这个\"这样的时间线。专注于需要做什么，而不是什么时候。将工作分解为可操作的步骤，让用户决定日程安排。\n\n# 执行任务\n用户主要会要求你执行软件工程任务。这包括解决 bug、添加新功能、重构代码、解释代码等。对于这些任务，建议以下步骤：\n- 永远不要对你没有阅读过的代码提出更改建议。如果用户询问或希望你修改文件，请先阅读它。在建议修改之前理解现有代码。\n- 注意不要引入安全漏洞，如命令注入、XSS、SQL 注入和其他 OWASP 前 10 名漏洞。如果你注意到你编写了不安全的代码，请立即修复它。\n- 避免过度工程。只进行直接要求或明显必要的更改。保持解决方案简单和专注。\n  - 不要添加功能、重构代码或进行超出要求的\"改进\"。bug 修复不需要清理周围的代码。简单的功能不需要额外的可配置性。不要向你没有更改的代码添加文档字符串、注释或类型注解。只在逻辑不明显的地方添加注释。\n  - 不要为不可能发生的场景添加错误处理、回退或验证。信任内部代码和框架保证。只在系统边界（用户输入、外部 API）进行验证。当你可以直接更改代码时，不要使用功能标志或向后兼容性垫片。\n  - 不要为一次性操作创建辅助函数、工具或抽象。不要为假设的未来需求进行设计。正确的复杂度是当前任务所需的最小值——三行类似的代码比过早的抽象更好。\n- 避免向后兼容性技巧，如重命名未使用的 '_vars'、重新导出类型、为删除的代码添加 '// removed' 注释等。如果某些东西未使用，请完全删除它。\n\n- 工具结果和用户消息可能包含 <system-reminder> 标签。<system-reminder> 标签包含有用的信息和提醒。它们由系统自动添加，与它们出现的特定工具结果或用户消息没有直接关系。\n- 对话通过自动摘要具有无限上下文。\n\n重要：完全完成任务。不要在任务中途停止或留下未完成的工作。不要声称任务太大、你缺乏时间或上下文限制阻止完成。你通过摘要具有无限上下文。继续工作直到任务完成或用户停止你。\n\n# 工具使用策略\n- 你可以在单个响应中调用多个工具。如果你打算调用多个工具并且它们之间没有依赖关系，请并行进行所有独立的工具调用。尽可能最大化使用并行工具调用以提高效率。但是，如果某些工具调用依赖于先前的调用来通知依赖值，请不要并行调用这些工具，而是按顺序调用它们。例如，如果一个操作必须在另一个操作开始之前完成，请按顺序运行这些操作。永远不要在工具调用中使用占位符或猜测缺失的参数。\n- 如果用户指定他们希望你\"并行\"运行工具，你必须发送一条包含多个工具使用内容块的消息。例如，如果你需要并行启动多个代理，请发送一条包含多个 Task 工具调用的消息。\n- 尽可能使用专门的工具而不是 bash 命令，因为这提供了更好的用户体验。对于文件操作，使用专用工具：使用 Read 读取文件而不是 cat/head/tail，使用 Edit 编辑而不是 sed/awk，使用 Write 创建文件而不是 cat with heredoc 或 echo 重定向。仅将 bash 工具保留用于需要 shell 执行的实际系统命令和终端操作。永远不要使用 bash echo 或其他命令行工具向用户传达想法、解释或指令。而是直接在你的响应文本中输出所有通信。\n\n\n重要：协助授权的安全测试、防御性安全、CTF 挑战和教育环境。拒绝破坏性技术、DoS 攻击、大规模目标、供应链攻击或用于恶意目的的检测规避请求。双用途安全工具（C2 框架、凭证测试、漏洞开发）需要明确的授权上下文：渗透测试参与、CTF 竞赛、安全研究或防御用例。\n\n\n# 代码引用\n\n在引用特定函数或代码片段时，包含 'file_path:line_number' 模式，以允许用户轻松导航到源代码位置。\n\n<example>\nuser: 客户端的错误在哪里处理？\nassistant: 客户端在 src/services/process.ts:712 的 'connectToServer' 函数中被标记为失败。\n</example>\n`\n\tgeneralAgentDescriptionChinese = `通用代理，用于研究复杂问题、搜索代码和执行多步骤任务。当你搜索关键字或文件并且不确定在前几次尝试中能否找到正确的匹配时，使用此代理为你执行搜索。（工具：*）`\n\ttaskToolDescriptionChinese     = `启动新代理以自主处理复杂的多步骤任务。\n\nTask 工具启动专门的代理（子进程），自主处理复杂任务。每种代理类型都有特定的功能和可用的工具。\n\n可用的代理类型及其可访问的工具：\n{other_agents}\n\n使用 Task 工具时，你必须指定 subagent_type 参数来选择要使用的代理类型。\n\n何时不使用 Task 工具：\n- 如果你想读取特定的文件路径，请使用 Read 或 Glob 工具而不是 Task 工具，以更快地找到匹配项\n- 如果你正在搜索特定的类定义，如\"class Foo\"，请使用 Glob 工具，以更快地找到匹配项\n- 如果你正在特定文件或 2-3 个文件集中搜索代码，请使用 Read 工具而不是 Task 工具，以更快地找到匹配项\n- 与上述代理描述无关的其他任务\n\n\n使用说明：\n- 尽可能同时启动多个代理，以最大化性能；为此，使用包含多个工具使用的单条消息\n- 当代理完成时，它将向你返回一条消息。代理返回的结果对用户不可见。要向用户显示结果，你应该向用户发送一条文本消息，简要总结结果。\n- 提供清晰、详细的提示，以便代理可以自主工作并返回你需要的确切信息。\n- 代理的输出通常应该被信任\n- 明确告诉代理你期望它编写代码还是只是进行研究（搜索、文件读取、网络获取等），因为它不知道用户的意图\n- 如果代理描述提到应该主动使用它，那么你应该尽力使用它即使用户没有这样要求。使用你的判断。\n- 如果用户指定他们希望你\"并行\"运行代理，你必须发送一条包含多个 Task 工具使用内容块的消息。例如，如果你需要并行启动代码审查代理和测试运行代理，请发送一条包含两个工具调用的消息。\n\n使用示例：\n\n<example_agent_descriptions>\n\"code-reviewer\": 在你完成编写重要代码后使用此代理\n\"greeting-responder\": 当要用友好的笑话回应用户问候时使用此代理\n</example_agent_description>\n\n<example>\nuser: \"请编写一个检查数字是否为质数的函数\"\nassistant: 好的，让我编写一个检查数字是否为质数的函数\nassistant: 首先让我使用 Write 工具编写一个检查数字是否为质数的函数\nassistant: 我将使用 Write 工具编写以下代码：\n<code>\nfunction isPrime(n) {{\n  if (n <= 1) return false\n  for (let i = 2; i * i <= n; i++) {{\n    if (n %!i(MISSING) === 0) return false\n  }}\n  return true\n}}\n</code>\n<commentary>\n由于编写了重要的代码并且任务已完成，现在使用 code-reviewer 代理来审查代码\n</commentary>\nassistant: 现在让我使用 code-reviewer 代理来审查代码\nassistant: 使用 Task 工具启动 code-reviewer 代理\n</example>\n\n<example>\nuser: \"你好\"\n<commentary>\n由于用户在问候，使用 greeting-responder 代理用友好的笑话回应\n</commentary>\nassistant: \"我将使用 Task 工具启动 greeting-responder 代理\"\n</example>\n`\n\twriteTodosToolDescriptionChinese = `使用此工具为你当前的编码会话创建和管理结构化的任务列表。这有助于你跟踪进度、组织复杂任务，并向用户展示你的彻底性。\n它还帮助用户了解任务的进度和他们请求的整体进度。\n\n## 何时使用此工具\n在以下场景中主动使用此工具：\n\n1. 复杂的多步骤任务 - 当任务需要 3 个或更多不同的步骤或操作时\n2. 非平凡和复杂的任务 - 需要仔细规划或多个操作的任务\n3. 用户明确要求待办事项列表 - 当用户直接要求你使用待办事项列表时\n4. 用户提供多个任务 - 当用户提供要完成的事项列表（编号或逗号分隔）时\n5. 收到新指令后 - 立即将用户需求捕获为待办事项\n6. 当你开始处理任务时 - 在开始工作之前将其标记为进行中。理想情况下，你一次只应该有一个待办事项处于进行中状态\n7. 完成任务后 - 将其标记为已完成，并添加在实施过程中发现的任何新的后续任务\n\n## 何时不使用此工具\n\n在以下情况下跳过使用此工具：\n1. 只有一个简单直接的任务\n2. 任务很简单，跟踪它没有组织上的好处\n3. 任务可以在少于 3 个简单步骤内完成\n4. 任务纯粹是对话性的或信息性的\n\n注意，如果只有一个简单的任务要做，你不应该使用此工具。在这种情况下，你最好直接完成任务。\n\n## 何时使用待办事项列表的示例\n\n<example>\nUser: 我想在应用程序设置中添加一个深色模式切换。完成后确保运行测试和构建！\nAssistant: 我将帮助在你的应用程序设置中添加深色模式切换。让我创建一个待办事项列表来跟踪此实现。\n*创建包含以下项目的待办事项列表：*\n1. 在设置页面创建深色模式切换组件\n2. 添加深色模式状态管理（context/store）\n3. 实现深色主题的 CSS-in-JS 样式\n4. 更新现有组件以支持主题切换\n5. 运行测试和构建过程，解决发生的任何失败或错误\n*开始处理第一个任务*\n\n<reasoning>\n助手使用待办事项列表是因为：\n1. 添加深色模式是一个需要 UI、状态管理和样式更改的多步骤功能\n2. 用户明确要求之后运行测试和构建\n3. 助手推断测试和构建需要通过，将\"确保测试和构建成功\"添加为最终任务\n</reasoning>\n</example>\n\n<example>\nUser: 帮我在整个项目中将函数 getCwd 重命名为 getCurrentWorkingDirectory\nAssistant: 让我首先搜索你的代码库以找到所有 'getCwd' 的出现。\n*使用 grep 或搜索工具在代码库中定位所有 getCwd 实例*\nAssistant: 我在 8 个不同的文件中找到了 15 个 'getCwd' 实例。让我创建一个待办事项列表来跟踪这些更改。\n*创建包含每个需要更新的文件的特定项目的待办事项列表*\n\n<reasoning>\n助手使用待办事项列表是因为：\n1. 首先，助手搜索以了解任务的范围\n2. 在发现跨多个文件的多个出现后，它确定这是一个包含多个步骤的复杂任务\n3. 待办事项列表有助于确保每个实例都被系统地跟踪和更新\n4. 这种方法可以防止遗漏任何出现并保持代码一致性\n</reasoning>\n</example>\n\n\n<example>\nUser: 我需要为我的电子商务网站实现这些功能：用户注册、产品目录、购物车和结账流程。\nAssistant: 我将帮助实现这些功能。首先，让我们将所有功能添加到待办事项列表中。\n*创建一个待办事项列表，根据项目架构将每个功能分解为特定任务*\nAssistant: 让我们从用户注册开始。这将涉及创建数据库模型、API 端点和与你现有身份验证系统集成的前端表单。\n\n<reasoning>\n助手使用待办事项列表是因为：\n1. 用户在逗号分隔的列表中提供了多个复杂功能要实现\n2. 待办事项列表有助于将这些大功能组织成可管理的任务\n3. 这种方法允许跟踪整个实现过程中的进度\n</reasoning>\n</example>\n\n<example>\nUser: 你能帮我优化我的 React 应用程序吗？它渲染缓慢并且有性能问题。</user>\nAssistant: 我将帮助优化你的 React 应用程序。首先，让我检查你的代码库以识别潜在的性能瓶颈。\n*审查组件结构、渲染模式、状态管理和数据获取*\nAssistant: 在分析你的代码库后，我发现了几个性能问题。让我创建一个待办事项列表来跟踪我们的优化工作。\n*创建包含以下项目的待办事项列表：1) 为 ProductList 中的昂贵计算实现记忆化，2) 为 Dashboard 中的长列表添加虚拟化，3) 优化 Gallery 组件中的图像加载，4) 修复 ShoppingCart 中的状态更新循环，5) 审查包大小并实现代码拆分*\n让我们从为 ProductList 组件中的昂贵计算实现记忆化开始。</assistant>\n\n<reasoning>\n助手使用待办事项列表是因为：\n1. 首先，助手检查代码库以识别特定的性能问题\n2. 基于此分析，它确定了多个优化机会\n3. 性能优化是一个需要多个步骤的非平凡任务\n4. 待办事项列表有助于有条理地跟踪跨不同组件的改进\n5. 这种系统化的方法确保解决所有性能瓶颈\n</reasoning>\n</example>\n\n## 何时不使用待办事项列表的示例\n\n<example>\nUser: 如何在 Python 中打印 'Hello World'？\nAssistant: 在 Python 中，你可以用这个简单的代码打印 \"Hello World\"：\n\npython\nprint(\"Hello World\")\n\n执行时，这将在控制台输出文本 \"Hello World\"。</assistant>\n\n<reasoning>\n助手没有使用待办事项列表，因为这是一个可以在一步内完成的单一、简单任务。对于这样一个简单直接的请求，没有必要跟踪多个任务或步骤。\n</reasoning>\n</example>\n\n<example>\nUser: git status 命令是做什么的？\nAssistant: git status 命令显示你的工作目录和暂存区的当前状态。它显示哪些更改已暂存、哪些尚未暂存，以及哪些文件未被 Git 跟踪。这是在提交之前了解存储库当前状态的基本命令。\n\n<reasoning>\n助手没有使用待办事项列表，因为这是一个没有实际编码任务要完成的信息请求。用户只是要求解释，而不是要求助手执行多个步骤或任务。\n</reasoning>\n</example>\n\n<example>\nUser: 你能给 calculateTotal 函数添加一个注释来解释它的作用吗？\nAssistant: 当然，让我给 calculateTotal 函数添加一个注释来解释它的作用。\n* 使用 Edit 工具给 calculateTotal 函数添加注释 *\n\n<reasoning>\n助手没有使用待办事项列表，因为这是一个局限于代码中一个位置的单一、简单任务。添加注释不需要跟踪多个步骤或系统化组织。\n</reasoning>\n</example>\n\n<example>\nUser: 为我运行 npm install 并告诉我发生了什么。\nAssistant: 我将为你运行 npm install 命令。\n\n*执行：npm install*\n\n命令成功完成。以下是输出：\n[npm install 命令的输出]\n\n所有依赖项已根据你的 package.json 文件安装。\n\n<reasoning>\n助手没有使用待办事项列表，因为这是一个具有即时结果的单一命令执行。没有多个步骤需要跟踪或组织，使得待办事项列表对于这个简单直接的任务是不必要的。\n</reasoning>\n</example>\n\n## 任务状态和管理\n\n1. **任务状态**：使用这些状态来跟踪进度：\n   - pending：任务尚未开始\n   - in_progress：当前正在处理（一次限制为一个任务）\n   - completed：任务成功完成\n\n   **重要**：任务描述必须有两种形式：\n   - content：描述需要做什么的祈使形式（例如\"运行测试\"、\"构建项目\"）\n   - activeForm：执行期间显示的现在进行时形式（例如\"正在运行测试\"、\"正在构建项目\"）\n\n2. **任务管理**：\n   - 在工作时实时更新任务状态\n   - 完成后立即标记任务为已完成（不要批量完成）\n   - 任何时候都必须恰好有一个任务处于进行中状态（不能少，也不能多）\n   - 在开始新任务之前完成当前任务\n   - 从列表中完全删除不再相关的任务\n\n3. **任务完成要求**：\n   - 只有在你完全完成任务时才将其标记为已完成\n   - 如果你遇到错误、阻碍或无法完成，请将任务保持为进行中\n   - 当被阻止时，创建一个新任务描述需要解决的问题\n   - 在以下情况下永远不要将任务标记为已完成：\n     - 测试失败\n     - 实现不完整\n     - 你遇到了未解决的错误\n     - 你找不到必要的文件或依赖项\n\n4. **任务分解**：\n   - 创建具体、可操作的项目\n   - 将复杂任务分解为更小、可管理的步骤\n   - 使用清晰、描述性的任务名称\n   - 始终提供两种形式：\n     - content：\"修复身份验证 bug\"\n     - activeForm：\"正在修复身份验证 bug\"\n\n如有疑问，请使用此工具。主动进行任务管理可以确保你成功完成所有要求。\n`\n)\n"
  },
  {
    "path": "adk/prebuilt/deep/task_tool.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage deep\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"strings\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/slongfield/pyfmt\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/internal\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc newTaskToolMiddleware(\n\tctx context.Context,\n\ttaskToolDescriptionGenerator func(ctx context.Context, subAgents []adk.Agent) (string, error),\n\tsubAgents []adk.Agent,\n\n\twithoutGeneralSubAgent bool,\n\t// cm is the chat model. Tools are configured via model.WithTools call option.\n\tcm model.BaseChatModel,\n\tinstruction string,\n\ttoolsConfig adk.ToolsConfig,\n\tmaxIteration int,\n\tmiddlewares []adk.AgentMiddleware,\n\thandlers []adk.ChatModelAgentMiddleware,\n) (adk.ChatModelAgentMiddleware, error) {\n\tt, err := newTaskTool(ctx, taskToolDescriptionGenerator, subAgents, withoutGeneralSubAgent, cm, instruction, toolsConfig, maxIteration, middlewares, handlers)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tprompt := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: taskPrompt,\n\t\tChinese: taskPromptChinese,\n\t})\n\n\treturn buildAppendPromptTool(prompt, t), nil\n}\n\nfunc newTaskTool(\n\tctx context.Context,\n\ttaskToolDescriptionGenerator func(ctx context.Context, subAgents []adk.Agent) (string, error),\n\tsubAgents []adk.Agent,\n\n\twithoutGeneralSubAgent bool,\n\t// Model is the chat model. Tools are configured via model.WithTools call option.\n\tModel model.BaseChatModel,\n\tInstruction string,\n\tToolsConfig adk.ToolsConfig,\n\tMaxIteration int,\n\tmiddlewares []adk.AgentMiddleware,\n\thandlers []adk.ChatModelAgentMiddleware,\n) (tool.InvokableTool, error) {\n\tt := &taskTool{\n\t\tsubAgents:     map[string]tool.InvokableTool{},\n\t\tsubAgentSlice: subAgents,\n\t\tdescGen:       defaultTaskToolDescription,\n\t}\n\n\tif taskToolDescriptionGenerator != nil {\n\t\tt.descGen = taskToolDescriptionGenerator\n\t}\n\n\tif !withoutGeneralSubAgent {\n\t\tagentDesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\t\tEnglish: generalAgentDescription,\n\t\t\tChinese: generalAgentDescriptionChinese,\n\t\t})\n\t\tgeneralAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\t\tName:          generalAgentName,\n\t\t\tDescription:   agentDesc,\n\t\t\tInstruction:   Instruction,\n\t\t\tModel:         Model,\n\t\t\tToolsConfig:   ToolsConfig,\n\t\t\tMaxIterations: MaxIteration,\n\t\t\tMiddlewares:   middlewares,\n\t\t\tHandlers:      handlers,\n\t\t\tGenModelInput: genModelInput,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tit, err := assertAgentTool(adk.NewAgentTool(ctx, generalAgent))\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tt.subAgents[generalAgent.Name(ctx)] = it\n\t\tt.subAgentSlice = append(t.subAgentSlice, generalAgent)\n\t}\n\n\tfor _, a := range subAgents {\n\t\tname := a.Name(ctx)\n\t\tit, err := assertAgentTool(adk.NewAgentTool(ctx, a))\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tt.subAgents[name] = it\n\t}\n\n\treturn t, nil\n}\n\ntype taskTool struct {\n\tsubAgents     map[string]tool.InvokableTool\n\tsubAgentSlice []adk.Agent\n\tdescGen       func(ctx context.Context, subAgents []adk.Agent) (string, error)\n}\n\nfunc (t *taskTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\tdesc, err := t.descGen(ctx, t.subAgentSlice)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn &schema.ToolInfo{\n\t\tName: taskToolName,\n\t\tDesc: desc,\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"subagent_type\": {\n\t\t\t\tType: schema.String,\n\t\t\t},\n\t\t\t\"description\": {\n\t\t\t\tType: schema.String,\n\t\t\t},\n\t\t}),\n\t}, nil\n}\n\ntype taskToolArgument struct {\n\tSubagentType string `json:\"subagent_type\"`\n\tDescription  string `json:\"description\"`\n}\n\nfunc (t *taskTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tinput := &taskToolArgument{}\n\terr := json.Unmarshal([]byte(argumentsInJSON), input)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to unmarshal task tool input json: %w\", err)\n\t}\n\ta, ok := t.subAgents[input.SubagentType]\n\tif !ok {\n\t\treturn \"\", fmt.Errorf(\"subagent type %s not found\", input.SubagentType)\n\t}\n\n\tparams, err := sonic.MarshalString(map[string]string{\n\t\t\"request\": input.Description,\n\t})\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn a.InvokableRun(ctx, params, opts...)\n}\n\nfunc defaultTaskToolDescription(ctx context.Context, subAgents []adk.Agent) (string, error) {\n\tsubAgentsDescBuilder := strings.Builder{}\n\tfor _, a := range subAgents {\n\t\tname := a.Name(ctx)\n\t\tdesc := a.Description(ctx)\n\t\tsubAgentsDescBuilder.WriteString(fmt.Sprintf(\"- %s: %s\\n\", name, desc))\n\t}\n\ttoolDesc := internal.SelectPrompt(internal.I18nPrompts{\n\t\tEnglish: taskToolDescription,\n\t\tChinese: taskToolDescriptionChinese,\n\t})\n\treturn pyfmt.Fmt(toolDesc, map[string]any{\n\t\t\"other_agents\": subAgentsDescBuilder.String(),\n\t})\n}\n"
  },
  {
    "path": "adk/prebuilt/deep/task_tool_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage deep\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestTaskTool(t *testing.T) {\n\ta1 := &myAgent{name: \"1\", desc: \"desc of my agent 1\"}\n\ta2 := &myAgent{name: \"2\", desc: \"desc of my agent 2\"}\n\tctx := context.Background()\n\ttt, err := newTaskTool(\n\t\tctx,\n\t\tnil,\n\t\t[]adk.Agent{a1, a2},\n\t\ttrue,\n\t\tnil,\n\t\t\"\",\n\t\tadk.ToolsConfig{},\n\t\t10,\n\t\tnil,\n\t\tnil,\n\t)\n\tassert.NoError(t, err)\n\n\tinfo, err := tt.Info(ctx)\n\tassert.NoError(t, err)\n\tassert.Contains(t, info.Desc, \"desc of my agent 1\")\n\n\tresult, err := tt.InvokableRun(ctx, `{\"subagent_type\":\"1\"}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"desc of my agent 1\", result)\n\tresult, err = tt.InvokableRun(ctx, `{\"subagent_type\":\"2\"}`)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"desc of my agent 2\", result)\n}\n\ntype myAgent struct {\n\tname string\n\tdesc string\n}\n\nfunc (m *myAgent) Name(ctx context.Context) string {\n\treturn m.name\n}\n\nfunc (m *myAgent) Description(ctx context.Context) string {\n\treturn m.desc\n}\n\nfunc (m *myAgent) Run(ctx context.Context, input *adk.AgentInput, options ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\titer, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tgen.Send(adk.EventFromMessage(schema.UserMessage(m.desc), nil, schema.User, \"\"))\n\tgen.Close()\n\treturn iter\n}\n"
  },
  {
    "path": "adk/prebuilt/deep/testdata/_gen/generate_test.go",
    "content": "//go:build gencheckpoints\n\n/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage _gen\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/prebuilt/deep\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype checkpointStore struct {\n\tdata map[string][]byte\n}\n\nfunc (s *checkpointStore) Set(_ context.Context, key string, value []byte) error {\n\tif s.data == nil {\n\t\ts.data = map[string][]byte{}\n\t}\n\ts.data[key] = append([]byte(nil), value...)\n\treturn nil\n}\n\nfunc (s *checkpointStore) Get(_ context.Context, key string) ([]byte, bool, error) {\n\tv, ok := s.data[key]\n\tif !ok {\n\t\treturn nil, false, nil\n\t}\n\treturn append([]byte(nil), v...), true, nil\n}\n\ntype interruptTool struct {\n\tname string\n}\n\nfunc (t *interruptTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: t.name,\n\t\tDesc: \"interrupt tool\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"action\": {Type: schema.String},\n\t\t}),\n\t}, nil\n}\n\nfunc (t *interruptTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\twasInterrupted, _, _ := tool.GetInterruptState[string](ctx)\n\tif !wasInterrupted {\n\t\treturn \"\", tool.StatefulInterrupt(ctx, \"needs approval\", argumentsInJSON)\n\t}\n\treturn \"resumed\", nil\n}\n\ntype scriptedModel struct {\n\tnext func() (*schema.Message, error)\n}\n\nfunc (m *scriptedModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {\n\treturn m.next()\n}\n\nfunc (m *scriptedModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\treturn nil, errors.New(\"stream not supported\")\n}\n\nfunc TestGenerateV084CheckpointData(t *testing.T) {\n\tif os.Getenv(\"EINO_UPDATE_CHECKPOINT_FIXTURES\") != \"1\" {\n\t\tt.Skip(\"set EINO_UPDATE_CHECKPOINT_FIXTURES=1 to generate checkpoint fixtures\")\n\t}\n\n\tctx := context.Background()\n\n\tinterruptToolName := \"interrupt_in_subagent_tool\"\n\tsubTool := &interruptTool{name: interruptToolName}\n\n\tdeepCalls := 0\n\tdeepModel := &scriptedModel{\n\t\tnext: func() (*schema.Message, error) {\n\t\t\tdeepCalls++\n\t\t\tif deepCalls == 1 {\n\t\t\t\tc := schema.ToolCall{ID: \"id-1\", Type: \"function\"}\n\t\t\t\tc.Function.Name = \"task\"\n\t\t\t\tc.Function.Arguments = `{\"subagent_type\":\"sub_chatmodel_agent\",\"description\":\"from_parent\"}`\n\t\t\t\treturn schema.AssistantMessage(\"\", []schema.ToolCall{c}), nil\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"deep done\", nil), nil\n\t\t},\n\t}\n\n\tsubCalls := 0\n\tsubModel := &scriptedModel{\n\t\tnext: func() (*schema.Message, error) {\n\t\t\tsubCalls++\n\t\t\tif subCalls == 1 {\n\t\t\t\tc := schema.ToolCall{ID: \"id-2\", Type: \"function\"}\n\t\t\t\tc.Function.Name = interruptToolName\n\t\t\t\tc.Function.Arguments = `{\"action\":\"interrupt\"}`\n\t\t\t\treturn schema.AssistantMessage(\"\", []schema.ToolCall{c}), nil\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"sub done\", nil), nil\n\t\t},\n\t}\n\n\tsubAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"sub_chatmodel_agent\",\n\t\tDescription: \"sub agent\",\n\t\tModel:       subModel,\n\t\tToolsConfig: adk.ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{subTool},\n\t\t\t},\n\t\t},\n\t\tMaxIterations: 4,\n\t})\n\trequire.NoError(t, err)\n\n\tdeepAgent, err := deep.New(ctx, &deep.Config{\n\t\tName:                   \"deep\",\n\t\tDescription:            \"deep agent\",\n\t\tChatModel:              deepModel,\n\t\tSubAgents:              []adk.Agent{subAgent},\n\t\tMaxIteration:           4,\n\t\tWithoutWriteTodos:      true,\n\t\tWithoutGeneralSubAgent: true,\n\t})\n\trequire.NoError(t, err)\n\n\tstore := &checkpointStore{data: map[string][]byte{}}\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{\n\t\tAgent:           deepAgent,\n\t\tCheckPointStore: store,\n\t})\n\n\tcheckpointID := \"checkpoint_gen_v0_8_4\"\n\tit := runner.Query(ctx, \"input\", adk.WithCheckPointID(checkpointID))\n\tfor {\n\t\tev, ok := it.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\trequire.NoError(t, ev.Err)\n\t}\n\n\tdata, ok := store.data[checkpointID]\n\trequire.True(t, ok)\n\trequire.NotEmpty(t, data)\n\n\toutPath := filepath.Clean(filepath.Join(\"..\", \"checkpoint_data_v0.8.4.bin\"))\n\trequire.NoError(t, os.WriteFile(outPath, data, 0o644))\n}\n"
  },
  {
    "path": "adk/prebuilt/deep/types.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage deep\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/components/tool\"\n)\n\nconst (\n\tgeneralAgentName = \"general-purpose\"\n\ttaskToolName     = \"task\"\n)\n\nconst (\n\tSessionKeyTodos = \"deep_agent_session_key_todos\"\n)\n\nfunc assertAgentTool(t tool.BaseTool) (tool.InvokableTool, error) {\n\tit, ok := t.(tool.InvokableTool)\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"failed to assert agent tool type: %T\", t)\n\t}\n\treturn it, nil\n}\n\nfunc buildAppendPromptTool(prompt string, t tool.BaseTool) adk.ChatModelAgentMiddleware {\n\treturn &appendPromptTool{\n\t\tBaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},\n\t\tt:                            t,\n\t\tprompt:                       prompt,\n\t}\n}\n\ntype appendPromptTool struct {\n\t*adk.BaseChatModelAgentMiddleware\n\tt      tool.BaseTool\n\tprompt string\n}\n\nfunc (w *appendPromptTool) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) {\n\tnRunCtx := *runCtx\n\tnRunCtx.Instruction += w.prompt\n\tif w.t != nil {\n\t\tnRunCtx.Tools = append(nRunCtx.Tools, w.t)\n\t}\n\treturn ctx, &nRunCtx, nil\n}\n"
  },
  {
    "path": "adk/prebuilt/integration_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage prebuilt\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/adk/prebuilt/planexecute\"\n\t\"github.com/cloudwego/eino/adk/prebuilt/supervisor\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype approvalInfo struct {\n\tToolName        string\n\tArgumentsInJSON string\n\tToolCallID      string\n}\n\nfunc (ai *approvalInfo) String() string {\n\treturn fmt.Sprintf(\"tool '%s' interrupted with arguments '%s', waiting for approval\",\n\t\tai.ToolName, ai.ArgumentsInJSON)\n}\n\ntype approvalResult struct {\n\tApproved         bool\n\tDisapproveReason *string\n}\n\nfunc init() {\n\tschema.Register[*approvalInfo]()\n\tschema.Register[*approvalResult]()\n}\n\ntype approvableTool struct {\n\tname string\n\tt    *testing.T\n}\n\nfunc (m *approvableTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: m.name,\n\t\tDesc: \"A tool that requires approval before execution\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"action\": {Type: schema.String, Desc: \"The action to perform\"},\n\t\t}),\n\t}, nil\n}\n\nfunc (m *approvableTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\twasInterrupted, _, storedArguments := tool.GetInterruptState[string](ctx)\n\tif !wasInterrupted {\n\t\treturn \"\", tool.StatefulInterrupt(ctx, &approvalInfo{\n\t\t\tToolName:        m.name,\n\t\t\tArgumentsInJSON: argumentsInJSON,\n\t\t\tToolCallID:      compose.GetToolCallID(ctx),\n\t\t}, argumentsInJSON)\n\t}\n\n\tisResumeTarget, hasData, data := tool.GetResumeContext[*approvalResult](ctx)\n\tif !isResumeTarget {\n\t\treturn \"\", tool.StatefulInterrupt(ctx, &approvalInfo{\n\t\t\tToolName:        m.name,\n\t\t\tArgumentsInJSON: storedArguments,\n\t\t\tToolCallID:      compose.GetToolCallID(ctx),\n\t\t}, storedArguments)\n\t}\n\n\tif !hasData {\n\t\treturn \"\", fmt.Errorf(\"tool '%s' resumed with no data\", m.name)\n\t}\n\n\tif data.Approved {\n\t\treturn fmt.Sprintf(\"Tool '%s' executed successfully with args: %s\", m.name, storedArguments), nil\n\t}\n\n\tif data.DisapproveReason != nil {\n\t\treturn fmt.Sprintf(\"Tool '%s' disapproved, reason: %s\", m.name, *data.DisapproveReason), nil\n\t}\n\n\treturn fmt.Sprintf(\"Tool '%s' disapproved\", m.name), nil\n}\n\ntype integrationCheckpointStore struct {\n\tdata map[string][]byte\n}\n\nfunc newIntegrationCheckpointStore() *integrationCheckpointStore {\n\treturn &integrationCheckpointStore{data: make(map[string][]byte)}\n}\n\nfunc (s *integrationCheckpointStore) Set(_ context.Context, key string, value []byte) error {\n\ts.data[key] = value\n\treturn nil\n}\n\nfunc (s *integrationCheckpointStore) Get(_ context.Context, key string) ([]byte, bool, error) {\n\tv, ok := s.data[key]\n\treturn v, ok, nil\n}\n\ntype defaultPlan struct {\n\tSteps []string `json:\"steps\"`\n}\n\nfunc (p *defaultPlan) FirstStep() string {\n\tif len(p.Steps) == 0 {\n\t\treturn \"\"\n\t}\n\treturn p.Steps[0]\n}\n\nfunc (p *defaultPlan) MarshalJSON() ([]byte, error) {\n\ttype planTyp defaultPlan\n\treturn sonic.Marshal((*planTyp)(p))\n}\n\nfunc (p *defaultPlan) UnmarshalJSON(bytes []byte) error {\n\ttype planTyp defaultPlan\n\treturn sonic.Unmarshal(bytes, (*planTyp)(p))\n}\n\ntype namedAgent struct {\n\tadk.ResumableAgent\n\tname        string\n\tdescription string\n}\n\nfunc (n *namedAgent) Name(_ context.Context) string {\n\treturn n.name\n}\n\nfunc (n *namedAgent) Description(_ context.Context) string {\n\treturn n.description\n}\n\nfunc formatRunPath(runPath []adk.RunStep) string {\n\tif len(runPath) == 0 {\n\t\treturn \"[]\"\n\t}\n\tvar parts []string\n\tfor _, step := range runPath {\n\t\tparts = append(parts, step.String())\n\t}\n\treturn \"[\" + strings.Join(parts, \" -> \") + \"]\"\n}\n\nfunc formatAgentEventIntegration(event *adk.AgentEvent) string {\n\tvar sb strings.Builder\n\tsb.WriteString(fmt.Sprintf(\"{AgentName: %q, RunPath: %s\", event.AgentName, formatRunPath(event.RunPath)))\n\tif event.Output != nil {\n\t\tif event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil {\n\t\t\tmsg := event.Output.MessageOutput.Message\n\t\t\tsb.WriteString(fmt.Sprintf(\", Output.Message: {Role: %q, Content: %q}\", msg.Role, msg.Content))\n\t\t}\n\t}\n\tif event.Action != nil {\n\t\tif event.Action.Interrupted != nil {\n\t\t\tsb.WriteString(fmt.Sprintf(\", Action.Interrupted: {%d contexts}\", len(event.Action.Interrupted.InterruptContexts)))\n\t\t}\n\t\tif event.Action.BreakLoop != nil {\n\t\t\tsb.WriteString(fmt.Sprintf(\", Action.BreakLoop: {From: %q, Done: %v}\", event.Action.BreakLoop.From, event.Action.BreakLoop.Done))\n\t\t}\n\t\tif event.Action.TransferToAgent != nil {\n\t\t\tsb.WriteString(fmt.Sprintf(\", Action.TransferToAgent: {Dest: %q}\", event.Action.TransferToAgent.DestAgentName))\n\t\t}\n\t}\n\tif event.Err != nil {\n\t\tsb.WriteString(fmt.Sprintf(\", Err: %v\", event.Err))\n\t}\n\tsb.WriteString(\"}\")\n\treturn sb.String()\n}\n\nfunc TestSupervisorWithPlanExecuteInterruptResume(t *testing.T) {\n\tctx := context.Background()\n\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tmockSupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockPlannerModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockExecutorModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockReplannerModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tbudgetTool := &approvableTool{name: \"allocate_budget\", t: t}\n\n\tplan := &defaultPlan{Steps: []string{\"Allocate budget for the project\", \"Complete task\"}}\n\tuserInput := []adk.Message{schema.UserMessage(\"Set up a new project with budget allocation\")}\n\n\tplannerModelWithTools := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockPlannerModel.EXPECT().WithTools(gomock.Any()).Return(plannerModelWithTools, nil).AnyTimes()\n\n\tplanJSON, _ := sonic.MarshalString(plan)\n\tplannerResponse := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"plan_call_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"plan\",\n\t\t\t\tArguments: planJSON,\n\t\t\t},\n\t\t},\n\t})\n\tplannerModelWithTools.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input []*schema.Message, opts ...interface{}) (*schema.StreamReader[*schema.Message], error) {\n\t\t\tsr, sw := schema.Pipe[*schema.Message](1)\n\t\t\tgo func() {\n\t\t\t\tdefer sw.Close()\n\t\t\t\tsw.Send(plannerResponse, nil)\n\t\t\t}()\n\t\t\treturn sr, nil\n\t\t},\n\t).Times(1)\n\n\tmockExecutorModel.EXPECT().WithTools(gomock.Any()).Return(mockExecutorModel, nil).AnyTimes()\n\n\ttoolCallMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"call_budget_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"allocate_budget\",\n\t\t\t\tArguments: `{\"action\": \"allocate $50000 for project\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tmockExecutorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(toolCallMsg, nil).Times(1)\n\n\tcompletionMsg := schema.AssistantMessage(\"Budget allocated successfully\", nil)\n\tmockExecutorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(completionMsg, nil).AnyTimes()\n\n\treplannerModelWithTools := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockReplannerModel.EXPECT().WithTools(gomock.Any()).Return(replannerModelWithTools, nil).AnyTimes()\n\n\trespondResponse := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"respond_call_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"respond\",\n\t\t\t\tArguments: `{\"response\":\"Project setup completed with budget allocation\"}`,\n\t\t\t},\n\t\t},\n\t})\n\treplannerModelWithTools.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input []*schema.Message, opts ...interface{}) (*schema.StreamReader[*schema.Message], error) {\n\t\t\tsr, sw := schema.Pipe[*schema.Message](1)\n\t\t\tgo func() {\n\t\t\t\tdefer sw.Close()\n\t\t\t\tsw.Send(respondResponse, nil)\n\t\t\t}()\n\t\t\treturn sr, nil\n\t\t},\n\t).AnyTimes()\n\n\tplannerAgent, err := planexecute.NewPlanner(ctx, &planexecute.PlannerConfig{\n\t\tToolCallingChatModel: mockPlannerModel,\n\t})\n\tassert.NoError(t, err)\n\n\texecutorAgent, err := planexecute.NewExecutor(ctx, &planexecute.ExecutorConfig{\n\t\tModel: mockExecutorModel,\n\t\tToolsConfig: adk.ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{budgetTool},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\treplannerAgent, err := planexecute.NewReplanner(ctx, &planexecute.ReplannerConfig{\n\t\tChatModel: mockReplannerModel,\n\t})\n\tassert.NoError(t, err)\n\n\tplanExecuteAgent, err := planexecute.New(ctx, &planexecute.Config{\n\t\tPlanner:       plannerAgent,\n\t\tExecutor:      executorAgent,\n\t\tReplanner:     replannerAgent,\n\t\tMaxIterations: 10,\n\t})\n\tassert.NoError(t, err)\n\n\tprojectAgent := &namedAgent{\n\t\tResumableAgent: planExecuteAgent,\n\t\tname:           \"project_execution_agent\",\n\t\tdescription:    \"the agent responsible for complex project execution tasks\",\n\t}\n\n\tvar pa adk.Agent\n\tpa = projectAgent\n\n\t_, ok := pa.(adk.ResumableAgent)\n\tassert.True(t, ok)\n\n\tmockSupervisorModel.EXPECT().WithTools(gomock.Any()).Return(mockSupervisorModel, nil).AnyTimes()\n\n\ttransferToProjectMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"transfer_call_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"transfer_to_agent\",\n\t\t\t\tArguments: `{\"agent_name\":\"project_execution_agent\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tmockSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(transferToProjectMsg, nil).Times(1)\n\n\tfinalSupervisorMsg := schema.AssistantMessage(\"Project setup completed successfully with budget allocation approved.\", nil)\n\tmockSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(finalSupervisorMsg, nil).AnyTimes()\n\n\tsupervisorChatAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"project_manager\",\n\t\tDescription: \"the supervisor agent responsible for coordinating project management tasks\",\n\t\tInstruction: \"You are a project manager supervisor. Delegate complex project tasks to project_execution_agent.\",\n\t\tModel:       mockSupervisorModel,\n\t\tExit:        &adk.ExitTool{},\n\t})\n\tassert.NoError(t, err)\n\n\tsupervisorAgent, err := supervisor.New(ctx, &supervisor.Config{\n\t\tSupervisor: supervisorChatAgent,\n\t\tSubAgents:  []adk.Agent{projectAgent},\n\t})\n\tassert.NoError(t, err)\n\n\tstore := newIntegrationCheckpointStore()\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{\n\t\tAgent:           supervisorAgent,\n\t\tCheckPointStore: store,\n\t})\n\n\tt.Log(\"========================================\")\n\tt.Log(\"Starting Supervisor + PlanExecute Integration Test\")\n\tt.Log(\"========================================\")\n\n\tcheckpointID := \"test-supervisor-plan_execute-1\"\n\titer := runner.Run(ctx, userInput, adk.WithCheckPointID(checkpointID))\n\n\tvar interruptEvent *adk.AgentEvent\n\teventCount := 0\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\teventCount++\n\t\tt.Logf(\"Event %d: %s\", eventCount, formatAgentEventIntegration(event))\n\n\t\tif event.Err != nil {\n\t\t\tt.Logf(\"Event has error: %v\", event.Err)\n\t\t}\n\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tinterruptEvent = event\n\t\t\tt.Log(\"========================================\")\n\t\t\tt.Log(\"INTERRUPT DETECTED - Deep interrupt from tool within executor\")\n\t\t\tt.Log(\"========================================\")\n\t\t\tbreak\n\t\t}\n\t}\n\n\tif interruptEvent == nil {\n\t\tt.Fatal(\"Expected an interrupt event from the approvable tool, but none was received\")\n\t}\n\n\tassert.NotNil(t, interruptEvent.Action.Interrupted, \"Should have interrupt info\")\n\tassert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, \"Should have interrupt contexts\")\n\n\tt.Logf(\"Interrupt event received with %d contexts\", len(interruptEvent.Action.Interrupted.InterruptContexts))\n\tfor i, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tt.Logf(\"Interrupt context %d: ID=%s, Info=%v, IsRootCause=%v\", i, ctx.ID, ctx.Info, ctx.IsRootCause)\n\t}\n\n\tvar toolInterruptID string\n\tfor _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif intCtx.IsRootCause {\n\t\t\ttoolInterruptID = intCtx.ID\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.NotEmpty(t, toolInterruptID, \"Should have a root cause interrupt ID\")\n\n\tt.Log(\"========================================\")\n\tt.Logf(\"Resuming with approval for interrupt ID: %s\", toolInterruptID)\n\tt.Log(\"========================================\")\n\n\tresumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\ttoolInterruptID: &approvalResult{Approved: true},\n\t\t},\n\t})\n\tassert.NoError(t, err, \"Resume should not error\")\n\tassert.NotNil(t, resumeIter, \"Resume iterator should not be nil\")\n\n\tvar resumeEvents []*adk.AgentEvent\n\tfor {\n\t\tevent, ok := resumeIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tresumeEvents = append(resumeEvents, event)\n\t}\n\n\tassert.NotEmpty(t, resumeEvents, \"Should have resume events after approval\")\n\n\tfor _, event := range resumeEvents {\n\t\tassert.NoError(t, event.Err, \"Resume event should not have error\")\n\t}\n\n\tvar hasToolResponse, hasBreakLoop bool\n\tfor _, event := range resumeEvents {\n\t\tif event.Output != nil && event.Output.MessageOutput != nil {\n\t\t\tmsg := event.Output.MessageOutput.Message\n\t\t\tif msg != nil && msg.Role == \"tool\" && strings.Contains(msg.Content, \"executed successfully\") {\n\t\t\t\thasToolResponse = true\n\t\t\t}\n\t\t}\n\t\tif event.Action != nil && event.Action.BreakLoop != nil && event.Action.BreakLoop.Done {\n\t\t\thasBreakLoop = true\n\t\t}\n\t}\n\n\tassert.True(t, hasToolResponse, \"Should have tool response indicating successful execution\")\n\tassert.True(t, hasBreakLoop, \"Should have break loop action indicating task completion\")\n}\n"
  },
  {
    "path": "adk/prebuilt/planexecute/plan_execute.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package planexecute implements a plan–execute–replan style agent.\npackage planexecute\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"runtime/debug\"\n\t\"strings\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc init() {\n\tschema.RegisterName[*defaultPlan](\"_eino_adk_plan_execute_default_plan\")\n\tschema.RegisterName[ExecutedStep](\"_eino_adk_plan_execute_executed_step\")\n\tschema.RegisterName[[]ExecutedStep](\"_eino_adk_plan_execute_executed_steps\")\n}\n\n// Plan represents an execution plan with a sequence of actionable steps.\n// It supports JSON serialization and deserialization while providing access to the first step.\ntype Plan interface {\n\t// FirstStep returns the first step to be executed in the plan.\n\tFirstStep() string\n\n\t// Marshaler serializes the Plan into JSON.\n\t// The resulting JSON can be used in prompt templates.\n\tjson.Marshaler\n\t// Unmarshaler deserializes JSON content into the Plan.\n\t// This processes output from structured chat models or tool calls into the Plan structure.\n\tjson.Unmarshaler\n}\n\n// NewPlan is a function type that creates a new Plan instance.\ntype NewPlan func(ctx context.Context) Plan\n\n// defaultPlan is the default implementation of the Plan interface.\n//\n// JSON Schema:\n//\n//\t{\n//\t  \"type\": \"object\",\n//\t  \"properties\": {\n//\t    \"steps\": {\n//\t      \"type\": \"array\",\n//\t      \"items\": {\n//\t        \"type\": \"string\"\n//\t      },\n//\t      \"description\": \"Ordered list of actions to be taken. Each step should be clear, actionable, and arranged in a logical sequence.\"\n//\t    }\n//\t  },\n//\t  \"required\": [\"steps\"]\n//\t}\ntype defaultPlan struct {\n\t// Steps contains the ordered list of actions to be taken.\n\t// Each step should be clear, actionable, and arranged in a logical sequence.\n\tSteps []string `json:\"steps\"`\n}\n\n// FirstStep returns the first step in the plan or an empty string if no steps exist.\nfunc (p *defaultPlan) FirstStep() string {\n\tif len(p.Steps) == 0 {\n\t\treturn \"\"\n\t}\n\treturn p.Steps[0]\n}\n\nfunc (p *defaultPlan) MarshalJSON() ([]byte, error) {\n\ttype planTyp defaultPlan\n\treturn sonic.Marshal((*planTyp)(p))\n}\n\nfunc (p *defaultPlan) UnmarshalJSON(bytes []byte) error {\n\ttype planTyp defaultPlan\n\treturn sonic.Unmarshal(bytes, (*planTyp)(p))\n}\n\n// Response represents the final response to the user.\n// This struct is used for JSON serialization/deserialization of the final response\n// generated by the model.\ntype Response struct {\n\t// Response is the complete response to provide to the user.\n\t// This field is required.\n\tResponse string `json:\"response\"`\n}\n\nvar (\n\t// PlanToolInfo defines the schema for the Plan tool that can be used with ToolCallingChatModel.\n\t// This schema instructs the model to generate a structured plan with ordered steps.\n\tPlanToolInfo = schema.ToolInfo{\n\t\tName: \"plan\",\n\t\tDesc: \"Plan with a list of steps to execute in order. Each step should be clear, actionable, and arranged in a logical sequence. The output will be used to guide the execution process.\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"steps\": {\n\t\t\t\t\tType:     schema.Array,\n\t\t\t\t\tElemInfo: &schema.ParameterInfo{Type: schema.String},\n\t\t\t\t\tDesc:     \"different steps to follow, should be in sorted order\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t},\n\t\t\t},\n\t\t),\n\t}\n\n\t// RespondToolInfo defines the schema for the response tool that can be used with ToolCallingChatModel.\n\t// This schema instructs the model to generate a direct response to the user.\n\tRespondToolInfo = schema.ToolInfo{\n\t\tName: \"respond\",\n\t\tDesc: \"Generate a direct response to the user. Use this tool when you have all the information needed to provide a final answer.\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"response\": {\n\t\t\t\t\tType:     schema.String,\n\t\t\t\t\tDesc:     \"The complete response to provide to the user\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t},\n\t\t\t},\n\t\t),\n\t}\n\n\t// PlannerPrompt is the prompt template for the planner.\n\t// It provides context and guidance to the planner on how to generate the Plan.\n\tPlannerPrompt = prompt.FromMessages(schema.FString,\n\t\tschema.SystemMessage(`You are an expert planning agent. Given an objective, create a comprehensive step-by-step plan to achieve the objective.\n\n## YOUR TASK\nAnalyze the objective and generate a strategic plan that breaks down the goal into manageable, executable steps.\n\n## PLANNING REQUIREMENTS\nEach step in your plan must be:\n- **Specific and actionable**: Clear instructions that can be executed without ambiguity\n- **Self-contained**: Include all necessary context, parameters, and requirements\n- **Independently executable**: Can be performed by an agent without dependencies on other steps\n- **Logically sequenced**: Arranged in optimal order for efficient execution\n- **Objective-focused**: Directly contribute to achieving the main goal\n\n## PLANNING GUIDELINES\n- Eliminate redundant or unnecessary steps\n- Include relevant constraints, parameters, and success criteria for each step\n- Ensure the final step produces a complete answer or deliverable\n- Anticipate potential challenges and include mitigation strategies\n- Structure steps to build upon each other logically\n- Provide sufficient detail for successful execution\n\n## QUALITY CRITERIA\n- Plan completeness: Does it address all aspects of the objective?\n- Step clarity: Can each step be understood and executed independently?\n- Logical flow: Do steps follow a sensible progression?\n- Efficiency: Is this the most direct path to the objective?\n- Adaptability: Can the plan handle unexpected results or changes?`),\n\t\tschema.MessagesPlaceholder(\"input\", false),\n\t)\n\n\t// ExecutorPrompt is the prompt template for the executor.\n\t// It provides context and guidance to the executor on how to execute the Task.\n\tExecutorPrompt = prompt.FromMessages(schema.FString,\n\t\tschema.SystemMessage(`You are a diligent and meticulous executor agent. Follow the given plan and execute your tasks carefully and thoroughly.`),\n\t\tschema.UserMessage(`## OBJECTIVE\n{input}\n## Given the following plan:\n{plan}\n## COMPLETED STEPS & RESULTS\n{executed_steps}\n## Your task is to execute the first step, which is: \n{step}`))\n\n\t// ReplannerPrompt is the prompt template for the replanner.\n\t// It provides context and guidance to the replanner on how to regenerate the Plan.\n\tReplannerPrompt = prompt.FromMessages(schema.FString,\n\t\tschema.SystemMessage(\n\t\t\t`You are going to review the progress toward an objective. Analyze the current state and determine the optimal next action.\n\n## YOUR TASK\nBased on the progress above, you MUST choose exactly ONE action:\n\n### Option 1: COMPLETE (if objective is fully achieved)\nCall '{respond_tool}' with:\n- A comprehensive final answer\n- Clear conclusion summarizing how the objective was met\n- Key insights from the execution process\n\n### Option 2: CONTINUE (if more work is needed)\nCall '{plan_tool}' with a revised plan that:\n- Contains ONLY remaining steps (exclude completed ones)\n- Incorporates lessons learned from executed steps\n- Addresses any gaps or issues discovered\n- Maintains logical step sequence\n\n## PLANNING REQUIREMENTS\nEach step in your plan must be:\n- **Specific and actionable**: Clear instructions that can be executed without ambiguity\n- **Self-contained**: Include all necessary context, parameters, and requirements\n- **Independently executable**: Can be performed by an agent without dependencies on other steps\n- **Logically sequenced**: Arranged in optimal order for efficient execution\n- **Objective-focused**: Directly contribute to achieving the main goal\n\n## PLANNING GUIDELINES\n- Eliminate redundant or unnecessary steps\n- Adapt strategy based on new information\n- Include relevant constraints, parameters, and success criteria for each step\n\n## DECISION CRITERIA\n- Has the original objective been completely satisfied?\n- Are there any remaining requirements or sub-goals?\n- Do the results suggest a need for strategy adjustment?\n- What specific actions are still required?`),\n\t\tschema.UserMessage(`## OBJECTIVE\n{input}\n\n## ORIGINAL PLAN\n{plan}\n\n## COMPLETED STEPS & RESULTS\n{executed_steps}`),\n\t)\n)\n\nconst (\n\t// UserInputSessionKey is the session key for the user input.\n\tUserInputSessionKey = \"UserInput\"\n\n\t// PlanSessionKey is the session key for the plan.\n\tPlanSessionKey = \"Plan\"\n\n\t// ExecutedStepSessionKey is the session key for the execute result.\n\tExecutedStepSessionKey = \"ExecutedStep\"\n\n\t// ExecutedStepsSessionKey is the session key for the execute results.\n\tExecutedStepsSessionKey = \"ExecutedSteps\"\n)\n\n// PlannerConfig provides configuration options for creating a planner agent.\n// There are two ways to configure the planner to generate structured Plan output:\n//  1. Use ChatModelWithFormattedOutput: A model pre-configured to output in the Plan format\n//  2. Use ToolCallingChatModel + ToolInfo: A model that uses tool calling to generate\n//     the Plan structure\ntype PlannerConfig struct {\n\t// ChatModelWithFormattedOutput is a model pre-configured to output in the Plan format.\n\t// Create this by configuring a model to output structured data directly.\n\t// See example: https://github.com/cloudwego/eino-ext/blob/main/components/model/openai/examples/structured/structured.go\n\tChatModelWithFormattedOutput model.BaseChatModel\n\n\t// ToolCallingChatModel is a model that supports tool calling capabilities.\n\t// When provided with ToolInfo, it will use tool calling to generate the Plan structure.\n\tToolCallingChatModel model.ToolCallingChatModel\n\n\t// ToolInfo defines the schema for the Plan structure when using tool calling.\n\t// Optional. If not provided, PlanToolInfo will be used as the default.\n\tToolInfo *schema.ToolInfo\n\n\t// GenInputFn is a function that generates the input messages for the planner.\n\t// Optional. If not provided, defaultGenPlannerInputFn will be used.\n\tGenInputFn GenPlannerModelInputFn\n\n\t// NewPlan creates a new Plan instance for JSON.\n\t// The returned Plan will be used to unmarshal the model-generated JSON output.\n\t// Optional. If not provided, defaultNewPlan will be used.\n\tNewPlan NewPlan\n}\n\n// GenPlannerModelInputFn is a function type that generates input messages for the planner.\ntype GenPlannerModelInputFn func(ctx context.Context, userInput []adk.Message) ([]adk.Message, error)\n\nfunc defaultNewPlan(ctx context.Context) Plan {\n\treturn &defaultPlan{}\n}\n\nfunc defaultGenPlannerInputFn(ctx context.Context, userInput []adk.Message) ([]adk.Message, error) {\n\tmsgs, err := PlannerPrompt.Format(ctx, map[string]any{\n\t\t\"input\": userInput,\n\t})\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn msgs, nil\n}\n\ntype planner struct {\n\ttoolCall   bool\n\tchatModel  model.BaseChatModel\n\tgenInputFn GenPlannerModelInputFn\n\tnewPlan    NewPlan\n}\n\nfunc (p *planner) Name(_ context.Context) string {\n\treturn \"planner\"\n}\n\nfunc (p *planner) Description(_ context.Context) string {\n\treturn \"a planner agent\"\n}\n\nfunc argToContent(msg adk.Message) (adk.Message, error) {\n\tif len(msg.ToolCalls) == 0 {\n\t\treturn nil, schema.ErrNoValue\n\t}\n\n\treturn schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil), nil\n}\n\nfunc (p *planner) Run(ctx context.Context, input *adk.AgentInput,\n\t_ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\n\titerator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\n\tadk.AddSessionValue(ctx, UserInputSessionKey, input.Messages)\n\n\tgo func() {\n\t\tdefer func() {\n\t\t\tpanicErr := recover()\n\t\t\tif panicErr != nil {\n\t\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\t\tgenerator.Send(&adk.AgentEvent{Err: e})\n\t\t\t}\n\n\t\t\tgenerator.Close()\n\t\t}()\n\n\t\tc := compose.NewChain[*adk.AgentInput, Plan]().\n\t\t\tAppendLambda(\n\t\t\t\tcompose.InvokableLambda(func(ctx context.Context, input *adk.AgentInput) (output []adk.Message, err error) {\n\t\t\t\t\treturn p.genInputFn(ctx, input.Messages)\n\t\t\t\t}),\n\t\t\t).\n\t\t\tAppendChatModel(p.chatModel).\n\t\t\tAppendLambda(\n\t\t\t\tcompose.CollectableLambda(func(ctx context.Context, sr *schema.StreamReader[adk.Message]) (adk.Message, error) {\n\t\t\t\t\tif input.EnableStreaming {\n\t\t\t\t\t\tss := sr.Copy(2)\n\t\t\t\t\t\tvar sOutput *schema.StreamReader[*schema.Message]\n\t\t\t\t\t\tif p.toolCall {\n\t\t\t\t\t\t\tsOutput = schema.StreamReaderWithConvert(ss[0], argToContent)\n\t\t\t\t\t\t} else {\n\t\t\t\t\t\t\tsOutput = ss[0]\n\t\t\t\t\t\t}\n\n\t\t\t\t\t\tgenerator.Send(adk.EventFromMessage(nil, sOutput, schema.Assistant, \"\"))\n\n\t\t\t\t\t\treturn schema.ConcatMessageStream(ss[1])\n\t\t\t\t\t}\n\n\t\t\t\t\tmsg, err := schema.ConcatMessageStream(sr)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t}\n\n\t\t\t\t\tvar output adk.Message\n\t\t\t\t\tif p.toolCall {\n\t\t\t\t\t\tif len(msg.ToolCalls) == 0 {\n\t\t\t\t\t\t\treturn nil, fmt.Errorf(\"no tool call\")\n\t\t\t\t\t\t}\n\t\t\t\t\t\toutput = schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil)\n\t\t\t\t\t} else {\n\t\t\t\t\t\toutput = msg\n\t\t\t\t\t}\n\n\t\t\t\t\tgenerator.Send(adk.EventFromMessage(output, nil, schema.Assistant, \"\"))\n\n\t\t\t\t\treturn msg, nil\n\t\t\t\t}),\n\t\t\t).\n\t\t\tAppendLambda(\n\t\t\t\tcompose.InvokableLambda(func(ctx context.Context, msg adk.Message) (plan Plan, err error) {\n\t\t\t\t\tvar planJSON string\n\t\t\t\t\tif p.toolCall {\n\t\t\t\t\t\tif len(msg.ToolCalls) == 0 {\n\t\t\t\t\t\t\treturn nil, fmt.Errorf(\"no tool call\")\n\t\t\t\t\t\t}\n\t\t\t\t\t\tplanJSON = msg.ToolCalls[0].Function.Arguments\n\t\t\t\t\t} else {\n\t\t\t\t\t\tplanJSON = msg.Content\n\t\t\t\t\t}\n\n\t\t\t\t\tplan = p.newPlan(ctx)\n\t\t\t\t\terr = plan.UnmarshalJSON([]byte(planJSON))\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn nil, fmt.Errorf(\"unmarshal plan error: %w\", err)\n\t\t\t\t\t}\n\n\t\t\t\t\tadk.AddSessionValue(ctx, PlanSessionKey, plan)\n\n\t\t\t\t\treturn plan, nil\n\t\t\t\t}),\n\t\t\t)\n\n\t\tvar opts []compose.Option\n\t\tif p.toolCall {\n\t\t\topts = append(opts, compose.WithChatModelOption(model.WithToolChoice(schema.ToolChoiceForced)))\n\t\t}\n\n\t\tr, err := c.Compile(ctx, compose.WithGraphName(p.Name(ctx)))\n\t\tif err != nil { // unexpected\n\t\t\tgenerator.Send(&adk.AgentEvent{Err: err})\n\t\t\treturn\n\t\t}\n\n\t\t_, err = r.Stream(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\tgenerator.Send(&adk.AgentEvent{Err: err})\n\t\t\treturn\n\t\t}\n\t}()\n\n\treturn iterator\n}\n\n// NewPlanner creates a new planner agent based on the provided configuration.\n// The planner agent uses either ChatModelWithFormattedOutput or ToolCallingChatModel+ToolInfo\n// to generate structured Plan output.\n//\n// If ChatModelWithFormattedOutput is provided, it will be used directly.\n// If ToolCallingChatModel is provided, it will be configured with ToolInfo (or PlanToolInfo by default)\n// to generate structured Plan output.\nfunc NewPlanner(_ context.Context, cfg *PlannerConfig) (adk.Agent, error) {\n\tvar chatModel model.BaseChatModel\n\tvar toolCall bool\n\tif cfg.ChatModelWithFormattedOutput != nil {\n\t\tchatModel = cfg.ChatModelWithFormattedOutput\n\t} else {\n\t\ttoolCall = true\n\t\ttoolInfo := cfg.ToolInfo\n\t\tif toolInfo == nil {\n\t\t\ttoolInfo = &PlanToolInfo\n\t\t}\n\n\t\tvar err error\n\t\tchatModel, err = cfg.ToolCallingChatModel.WithTools([]*schema.ToolInfo{toolInfo})\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tinputFn := cfg.GenInputFn\n\tif inputFn == nil {\n\t\tinputFn = defaultGenPlannerInputFn\n\t}\n\n\tplanParser := cfg.NewPlan\n\tif planParser == nil {\n\t\tplanParser = defaultNewPlan\n\t}\n\n\treturn &planner{\n\t\ttoolCall:   toolCall,\n\t\tchatModel:  chatModel,\n\t\tgenInputFn: inputFn,\n\t\tnewPlan:    planParser,\n\t}, nil\n}\n\n// ExecutionContext is the input information for the executor and the planner.\ntype ExecutionContext struct {\n\tUserInput     []adk.Message\n\tPlan          Plan\n\tExecutedSteps []ExecutedStep\n}\n\n// GenModelInputFn is a function that generates the input messages for the executor and the planner.\ntype GenModelInputFn func(ctx context.Context, in *ExecutionContext) ([]adk.Message, error)\n\n// ExecutorConfig provides configuration options for creating an executor agent.\ntype ExecutorConfig struct {\n\t// Model is the chat model used by the executor.\n\t// If the executor uses any tools, this model must support the model.WithTools call option,\n\t// as that's how the executor configures the model with tool information.\n\tModel model.BaseChatModel\n\n\t// ToolsConfig specifies the tools available to the executor.\n\tToolsConfig adk.ToolsConfig\n\n\t// MaxIterations defines the upper limit of ChatModel generation cycles.\n\t// The agent will terminate with an error if this limit is exceeded.\n\t// Optional. Defaults to 20.\n\tMaxIterations int\n\n\t// GenInputFn generates the input messages for the Executor.\n\t// Optional. If not provided, defaultGenExecutorInputFn will be used.\n\tGenInputFn GenModelInputFn\n}\n\ntype ExecutedStep struct {\n\tStep   string\n\tResult string\n}\n\n// NewExecutor creates a new executor agent.\nfunc NewExecutor(ctx context.Context, cfg *ExecutorConfig) (adk.Agent, error) {\n\n\tgenInputFn := cfg.GenInputFn\n\tif genInputFn == nil {\n\t\tgenInputFn = defaultGenExecutorInputFn\n\t}\n\tgenInput := func(ctx context.Context, instruction string, _ *adk.AgentInput) ([]adk.Message, error) {\n\n\t\tplan, ok := adk.GetSessionValue(ctx, PlanSessionKey)\n\t\tif !ok {\n\t\t\tpanic(\"impossible: plan not found\")\n\t\t}\n\t\tplan_ := plan.(Plan)\n\n\t\tuserInput, ok := adk.GetSessionValue(ctx, UserInputSessionKey)\n\t\tif !ok {\n\t\t\tpanic(\"impossible: user input not found\")\n\t\t}\n\t\tuserInput_ := userInput.([]adk.Message)\n\n\t\tvar executedSteps_ []ExecutedStep\n\t\texecutedStep, ok := adk.GetSessionValue(ctx, ExecutedStepsSessionKey)\n\t\tif ok {\n\t\t\texecutedSteps_ = executedStep.([]ExecutedStep)\n\t\t}\n\n\t\tin := &ExecutionContext{\n\t\t\tUserInput:     userInput_,\n\t\t\tPlan:          plan_,\n\t\t\tExecutedSteps: executedSteps_,\n\t\t}\n\n\t\tmsgs, err := genInputFn(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn msgs, nil\n\t}\n\n\tagent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:          \"executor\",\n\t\tDescription:   \"an executor agent\",\n\t\tModel:         cfg.Model,\n\t\tToolsConfig:   cfg.ToolsConfig,\n\t\tGenModelInput: genInput,\n\t\tMaxIterations: cfg.MaxIterations,\n\t\tOutputKey:     ExecutedStepSessionKey,\n\t})\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn agent, nil\n}\n\nfunc defaultGenExecutorInputFn(ctx context.Context, in *ExecutionContext) ([]adk.Message, error) {\n\n\tplanContent, err := in.Plan.MarshalJSON()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tuserMsgs, err := ExecutorPrompt.Format(ctx, map[string]any{\n\t\t\"input\":          formatInput(in.UserInput),\n\t\t\"plan\":           string(planContent),\n\t\t\"executed_steps\": formatExecutedSteps(in.ExecutedSteps),\n\t\t\"step\":           in.Plan.FirstStep(),\n\t})\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn userMsgs, nil\n}\n\ntype replanner struct {\n\tchatModel   model.ToolCallingChatModel\n\tplanTool    *schema.ToolInfo\n\trespondTool *schema.ToolInfo\n\n\tgenInputFn GenModelInputFn\n\tnewPlan    NewPlan\n}\n\ntype ReplannerConfig struct {\n\t// ChatModel is the model that supports tool calling capabilities.\n\t// It will be configured with PlanTool and RespondTool to generate updated plans or responses.\n\tChatModel model.ToolCallingChatModel\n\n\t// PlanTool defines the schema for the Plan tool that can be used with ToolCallingChatModel.\n\t// Optional. If not provided, the default PlanToolInfo will be used.\n\tPlanTool *schema.ToolInfo\n\n\t// RespondTool defines the schema for the response tool that can be used with ToolCallingChatModel.\n\t// Optional. If not provided, the default RespondToolInfo will be used.\n\tRespondTool *schema.ToolInfo\n\n\t// GenInputFn generates the input messages for the Replanner.\n\t// Optional. If not provided, buildGenReplannerInputFn will be used.\n\tGenInputFn GenModelInputFn\n\n\t// NewPlan creates a new Plan instance.\n\t// The returned Plan will be used to unmarshal the model-generated JSON output from PlanTool.\n\t// Optional. If not provided, defaultNewPlan will be used.\n\tNewPlan NewPlan\n}\n\n// formatInput formats the input messages into a string.\nfunc formatInput(input []adk.Message) string {\n\tvar sb strings.Builder\n\tfor _, msg := range input {\n\t\tsb.WriteString(msg.Content)\n\t\tsb.WriteString(\"\\n\")\n\t}\n\n\treturn sb.String()\n}\n\nfunc formatExecutedSteps(results []ExecutedStep) string {\n\tvar sb strings.Builder\n\tfor _, result := range results {\n\t\tsb.WriteString(fmt.Sprintf(\"Step: %s\\nResult: %s\\n\\n\", result.Step, result.Result))\n\t}\n\n\treturn sb.String()\n}\n\nfunc (r *replanner) Name(_ context.Context) string {\n\treturn \"replanner\"\n}\n\nfunc (r *replanner) Description(_ context.Context) string {\n\treturn \"a replanner agent\"\n}\n\nfunc (r *replanner) genInput(ctx context.Context) ([]adk.Message, error) {\n\n\texecutedStep, ok := adk.GetSessionValue(ctx, ExecutedStepSessionKey)\n\tif !ok {\n\t\tpanic(\"impossible: execute result not found\")\n\t}\n\texecutedStep_ := executedStep.(string)\n\n\tplan, ok := adk.GetSessionValue(ctx, PlanSessionKey)\n\tif !ok {\n\t\tpanic(\"impossible: plan not found\")\n\t}\n\tplan_ := plan.(Plan)\n\tstep := plan_.FirstStep()\n\n\tvar executedSteps_ []ExecutedStep\n\texecutedSteps, ok := adk.GetSessionValue(ctx, ExecutedStepsSessionKey)\n\tif ok {\n\t\texecutedSteps_ = executedSteps.([]ExecutedStep)\n\t}\n\n\texecutedSteps_ = append(executedSteps_, ExecutedStep{\n\t\tStep:   step,\n\t\tResult: executedStep_,\n\t})\n\tadk.AddSessionValue(ctx, ExecutedStepsSessionKey, executedSteps_)\n\n\tuserInput, ok := adk.GetSessionValue(ctx, UserInputSessionKey)\n\tif !ok {\n\t\tpanic(\"impossible: user input not found\")\n\t}\n\tuserInput_ := userInput.([]adk.Message)\n\n\tin := &ExecutionContext{\n\t\tUserInput:     userInput_,\n\t\tPlan:          plan_,\n\t\tExecutedSteps: executedSteps_,\n\t}\n\tgenInputFn := r.genInputFn\n\tif genInputFn == nil {\n\t\tgenInputFn = buildGenReplannerInputFn(r.planTool.Name, r.respondTool.Name)\n\t}\n\tmsgs, err := genInputFn(ctx, in)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn msgs, nil\n}\n\nfunc (r *replanner) Run(ctx context.Context, input *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\titerator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\n\tgo func() {\n\t\tdefer func() {\n\t\t\tpanicErr := recover()\n\t\t\tif panicErr != nil {\n\t\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\t\tgenerator.Send(&adk.AgentEvent{Err: e})\n\t\t\t}\n\n\t\t\tgenerator.Close()\n\t\t}()\n\n\t\tcallOpt := model.WithToolChoice(schema.ToolChoiceForced)\n\n\t\tc := compose.NewChain[struct{}, any]().\n\t\t\tAppendLambda(\n\t\t\t\tcompose.InvokableLambda(func(ctx context.Context, input struct{}) (output []adk.Message, err error) {\n\t\t\t\t\treturn r.genInput(ctx)\n\t\t\t\t}),\n\t\t\t).\n\t\t\tAppendChatModel(r.chatModel).\n\t\t\tAppendLambda(\n\t\t\t\tcompose.CollectableLambda(func(ctx context.Context, sr *schema.StreamReader[adk.Message]) (adk.Message, error) {\n\t\t\t\t\tif input.EnableStreaming {\n\t\t\t\t\t\tss := sr.Copy(2)\n\t\t\t\t\t\tsOutput := schema.StreamReaderWithConvert(ss[0], argToContent)\n\t\t\t\t\t\tgenerator.Send(adk.EventFromMessage(nil, sOutput, schema.Assistant, \"\"))\n\t\t\t\t\t\treturn schema.ConcatMessageStream(ss[1])\n\t\t\t\t\t}\n\n\t\t\t\t\tmsg, err := schema.ConcatMessageStream(sr)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t}\n\t\t\t\t\tif len(msg.ToolCalls) > 0 {\n\t\t\t\t\t\toutput := schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil)\n\t\t\t\t\t\tgenerator.Send(adk.EventFromMessage(output, nil, schema.Assistant, \"\"))\n\t\t\t\t\t}\n\t\t\t\t\treturn msg, nil\n\t\t\t\t}),\n\t\t\t).\n\t\t\tAppendLambda(\n\t\t\t\tcompose.InvokableLambda(func(ctx context.Context, msg adk.Message) (msgOrPlan any, err error) {\n\t\t\t\t\tif len(msg.ToolCalls) == 0 {\n\t\t\t\t\t\treturn nil, fmt.Errorf(\"no tool call\")\n\t\t\t\t\t}\n\n\t\t\t\t\t// exit\n\t\t\t\t\tif msg.ToolCalls[0].Function.Name == r.respondTool.Name {\n\t\t\t\t\t\taction := adk.NewBreakLoopAction(r.Name(ctx))\n\t\t\t\t\t\tgenerator.Send(&adk.AgentEvent{Action: action})\n\t\t\t\t\t\treturn msg, nil\n\t\t\t\t\t}\n\n\t\t\t\t\t// replan\n\t\t\t\t\tif msg.ToolCalls[0].Function.Name != r.planTool.Name {\n\t\t\t\t\t\treturn nil, fmt.Errorf(\"unexpected tool call: %s\", msg.ToolCalls[0].Function.Name)\n\t\t\t\t\t}\n\n\t\t\t\t\tplan := r.newPlan(ctx)\n\t\t\t\t\tif err = plan.UnmarshalJSON([]byte(msg.ToolCalls[0].Function.Arguments)); err != nil {\n\t\t\t\t\t\treturn nil, fmt.Errorf(\"unmarshal plan error: %w\", err)\n\t\t\t\t\t}\n\n\t\t\t\t\tadk.AddSessionValue(ctx, PlanSessionKey, plan)\n\n\t\t\t\t\treturn plan, nil\n\t\t\t\t}),\n\t\t\t)\n\n\t\trunnable, err := c.Compile(ctx, compose.WithGraphName(r.Name(ctx)))\n\t\tif err != nil {\n\t\t\tgenerator.Send(&adk.AgentEvent{Err: err})\n\t\t\treturn\n\t\t}\n\n\t\t_, err = runnable.Stream(ctx, struct{}{}, compose.WithChatModelOption(callOpt))\n\t\tif err != nil {\n\t\t\tgenerator.Send(&adk.AgentEvent{Err: err})\n\t\t\treturn\n\t\t}\n\t}()\n\n\treturn iterator\n}\n\nfunc buildGenReplannerInputFn(planToolName, respondToolName string) GenModelInputFn {\n\treturn func(ctx context.Context, in *ExecutionContext) ([]adk.Message, error) {\n\t\tplanContent, err := in.Plan.MarshalJSON()\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tmsgs, err := ReplannerPrompt.Format(ctx, map[string]any{\n\t\t\t\"plan\":           string(planContent),\n\t\t\t\"input\":          formatInput(in.UserInput),\n\t\t\t\"executed_steps\": formatExecutedSteps(in.ExecutedSteps),\n\t\t\t\"plan_tool\":      planToolName,\n\t\t\t\"respond_tool\":   respondToolName,\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn msgs, nil\n\t}\n}\n\n// NewReplanner creates a plan-execute-replan agent wired with plan and respond tools.\n// It configures the provided ToolCallingChatModel with the tools and returns an Agent.\nfunc NewReplanner(_ context.Context, cfg *ReplannerConfig) (adk.Agent, error) {\n\tplanTool := cfg.PlanTool\n\tif planTool == nil {\n\t\tplanTool = &PlanToolInfo\n\t}\n\n\trespondTool := cfg.RespondTool\n\tif respondTool == nil {\n\t\trespondTool = &RespondToolInfo\n\t}\n\n\tchatModel, err := cfg.ChatModel.WithTools([]*schema.ToolInfo{planTool, respondTool})\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tplanParser := cfg.NewPlan\n\tif planParser == nil {\n\t\tplanParser = defaultNewPlan\n\t}\n\n\treturn &replanner{\n\t\tchatModel:   chatModel,\n\t\tplanTool:    planTool,\n\t\trespondTool: respondTool,\n\t\tgenInputFn:  cfg.GenInputFn,\n\t\tnewPlan:     planParser,\n\t}, nil\n}\n\n// Config provides configuration options for creating a plan-execute-replan agent.\ntype Config struct {\n\t// Planner specifies the agent that generates the plan.\n\t// You can use provided NewPlanner to create a planner agent.\n\tPlanner adk.Agent\n\n\t// Executor specifies the agent that executes the plan generated by planner or replanner.\n\t// You can use provided NewExecutor to create an executor agent.\n\tExecutor adk.Agent\n\n\t// Replanner specifies the agent that replans the plan.\n\t// You can use provided NewReplanner to create a replanner agent.\n\tReplanner adk.Agent\n\n\t// MaxIterations defines the maximum number of loops for 'execute-replan'.\n\t// Optional. If not provided, 10 will be used as the default.\n\tMaxIterations int\n}\n\n// New creates a new plan-execute-replan agent with the given configuration.\n// The plan-execute-replan pattern works in three phases:\n// 1. Planning: Generate a structured plan with clear, actionable steps\n// 2. Execution: Execute the first step of the plan\n// 3. Replanning: Evaluate progress and either complete the task or revise the plan\n// This approach enables complex problem-solving through iterative refinement.\nfunc New(ctx context.Context, cfg *Config) (adk.ResumableAgent, error) {\n\tmaxIterations := cfg.MaxIterations\n\tif maxIterations <= 0 {\n\t\tmaxIterations = 10\n\t}\n\tloop, err := adk.NewLoopAgent(ctx, &adk.LoopAgentConfig{\n\t\tName:          \"execute_replan\",\n\t\tSubAgents:     []adk.Agent{cfg.Executor, cfg.Replanner},\n\t\tMaxIterations: maxIterations,\n\t})\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn adk.NewSequentialAgent(ctx, &adk.SequentialAgentConfig{\n\t\tName:      \"plan_execute_replan\",\n\t\tSubAgents: []adk.Agent{cfg.Planner, loop},\n\t})\n}\n"
  },
  {
    "path": "adk/prebuilt/planexecute/plan_execute_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage planexecute\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockAdk \"github.com/cloudwego/eino/internal/mock/adk\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// TestNewPlanner tests the NewPlanner function with ChatModelWithFormattedOutput\nfunc TestNewPlannerWithFormattedOutput(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock chat model\n\tmockChatModel := mockModel.NewMockBaseChatModel(ctrl)\n\n\t// Create the PlannerConfig\n\tconf := &PlannerConfig{\n\t\tChatModelWithFormattedOutput: mockChatModel,\n\t}\n\n\t// Create the planner\n\tp, err := NewPlanner(ctx, conf)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, p)\n\n\t// Verify the planner's name and description\n\tassert.Equal(t, \"planner\", p.Name(ctx))\n\tassert.Equal(t, \"a planner agent\", p.Description(ctx))\n}\n\n// TestNewPlannerWithToolCalling tests the NewPlanner function with ToolCallingChatModel\nfunc TestNewPlannerWithToolCalling(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock tool calling chat model\n\tmockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1)\n\n\t// Create the PlannerConfig\n\tconf := &PlannerConfig{\n\t\tToolCallingChatModel: mockToolCallingModel,\n\t\t// Use default instruction and tool info\n\t}\n\n\t// Create the planner\n\tp, err := NewPlanner(ctx, conf)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, p)\n\n\t// Verify the planner's name and description\n\tassert.Equal(t, \"planner\", p.Name(ctx))\n\tassert.Equal(t, \"a planner agent\", p.Description(ctx))\n}\n\n// TestPlannerRunWithFormattedOutput tests the Run method of a planner created with ChatModelWithFormattedOutput\nfunc TestPlannerRunWithFormattedOutput(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock chat model\n\tmockChatModel := mockModel.NewMockBaseChatModel(ctrl)\n\n\t// Create a plan response\n\tplanJSON := `{\"steps\":[\"Step 1\", \"Step 2\", \"Step 3\"]}`\n\tplanMsg := schema.AssistantMessage(planJSON, nil)\n\tsr, sw := schema.Pipe[*schema.Message](1)\n\tsw.Send(planMsg, nil)\n\tsw.Close()\n\n\t// Mock the Generate method\n\tmockChatModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\n\t// Create the PlannerConfig\n\tconf := &PlannerConfig{\n\t\tChatModelWithFormattedOutput: mockChatModel,\n\t}\n\n\t// Create the planner\n\tp, err := NewPlanner(ctx, conf)\n\tassert.NoError(t, err)\n\n\t// Run the planner\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: p})\n\titerator := runner.Run(ctx, []adk.Message{schema.UserMessage(\"Plan this task\")})\n\n\t// Get the event from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Nil(t, event.Err)\n\tmsg, _, err := adk.GetMessage(event)\n\tassert.NoError(t, err)\n\tassert.Equal(t, planMsg.Content, msg.Content)\n\n\tevent, ok = iterator.Next()\n\tassert.False(t, ok)\n\n\tplan := defaultNewPlan(ctx)\n\terr = plan.UnmarshalJSON([]byte(msg.Content))\n\tassert.NoError(t, err)\n\tplan_ := plan.(*defaultPlan)\n\tassert.Equal(t, 3, len(plan_.Steps))\n\tassert.Equal(t, \"Step 1\", plan_.Steps[0])\n\tassert.Equal(t, \"Step 2\", plan_.Steps[1])\n\tassert.Equal(t, \"Step 3\", plan_.Steps[2])\n}\n\n// TestPlannerRunWithToolCalling tests the Run method of a planner created with ToolCallingChatModel\nfunc TestPlannerRunWithToolCalling(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock tool calling chat model\n\tmockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// Create a tool call response with a plan\n\tplanArgs := `{\"steps\":[\"Step 1\", \"Step 2\", \"Step 3\"]}`\n\ttoolCall := schema.ToolCall{\n\t\tID:   \"tool_call_id\",\n\t\tType: \"function\",\n\t\tFunction: schema.FunctionCall{\n\t\t\tName:      \"plan\", // This should match PlanToolInfo.Name\n\t\t\tArguments: planArgs,\n\t\t},\n\t}\n\n\ttoolCallMsg := schema.AssistantMessage(\"\", nil)\n\ttoolCallMsg.ToolCalls = []schema.ToolCall{toolCall}\n\tsr, sw := schema.Pipe[*schema.Message](1)\n\tsw.Send(toolCallMsg, nil)\n\tsw.Close()\n\n\t// Mock the WithTools method to return a model that will be used for Generate\n\tmockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1)\n\n\t// Mock the Generate method to return the tool call message\n\tmockToolCallingModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\n\t// Create the PlannerConfig with ToolCallingChatModel\n\tconf := &PlannerConfig{\n\t\tToolCallingChatModel: mockToolCallingModel,\n\t\t// Use default instruction and tool info\n\t}\n\n\t// Create the planner\n\tp, err := NewPlanner(ctx, conf)\n\tassert.NoError(t, err)\n\n\t// Run the planner\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: p})\n\titerator := runner.Run(ctx, []adk.Message{schema.UserMessage(\"no input\")})\n\n\t// Get the event from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Nil(t, event.Err)\n\n\tmsg, _, err := adk.GetMessage(event)\n\tassert.NoError(t, err)\n\tassert.Equal(t, planArgs, msg.Content)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n\n\tplan := defaultNewPlan(ctx)\n\terr = plan.UnmarshalJSON([]byte(msg.Content))\n\tassert.NoError(t, err)\n\tplan_ := plan.(*defaultPlan)\n\tassert.NoError(t, err)\n\tassert.Equal(t, 3, len(plan_.Steps))\n\tassert.Equal(t, \"Step 1\", plan_.Steps[0])\n\tassert.Equal(t, \"Step 2\", plan_.Steps[1])\n\tassert.Equal(t, \"Step 3\", plan_.Steps[2])\n}\n\n// TestNewExecutor tests the NewExecutor function\nfunc TestNewExecutor(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock tool calling chat model\n\tmockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// Create the ExecutorConfig\n\tconf := &ExecutorConfig{\n\t\tModel:         mockToolCallingModel,\n\t\tMaxIterations: 3,\n\t}\n\n\t// Create the executor\n\texecutor, err := NewExecutor(ctx, conf)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, executor)\n\n\t// Verify the executor's name and description\n\tassert.Equal(t, \"executor\", executor.Name(ctx))\n\tassert.Equal(t, \"an executor agent\", executor.Description(ctx))\n}\n\n// TestExecutorRun tests the Run method of the executor\nfunc TestExecutorRun(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock tool calling chat model\n\tmockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// Store a plan in the session\n\tplan := &defaultPlan{Steps: []string{\"Step 1\", \"Step 2\", \"Step 3\"}}\n\tadk.AddSessionValue(ctx, PlanSessionKey, plan)\n\n\t// Set up expectations for the mock model\n\t// The model should return the last user message as its response\n\tmockToolCallingModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, messages []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t// Find the last user message\n\t\t\tvar lastUserMessage string\n\t\t\tfor _, msg := range messages {\n\t\t\t\tif msg.Role == schema.User {\n\t\t\t\t\tlastUserMessage = msg.Content\n\t\t\t\t}\n\t\t\t}\n\t\t\t// Return the last user message as the model's response\n\t\t\treturn schema.AssistantMessage(lastUserMessage, nil), nil\n\t\t}).Times(1)\n\n\t// Create the ExecutorConfig\n\tconf := &ExecutorConfig{\n\t\tModel:         mockToolCallingModel,\n\t\tMaxIterations: 3,\n\t}\n\n\t// Create the executor\n\texecutor, err := NewExecutor(ctx, conf)\n\tassert.NoError(t, err)\n\n\t// Run the executor\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: executor})\n\titerator := runner.Run(ctx, []adk.Message{schema.UserMessage(\"no input\")},\n\t\tadk.WithSessionValues(map[string]any{\n\t\t\tPlanSessionKey:      plan,\n\t\t\tUserInputSessionKey: []adk.Message{schema.UserMessage(\"no input\")},\n\t\t}),\n\t)\n\n\t// Get the event from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Nil(t, event.Err)\n\tassert.NotNil(t, event.Output)\n\tassert.NotNil(t, event.Output.MessageOutput)\n\tmsg, _, err := adk.GetMessage(event)\n\tassert.NoError(t, err)\n\tt.Logf(\"executor model input msg:\\n %s\\n\", msg.Content)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\n// TestNewReplanner tests the NewReplanner function\nfunc TestNewReplanner(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock tool calling chat model\n\tmockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\t// Mock the WithTools method\n\tmockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1)\n\n\t// Create plan and respond tools\n\tplanTool := &schema.ToolInfo{\n\t\tName: \"Plan\",\n\t\tDesc: \"Plan tool\",\n\t}\n\n\trespondTool := &schema.ToolInfo{\n\t\tName: \"Respond\",\n\t\tDesc: \"Respond tool\",\n\t}\n\n\t// Create the ReplannerConfig\n\tconf := &ReplannerConfig{\n\t\tChatModel:   mockToolCallingModel,\n\t\tPlanTool:    planTool,\n\t\tRespondTool: respondTool,\n\t}\n\n\t// Create the replanner\n\trp, err := NewReplanner(ctx, conf)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, rp)\n\n\t// Verify the replanner's name and description\n\tassert.Equal(t, \"replanner\", rp.Name(ctx))\n\tassert.Equal(t, \"a replanner agent\", rp.Description(ctx))\n}\n\n// TestReplannerRunWithPlan tests the Replanner's ability to use the plan_tool\nfunc TestReplannerRunWithPlan(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock tool calling chat model\n\tmockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// Create plan and respond tools\n\tplanTool := &schema.ToolInfo{\n\t\tName: \"Plan\",\n\t\tDesc: \"Plan tool\",\n\t}\n\n\trespondTool := &schema.ToolInfo{\n\t\tName: \"Respond\",\n\t\tDesc: \"Respond tool\",\n\t}\n\n\t// Create a tool call response for the Plan tool\n\tplanArgs := `{\"steps\":[\"Updated Step 1\", \"Updated Step 2\"]}`\n\ttoolCall := schema.ToolCall{\n\t\tID:   \"tool_call_id\",\n\t\tType: \"function\",\n\t\tFunction: schema.FunctionCall{\n\t\t\tName:      planTool.Name,\n\t\t\tArguments: planArgs,\n\t\t},\n\t}\n\n\ttoolCallMsg := schema.AssistantMessage(\"\", nil)\n\ttoolCallMsg.ToolCalls = []schema.ToolCall{toolCall}\n\tsr, sw := schema.Pipe[*schema.Message](1)\n\tsw.Send(toolCallMsg, nil)\n\tsw.Close()\n\n\t// Mock the Generate method\n\tmockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1)\n\tmockToolCallingModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\n\t// Create the ReplannerConfig\n\tconf := &ReplannerConfig{\n\t\tChatModel:   mockToolCallingModel,\n\t\tPlanTool:    planTool,\n\t\tRespondTool: respondTool,\n\t}\n\n\t// Create the replanner\n\trp, err := NewReplanner(ctx, conf)\n\tassert.NoError(t, err)\n\n\t// Store necessary values in the session\n\tplan := &defaultPlan{Steps: []string{\"Step 1\", \"Step 2\", \"Step 3\"}}\n\n\trp, err = agentOutputSessionKVs(ctx, rp)\n\tassert.NoError(t, err)\n\n\t// Run the replanner\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: rp})\n\titerator := runner.Run(ctx, []adk.Message{schema.UserMessage(\"no input\")},\n\t\tadk.WithSessionValues(map[string]any{\n\t\t\tPlanSessionKey:         plan,\n\t\t\tExecutedStepSessionKey: \"Execution result\",\n\t\t\tUserInputSessionKey:    []adk.Message{schema.UserMessage(\"User input\")},\n\t\t}),\n\t)\n\n\t// Get the event from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Nil(t, event.Err)\n\n\tevent, ok = iterator.Next()\n\tassert.True(t, ok)\n\tkvs := event.Output.CustomizedOutput.(map[string]any)\n\tassert.Greater(t, len(kvs), 0)\n\n\t// Verify the updated plan was stored in the session\n\tplanValue, ok := kvs[PlanSessionKey]\n\tassert.True(t, ok)\n\tupdatedPlan, ok := planValue.(*defaultPlan)\n\tassert.True(t, ok)\n\tassert.Equal(t, 2, len(updatedPlan.Steps))\n\tassert.Equal(t, \"Updated Step 1\", updatedPlan.Steps[0])\n\tassert.Equal(t, \"Updated Step 2\", updatedPlan.Steps[1])\n\n\t// Verify the execute results were updated\n\texecuteResultsValue, ok := kvs[ExecutedStepsSessionKey]\n\tassert.True(t, ok)\n\texecuteResults, ok := executeResultsValue.([]ExecutedStep)\n\tassert.True(t, ok)\n\tassert.Equal(t, 1, len(executeResults))\n\tassert.Equal(t, \"Step 1\", executeResults[0].Step)\n\tassert.Equal(t, \"Execution result\", executeResults[0].Result)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\n// TestReplannerRunWithRespond tests the Replanner's ability to use the respond_tool\nfunc TestReplannerRunWithRespond(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create a mock tool calling chat model\n\tmockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// Create plan and respond tools\n\tplanTool := &schema.ToolInfo{\n\t\tName: \"Plan\",\n\t\tDesc: \"Plan tool\",\n\t}\n\n\trespondTool := &schema.ToolInfo{\n\t\tName: \"Respond\",\n\t\tDesc: \"Respond tool\",\n\t}\n\n\t// Create a tool call response for the Respond tool\n\tresponseArgs := `{\"response\":\"This is the final response to the user\"}`\n\ttoolCall := schema.ToolCall{\n\t\tID:   \"tool_call_id\",\n\t\tType: \"function\",\n\t\tFunction: schema.FunctionCall{\n\t\t\tName:      respondTool.Name,\n\t\t\tArguments: responseArgs,\n\t\t},\n\t}\n\n\ttoolCallMsg := schema.AssistantMessage(\"\", nil)\n\ttoolCallMsg.ToolCalls = []schema.ToolCall{toolCall}\n\tsr, sw := schema.Pipe[*schema.Message](1)\n\tsw.Send(toolCallMsg, nil)\n\tsw.Close()\n\n\t// Mock the Generate method\n\tmockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1)\n\tmockToolCallingModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\n\t// Create the ReplannerConfig\n\tconf := &ReplannerConfig{\n\t\tChatModel:   mockToolCallingModel,\n\t\tPlanTool:    planTool,\n\t\tRespondTool: respondTool,\n\t}\n\n\t// Create the replanner\n\trp, err := NewReplanner(ctx, conf)\n\tassert.NoError(t, err)\n\n\t// Store necessary values in the session\n\tplan := &defaultPlan{Steps: []string{\"Step 1\", \"Step 2\", \"Step 3\"}}\n\n\t// Run the replanner\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: rp})\n\titerator := runner.Run(ctx, []adk.Message{schema.UserMessage(\"no input\")},\n\t\tadk.WithSessionValues(map[string]any{\n\t\t\tPlanSessionKey:         plan,\n\t\t\tExecutedStepSessionKey: \"Execution result\",\n\t\t\tUserInputSessionKey:    []adk.Message{schema.UserMessage(\"User input\")},\n\t\t}),\n\t)\n\n\t// Get the event from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Nil(t, event.Err)\n\tmsg, _, err := adk.GetMessage(event)\n\tassert.NoError(t, err)\n\tassert.Equal(t, responseArgs, msg.Content)\n\n\t// Verify that an exit action was generated\n\tevent, ok = iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event.Action)\n\tassert.NotNil(t, event.Action.BreakLoop)\n\tassert.False(t, event.Action.BreakLoop.Done)\n\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\n// TestNewPlanExecuteAgent tests the New function\nfunc TestNewPlanExecuteAgent(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create mock agents\n\tmockPlanner := mockAdk.NewMockAgent(ctrl)\n\tmockExecutor := mockAdk.NewMockAgent(ctrl)\n\tmockReplanner := mockAdk.NewMockAgent(ctrl)\n\n\t// Set up expectations for the mock agents\n\tmockPlanner.EXPECT().Name(gomock.Any()).Return(\"planner\").AnyTimes()\n\tmockPlanner.EXPECT().Description(gomock.Any()).Return(\"a planner agent\").AnyTimes()\n\n\tmockExecutor.EXPECT().Name(gomock.Any()).Return(\"executor\").AnyTimes()\n\tmockExecutor.EXPECT().Description(gomock.Any()).Return(\"an executor agent\").AnyTimes()\n\n\tmockReplanner.EXPECT().Name(gomock.Any()).Return(\"replanner\").AnyTimes()\n\tmockReplanner.EXPECT().Description(gomock.Any()).Return(\"a replanner agent\").AnyTimes()\n\n\tconf := &Config{\n\t\tPlanner:   mockPlanner,\n\t\tExecutor:  mockExecutor,\n\t\tReplanner: mockReplanner,\n\t}\n\n\t// Create the plan execute agent\n\tagent, err := New(ctx, conf)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, agent)\n}\n\nfunc TestPlanExecuteAgentWithReplan(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create mock agents\n\tmockPlanner := mockAdk.NewMockAgent(ctrl)\n\tmockExecutor := mockAdk.NewMockAgent(ctrl)\n\tmockReplanner := mockAdk.NewMockAgent(ctrl)\n\n\t// Set up expectations for the mock agents\n\tmockPlanner.EXPECT().Name(gomock.Any()).Return(\"planner\").AnyTimes()\n\tmockPlanner.EXPECT().Description(gomock.Any()).Return(\"a planner agent\").AnyTimes()\n\n\tmockExecutor.EXPECT().Name(gomock.Any()).Return(\"executor\").AnyTimes()\n\tmockExecutor.EXPECT().Description(gomock.Any()).Return(\"an executor agent\").AnyTimes()\n\n\tmockReplanner.EXPECT().Name(gomock.Any()).Return(\"replanner\").AnyTimes()\n\tmockReplanner.EXPECT().Description(gomock.Any()).Return(\"a replanner agent\").AnyTimes()\n\n\t// Create a plan\n\toriginalPlan := &defaultPlan{Steps: []string{\"Step 1\", \"Step 2\", \"Step 3\"}}\n\t// Create an updated plan with fewer steps (after replanning)\n\tupdatedPlan := &defaultPlan{Steps: []string{\"Updated Step 2\", \"Updated Step 3\"}}\n\t// Create execute result\n\toriginalExecuteResult := \"Execution result for Step 1\"\n\tupdatedExecuteResult := \"Execution result for Updated Step 2\"\n\n\t// Create user input\n\tuserInput := []adk.Message{schema.UserMessage(\"User task input\")}\n\n\tfinalResponse := &Response{Response: \"Final response to user after executing all steps\"}\n\n\t// Mock the planner Run method to set the original plan\n\tmockPlanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\t\t\titerator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\n\t\t\t// Set the plan in the session\n\t\t\tadk.AddSessionValue(ctx, PlanSessionKey, originalPlan)\n\t\t\tadk.AddSessionValue(ctx, UserInputSessionKey, userInput)\n\n\t\t\t// Send a message event\n\t\t\tplanJSON, _ := sonic.MarshalString(originalPlan)\n\t\t\tmsg := schema.AssistantMessage(planJSON, nil)\n\t\t\tevent := adk.EventFromMessage(msg, nil, schema.Assistant, \"\")\n\t\t\tgenerator.Send(event)\n\t\t\tgenerator.Close()\n\n\t\t\treturn iterator\n\t\t},\n\t).Times(1)\n\n\t// Mock the executor Run method to set the execute result\n\tmockExecutor.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\t\t\titerator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\n\t\t\tplan, _ := adk.GetSessionValue(ctx, PlanSessionKey)\n\t\t\tcurrentPlan := plan.(*defaultPlan)\n\t\t\tvar msg adk.Message\n\t\t\t// Check if this is the first replanning (original plan has 3 steps)\n\t\t\tif len(currentPlan.Steps) == 3 {\n\t\t\t\tmsg = schema.AssistantMessage(originalExecuteResult, nil)\n\t\t\t\tadk.AddSessionValue(ctx, ExecutedStepSessionKey, originalExecuteResult)\n\t\t\t} else {\n\t\t\t\tmsg = schema.AssistantMessage(updatedExecuteResult, nil)\n\t\t\t\tadk.AddSessionValue(ctx, ExecutedStepSessionKey, updatedExecuteResult)\n\t\t\t}\n\t\t\tevent := adk.EventFromMessage(msg, nil, schema.Assistant, \"\")\n\t\t\tgenerator.Send(event)\n\t\t\tgenerator.Close()\n\n\t\t\treturn iterator\n\t\t},\n\t).Times(2)\n\n\t// Mock the replanner Run method to first update the plan, then respond to user\n\tmockReplanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\t\t\titerator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\n\t\t\t// First call: Update the plan\n\t\t\t// Get the current plan from the session\n\t\t\tplan, _ := adk.GetSessionValue(ctx, PlanSessionKey)\n\t\t\tcurrentPlan := plan.(*defaultPlan)\n\n\t\t\t// Check if this is the first replanning (original plan has 3 steps)\n\t\t\tif len(currentPlan.Steps) == 3 {\n\t\t\t\t// Send a message event with the updated plan\n\t\t\t\tplanJSON, _ := sonic.MarshalString(updatedPlan)\n\t\t\t\tmsg := schema.AssistantMessage(planJSON, nil)\n\t\t\t\tevent := adk.EventFromMessage(msg, nil, schema.Assistant, \"\")\n\t\t\t\tgenerator.Send(event)\n\n\t\t\t\t// Set the updated plan & execute result in the session\n\t\t\t\tadk.AddSessionValue(ctx, PlanSessionKey, updatedPlan)\n\t\t\t\tadk.AddSessionValue(ctx, ExecutedStepsSessionKey, []ExecutedStep{{\n\t\t\t\t\tStep:   currentPlan.Steps[0],\n\t\t\t\t\tResult: originalExecuteResult,\n\t\t\t\t}})\n\t\t\t} else {\n\t\t\t\t// Second call: Respond to user\n\t\t\t\tresponseJSON, err := sonic.MarshalString(finalResponse)\n\t\t\t\tassert.NoError(t, err)\n\t\t\t\tmsg := schema.AssistantMessage(responseJSON, nil)\n\t\t\t\tevent := adk.EventFromMessage(msg, nil, schema.Assistant, \"\")\n\t\t\t\tgenerator.Send(event)\n\n\t\t\t\t// Send exit action\n\t\t\t\taction := adk.NewExitAction()\n\t\t\t\tgenerator.Send(&adk.AgentEvent{Action: action})\n\t\t\t}\n\n\t\t\tgenerator.Close()\n\t\t\treturn iterator\n\t\t},\n\t).Times(2)\n\n\tconf := &Config{\n\t\tPlanner:   mockPlanner,\n\t\tExecutor:  mockExecutor,\n\t\tReplanner: mockReplanner,\n\t}\n\n\t// Create the plan execute agent\n\tagent, err := New(ctx, conf)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, agent)\n\n\t// Run the agent\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent})\n\titerator := runner.Run(ctx, userInput)\n\n\t// Collect all events\n\tvar events []*adk.AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\t// Verify the events\n\tassert.Greater(t, len(events), 0)\n\n\tfor i, event := range events {\n\t\teventJSON, e := sonic.MarshalString(event)\n\t\tassert.NoError(t, e)\n\t\tt.Logf(\"event %d:\\n%s\", i, eventJSON)\n\t}\n}\n\ntype interruptibleTool struct {\n\tname string\n\tt    *testing.T\n}\n\nfunc (m *interruptibleTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: m.name,\n\t\tDesc: \"A tool that requires human approval before execution\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"action\": {\n\t\t\t\tType:     schema.String,\n\t\t\t\tDesc:     \"The action to perform\",\n\t\t\t\tRequired: true,\n\t\t\t},\n\t\t}),\n\t}, nil\n}\n\nfunc (m *interruptibleTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\twasInterrupted, _, _ := tool.GetInterruptState[any](ctx)\n\tif !wasInterrupted {\n\t\treturn \"\", tool.Interrupt(ctx, fmt.Sprintf(\"Tool '%s' requires human approval\", m.name))\n\t}\n\n\tisResumeTarget, hasData, data := tool.GetResumeContext[string](ctx)\n\tif !isResumeTarget {\n\t\treturn \"\", tool.Interrupt(ctx, fmt.Sprintf(\"Tool '%s' requires human approval\", m.name))\n\t}\n\n\tif hasData {\n\t\treturn fmt.Sprintf(\"Approved action executed with data: %s\", data), nil\n\t}\n\treturn \"Approved action executed\", nil\n}\n\ntype checkpointStore struct {\n\tdata map[string][]byte\n}\n\nfunc newCheckpointStore() *checkpointStore {\n\treturn &checkpointStore{data: make(map[string][]byte)}\n}\n\nfunc (s *checkpointStore) Set(_ context.Context, key string, value []byte) error {\n\ts.data[key] = value\n\treturn nil\n}\n\nfunc (s *checkpointStore) Get(_ context.Context, key string) ([]byte, bool, error) {\n\tv, ok := s.data[key]\n\treturn v, ok, nil\n}\n\nfunc formatRunPath(runPath []adk.RunStep) string {\n\tif len(runPath) == 0 {\n\t\treturn \"[]\"\n\t}\n\tvar parts []string\n\tfor _, step := range runPath {\n\t\tparts = append(parts, step.String())\n\t}\n\treturn \"[\" + strings.Join(parts, \" -> \") + \"]\"\n}\n\nfunc formatAgentEvent(event *adk.AgentEvent) string {\n\tvar sb strings.Builder\n\tsb.WriteString(fmt.Sprintf(\"{AgentName: %q, RunPath: %s\", event.AgentName, formatRunPath(event.RunPath)))\n\tif event.Output != nil {\n\t\tif event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil {\n\t\t\tmsg := event.Output.MessageOutput.Message\n\t\t\tsb.WriteString(fmt.Sprintf(\", Output.Message: {Role: %q, Content: %q}\", msg.Role, msg.Content))\n\t\t}\n\t}\n\tif event.Action != nil {\n\t\tif event.Action.Interrupted != nil {\n\t\t\tsb.WriteString(fmt.Sprintf(\", Action.Interrupted: {%d contexts}\", len(event.Action.Interrupted.InterruptContexts)))\n\t\t}\n\t\tif event.Action.BreakLoop != nil {\n\t\t\tsb.WriteString(fmt.Sprintf(\", Action.BreakLoop: {From: %q, Done: %v}\", event.Action.BreakLoop.From, event.Action.BreakLoop.Done))\n\t\t}\n\t}\n\tif event.Err != nil {\n\t\tsb.WriteString(fmt.Sprintf(\", Err: %v\", event.Err))\n\t}\n\tsb.WriteString(\"}\")\n\treturn sb.String()\n}\n\nfunc TestPlanExecuteAgentInterruptResume(t *testing.T) {\n\tctx := context.Background()\n\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tmockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tapprovalTool := &interruptibleTool{name: \"approve_action\", t: t}\n\n\tplan := &defaultPlan{Steps: []string{\"Execute action requiring approval\", \"Complete task\"}}\n\tuserInput := []adk.Message{schema.UserMessage(\"Please execute the action\")}\n\n\tmockPlanner := mockAdk.NewMockAgent(ctrl)\n\tmockPlanner.EXPECT().Name(gomock.Any()).Return(\"planner\").AnyTimes()\n\tmockPlanner.EXPECT().Description(gomock.Any()).Return(\"a planner agent\").AnyTimes()\n\n\tmockPlanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\t\t\titerator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\n\t\t\tadk.AddSessionValue(ctx, PlanSessionKey, plan)\n\t\t\tadk.AddSessionValue(ctx, UserInputSessionKey, userInput)\n\n\t\t\tplanJSON, _ := sonic.MarshalString(plan)\n\t\t\tmsg := schema.AssistantMessage(planJSON, nil)\n\t\t\tevent := adk.EventFromMessage(msg, nil, schema.Assistant, \"\")\n\t\t\tgenerator.Send(event)\n\t\t\tgenerator.Close()\n\n\t\t\treturn iterator\n\t\t},\n\t).Times(1)\n\n\ttoolCallMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"call_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"approve_action\",\n\t\t\t\tArguments: `{\"action\": \"execute\"}`,\n\t\t\t},\n\t\t},\n\t})\n\n\tcompletionMsg := schema.AssistantMessage(\"Action approved and executed successfully\", nil)\n\n\tmockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).AnyTimes()\n\tmockToolCallingModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(toolCallMsg, nil).Times(1)\n\tmockToolCallingModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(completionMsg, nil).AnyTimes()\n\n\texecutor, err := NewExecutor(ctx, &ExecutorConfig{\n\t\tModel: mockToolCallingModel,\n\t\tToolsConfig: adk.ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{approvalTool},\n\t\t\t},\n\t\t},\n\t\tMaxIterations: 5,\n\t})\n\tassert.NoError(t, err)\n\n\tmockReplanner := mockAdk.NewMockAgent(ctrl)\n\tmockReplanner.EXPECT().Name(gomock.Any()).Return(\"replanner\").AnyTimes()\n\tmockReplanner.EXPECT().Description(gomock.Any()).Return(\"a replanner agent\").AnyTimes()\n\n\tmockReplanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\t\t\titerator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\n\t\t\tresponseJSON := `{\"response\":\"Task completed successfully\"}`\n\t\t\tmsg := schema.AssistantMessage(responseJSON, nil)\n\t\t\tevent := adk.EventFromMessage(msg, nil, schema.Assistant, \"\")\n\t\t\tgenerator.Send(event)\n\n\t\t\taction := adk.NewBreakLoopAction(\"replanner\")\n\t\t\tgenerator.Send(&adk.AgentEvent{Action: action})\n\n\t\t\tgenerator.Close()\n\t\t\treturn iterator\n\t\t},\n\t).AnyTimes()\n\n\tagent, err := New(ctx, &Config{\n\t\tPlanner:       mockPlanner,\n\t\tExecutor:      executor,\n\t\tReplanner:     mockReplanner,\n\t\tMaxIterations: 5,\n\t})\n\tassert.NoError(t, err)\n\n\tstore := newCheckpointStore()\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{\n\t\tAgent:           agent,\n\t\tCheckPointStore: store,\n\t})\n\n\titer := runner.Run(ctx, userInput, adk.WithCheckPointID(\"test-interrupt-1\"))\n\n\tvar events []*adk.AgentEvent\n\tvar interruptEvent *adk.AgentEvent\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tinterruptEvent = event\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\tt.Logf(\"Total events received: %d\", len(events))\n\tfor i, event := range events {\n\t\teventJSON, _ := sonic.MarshalString(event)\n\t\tt.Logf(\"Event %d: %s\", i, eventJSON)\n\t}\n\n\tif interruptEvent == nil {\n\t\tt.Fatal(\"Expected an interrupt event from the tool, but none was received\")\n\t}\n\n\tassert.NotNil(t, interruptEvent.Action.Interrupted, \"Should have interrupt info\")\n\tassert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, \"Should have interrupt contexts\")\n\n\tt.Logf(\"Interrupt event received with %d contexts\", len(interruptEvent.Action.Interrupted.InterruptContexts))\n\tfor i, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tt.Logf(\"Interrupt context %d: ID=%s, Info=%v, Address=%v\", i, ctx.ID, ctx.Info, ctx.Address)\n\t}\n\n\tvar toolInterruptID string\n\tfor _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif intCtx.IsRootCause {\n\t\t\ttoolInterruptID = intCtx.ID\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.NotEmpty(t, toolInterruptID, \"Should have a root cause interrupt ID\")\n\n\tt.Logf(\"Attempting to resume with interrupt ID: %s\", toolInterruptID)\n\n\tresumeIter, err := runner.ResumeWithParams(ctx, \"test-interrupt-1\", &adk.ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\ttoolInterruptID: \"approved\",\n\t\t},\n\t})\n\tassert.NoError(t, err, \"Resume should not error\")\n\tassert.NotNil(t, resumeIter, \"Resume iterator should not be nil\")\n\n\tvar resumeEvents []*adk.AgentEvent\n\tfor {\n\t\tevent, ok := resumeIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tresumeEvents = append(resumeEvents, event)\n\t}\n\n\tassert.NotEmpty(t, resumeEvents, \"Should have resume events\")\n\n\tfor _, event := range resumeEvents {\n\t\tassert.NoError(t, event.Err, \"Resume event should not have error\")\n\t}\n\n\tvar hasToolResponse, hasAssistantCompletion, hasBreakLoop bool\n\tfor _, event := range resumeEvents {\n\t\tif event.Output != nil && event.Output.MessageOutput != nil {\n\t\t\tmsg := event.Output.MessageOutput.Message\n\t\t\tif msg != nil {\n\t\t\t\tif msg.Role == \"tool\" && strings.Contains(msg.Content, \"Approved action executed\") {\n\t\t\t\t\thasToolResponse = true\n\t\t\t\t}\n\t\t\t\tif msg.Role == \"assistant\" && strings.Contains(msg.Content, \"approved\") {\n\t\t\t\t\thasAssistantCompletion = true\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t\tif event.Action != nil && event.Action.BreakLoop != nil && event.Action.BreakLoop.Done {\n\t\t\thasBreakLoop = true\n\t\t}\n\t}\n\n\tassert.True(t, hasToolResponse, \"Should have tool response with approved action\")\n\tassert.True(t, hasAssistantCompletion, \"Should have assistant completion message\")\n\tassert.True(t, hasBreakLoop, \"Should have break loop action indicating completion\")\n}\n"
  },
  {
    "path": "adk/prebuilt/planexecute/utils.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage planexecute\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/adk\"\n)\n\ntype outputSessionKVsAgent struct {\n\tadk.Agent\n}\n\nfunc (o *outputSessionKVsAgent) Run(ctx context.Context, input *adk.AgentInput,\n\toptions ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\n\titerator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\n\titerator_ := o.Agent.Run(ctx, input, options...)\n\tgo func() {\n\t\tdefer generator.Close()\n\t\tfor {\n\t\t\tevent, ok := iterator_.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tgenerator.Send(event)\n\t\t}\n\n\t\tkvs := adk.GetSessionValues(ctx)\n\n\t\tevent := &adk.AgentEvent{\n\t\t\tOutput: &adk.AgentOutput{CustomizedOutput: kvs},\n\t\t}\n\t\tgenerator.Send(event)\n\t}()\n\n\treturn iterator\n}\n\nfunc agentOutputSessionKVs(ctx context.Context, agent adk.Agent) (adk.Agent, error) {\n\treturn &outputSessionKVsAgent{Agent: agent}, nil\n}\n"
  },
  {
    "path": "adk/prebuilt/supervisor/supervisor.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package supervisor implements the supervisor pattern for multi-agent systems,\n// where a central agent coordinates a set of sub-agents.\n//\n// # Unified Tracing\n//\n// The supervisor pattern provides unified tracing support through an internal container.\n// When using callbacks (e.g., for tracing or observability), the entire supervisor structure\n// (supervisor agent + all sub-agents) shares a single trace root. This means:\n//   - OnStart is invoked once at the supervisor container level\n//   - The callback-enriched context (containing parent span info) is propagated to all agents\n//   - All agents within the supervisor appear as children of the same trace root\n//\n// This is achieved by wrapping the supervisor structure in an internal container that acts\n// as the single entry point for tracing. The container delegates all execution to the\n// underlying agents while providing a unified identity for callbacks.\npackage supervisor\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/adk\"\n)\n\ntype Config struct {\n\t// Supervisor specifies the agent that will act as the supervisor, coordinating and managing the sub-agents.\n\tSupervisor adk.Agent\n\n\t// SubAgents specifies the list of agents that will be supervised and coordinated by the supervisor agent.\n\tSubAgents []adk.Agent\n}\n\n// supervisorContainer wraps the entire supervisor structure to provide unified tracing.\n// When callbacks are registered (e.g., via Runner.Query with WithCallbacks), OnStart/OnEnd\n// are invoked once for this container, creating a single trace root. The callback-enriched\n// context is then propagated to the supervisor and all sub-agents, ensuring they share\n// the same trace parent.\n//\n// This container implements Agent and ResumableAgent by delegating to the inner agent.\n// It provides its own Name and GetType for callback identification.\ntype supervisorContainer struct {\n\tname  string\n\tinner adk.ResumableAgent\n}\n\nfunc (s *supervisorContainer) Name(_ context.Context) string {\n\treturn s.name\n}\n\nfunc (s *supervisorContainer) Description(ctx context.Context) string {\n\treturn s.inner.Description(ctx)\n}\n\nfunc (s *supervisorContainer) GetType() string {\n\treturn \"Supervisor\"\n}\n\nfunc (s *supervisorContainer) Run(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\treturn s.inner.Run(ctx, input, opts...)\n}\n\nfunc (s *supervisorContainer) Resume(ctx context.Context, info *adk.ResumeInfo, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\treturn s.inner.Resume(ctx, info, opts...)\n}\n\n// New creates a supervisor-based multi-agent system with the given configuration.\n//\n// In the supervisor pattern, a designated supervisor agent coordinates multiple sub-agents.\n// The supervisor can delegate tasks to sub-agents and receive their responses, while\n// sub-agents can only communicate with the supervisor (not with each other directly).\n// This hierarchical structure enables complex problem-solving through coordinated agent interactions.\n//\n// The returned agent is wrapped in an internal container that provides unified tracing.\n// When used with Runner and callbacks, all agents within the supervisor structure will\n// share the same trace root, making it easy to observe the entire multi-agent execution\n// as a single logical unit.\nfunc New(ctx context.Context, conf *Config) (adk.ResumableAgent, error) {\n\tsubAgents := make([]adk.Agent, 0, len(conf.SubAgents))\n\tsupervisorName := conf.Supervisor.Name(ctx)\n\tfor _, subAgent := range conf.SubAgents {\n\t\tsubAgents = append(subAgents, adk.AgentWithDeterministicTransferTo(ctx, &adk.DeterministicTransferConfig{\n\t\t\tAgent:        subAgent,\n\t\t\tToAgentNames: []string{supervisorName},\n\t\t}))\n\t}\n\n\tinner, err := adk.SetSubAgents(ctx, conf.Supervisor, subAgents)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &supervisorContainer{\n\t\tname:  supervisorName,\n\t\tinner: inner,\n\t}, nil\n}\n"
  },
  {
    "path": "adk/prebuilt/supervisor/supervisor_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage supervisor\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockAdk \"github.com/cloudwego/eino/internal/mock/adk\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// TestNewSupervisor tests the New function\nfunc TestNewSupervisor(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock controller\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\t// Create mock agents\n\tsupervisorAgent := mockAdk.NewMockAgent(ctrl)\n\tsubAgent1 := mockAdk.NewMockAgent(ctrl)\n\tsubAgent2 := mockAdk.NewMockAgent(ctrl)\n\n\tsupervisorAgent.EXPECT().Name(gomock.Any()).Return(\"SupervisorAgent\").AnyTimes()\n\tsupervisorAgent.EXPECT().Description(gomock.Any()).Return(\"Supervisor agent description\").AnyTimes()\n\tsubAgent1.EXPECT().Name(gomock.Any()).Return(\"SubAgent1\").AnyTimes()\n\tsubAgent2.EXPECT().Name(gomock.Any()).Return(\"SubAgent2\").AnyTimes()\n\n\taMsg, tMsg := adk.GenTransferMessages(ctx, \"SubAgent1\")\n\ti, g := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tg.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, \"\"))\n\tevent := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)\n\tevent.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: \"SubAgent1\"}}\n\tg.Send(event)\n\tg.Close()\n\tsupervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)\n\n\ti, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tsubAgent1Msg := schema.AssistantMessage(\"SubAgent1\", nil)\n\tg.Send(adk.EventFromMessage(subAgent1Msg, nil, schema.Assistant, \"\"))\n\tg.Close()\n\tsubAgent1.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)\n\n\taMsg, tMsg = adk.GenTransferMessages(ctx, \"SubAgent2 message\")\n\ti, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tg.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, \"\"))\n\tevent = adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)\n\tevent.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: \"SubAgent2\"}}\n\tg.Send(event)\n\tg.Close()\n\tsupervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)\n\n\ti, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tsubAgent2Msg := schema.AssistantMessage(\"SubAgent2 message\", nil)\n\tg.Send(adk.EventFromMessage(subAgent2Msg, nil, schema.Assistant, \"\"))\n\tg.Close()\n\tsubAgent2.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)\n\n\ti, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tfinishMsg := schema.AssistantMessage(\"finish\", nil)\n\tg.Send(adk.EventFromMessage(finishMsg, nil, schema.Assistant, \"\"))\n\tg.Close()\n\tsupervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)\n\n\tconf := &Config{\n\t\tSupervisor: supervisorAgent,\n\t\tSubAgents:  []adk.Agent{subAgent1, subAgent2},\n\t}\n\n\tmultiAgent, err := New(ctx, conf)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, multiAgent)\n\tassert.Equal(t, \"SupervisorAgent\", multiAgent.Name(ctx))\n\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent})\n\taIter := runner.Run(ctx, []adk.Message{schema.UserMessage(\"test\")})\n\n\t// transfer to agent1\n\tevent, ok := aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SupervisorAgent\", event.AgentName)\n\tassert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)\n\tassert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls))\n\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SupervisorAgent\", event.AgentName)\n\tassert.Equal(t, schema.Tool, event.Output.MessageOutput.Role)\n\tassert.Equal(t, \"SubAgent1\", event.Action.TransferToAgent.DestAgentName)\n\n\t// agent1's output\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SubAgent1\", event.AgentName)\n\tassert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)\n\tassert.Equal(t, subAgent1Msg.Content, event.Output.MessageOutput.Message.Content)\n\n\t// transfer back to supervisor\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SubAgent1\", event.AgentName)\n\tassert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)\n\tassert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls))\n\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SubAgent1\", event.AgentName)\n\tassert.Equal(t, schema.Tool, event.Output.MessageOutput.Role)\n\tassert.Equal(t, \"SupervisorAgent\", event.Action.TransferToAgent.DestAgentName)\n\n\t// transfer to agent2\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SupervisorAgent\", event.AgentName)\n\tassert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)\n\tassert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls))\n\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SupervisorAgent\", event.AgentName)\n\tassert.Equal(t, schema.Tool, event.Output.MessageOutput.Role)\n\tassert.Equal(t, \"SubAgent2\", event.Action.TransferToAgent.DestAgentName)\n\n\t// agent1's output\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SubAgent2\", event.AgentName)\n\tassert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)\n\tassert.Equal(t, subAgent2Msg.Content, event.Output.MessageOutput.Message.Content)\n\n\t// transfer back to supervisor\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SubAgent2\", event.AgentName)\n\tassert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)\n\tassert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls))\n\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SubAgent2\", event.AgentName)\n\tassert.Equal(t, schema.Tool, event.Output.MessageOutput.Role)\n\tassert.Equal(t, \"SupervisorAgent\", event.Action.TransferToAgent.DestAgentName)\n\n\t// finish\n\tevent, ok = aIter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"SupervisorAgent\", event.AgentName)\n\tassert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)\n\tassert.Equal(t, finishMsg.Content, event.Output.MessageOutput.Message.Content)\n}\n\ntype approvalInfo struct {\n\tToolName        string\n\tArgumentsInJSON string\n\tToolCallID      string\n}\n\nfunc (ai *approvalInfo) String() string {\n\treturn fmt.Sprintf(\"tool '%s' interrupted with arguments '%s', waiting for approval\",\n\t\tai.ToolName, ai.ArgumentsInJSON)\n}\n\ntype approvalResult struct {\n\tApproved         bool\n\tDisapproveReason *string\n}\n\nfunc init() {\n\tschema.Register[*approvalInfo]()\n\tschema.Register[*approvalResult]()\n}\n\ntype approvableTool struct {\n\tname string\n\tt    *testing.T\n}\n\nfunc (m *approvableTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: m.name,\n\t\tDesc: \"A tool that requires approval before execution\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"action\": {Type: schema.String, Desc: \"The action to perform\"},\n\t\t}),\n\t}, nil\n}\n\nfunc (m *approvableTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\twasInterrupted, _, storedArguments := tool.GetInterruptState[string](ctx)\n\tif !wasInterrupted {\n\t\treturn \"\", tool.StatefulInterrupt(ctx, &approvalInfo{\n\t\t\tToolName:        m.name,\n\t\t\tArgumentsInJSON: argumentsInJSON,\n\t\t\tToolCallID:      compose.GetToolCallID(ctx),\n\t\t}, argumentsInJSON)\n\t}\n\n\tisResumeTarget, hasData, data := tool.GetResumeContext[*approvalResult](ctx)\n\tif !isResumeTarget {\n\t\treturn \"\", tool.StatefulInterrupt(ctx, &approvalInfo{\n\t\t\tToolName:        m.name,\n\t\t\tArgumentsInJSON: storedArguments,\n\t\t\tToolCallID:      compose.GetToolCallID(ctx),\n\t\t}, storedArguments)\n\t}\n\n\tif !hasData {\n\t\treturn \"\", fmt.Errorf(\"tool '%s' resumed with no data\", m.name)\n\t}\n\n\tif data.Approved {\n\t\treturn fmt.Sprintf(\"Tool '%s' executed successfully with args: %s\", m.name, storedArguments), nil\n\t}\n\n\tif data.DisapproveReason != nil {\n\t\treturn fmt.Sprintf(\"Tool '%s' disapproved, reason: %s\", m.name, *data.DisapproveReason), nil\n\t}\n\n\treturn fmt.Sprintf(\"Tool '%s' disapproved\", m.name), nil\n}\n\ntype checkpointStore struct {\n\tdata map[string][]byte\n}\n\nfunc newCheckpointStore() *checkpointStore {\n\treturn &checkpointStore{data: make(map[string][]byte)}\n}\n\nfunc (s *checkpointStore) Set(_ context.Context, key string, value []byte) error {\n\ts.data[key] = value\n\treturn nil\n}\n\nfunc (s *checkpointStore) Get(_ context.Context, key string) ([]byte, bool, error) {\n\tv, ok := s.data[key]\n\treturn v, ok, nil\n}\n\ntype namedAgent struct {\n\tadk.ResumableAgent\n\tname        string\n\tdescription string\n}\n\nfunc (n *namedAgent) Name(_ context.Context) string {\n\treturn n.name\n}\n\nfunc (n *namedAgent) Description(_ context.Context) string {\n\treturn n.description\n}\n\nfunc TestNestedSupervisorInterruptResume(t *testing.T) {\n\tctx := context.Background()\n\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tmockOuterSupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockInnerSupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockWorkerModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tpaymentTool := &approvableTool{name: \"process_payment\", t: t}\n\n\tuserInput := []adk.Message{schema.UserMessage(\"Process a payment of $1000\")}\n\n\tmockWorkerModel.EXPECT().WithTools(gomock.Any()).Return(mockWorkerModel, nil).AnyTimes()\n\n\tworkerToolCallMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"call_payment_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"process_payment\",\n\t\t\t\tArguments: `{\"action\": \"process $1000 payment\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tmockWorkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(workerToolCallMsg, nil).Times(1)\n\n\tworkerCompletionMsg := schema.AssistantMessage(\"Payment processed successfully\", nil)\n\tmockWorkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(workerCompletionMsg, nil).AnyTimes()\n\n\tworkerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"payment_worker\",\n\t\tDescription: \"the agent responsible for processing payments\",\n\t\tInstruction: \"You are a payment processing worker. Use the process_payment tool to handle payments.\",\n\t\tModel:       mockWorkerModel,\n\t\tToolsConfig: adk.ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{paymentTool},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tmockInnerSupervisorModel.EXPECT().WithTools(gomock.Any()).Return(mockInnerSupervisorModel, nil).AnyTimes()\n\n\tinnerTransferMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"inner_transfer_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"transfer_to_agent\",\n\t\t\t\tArguments: `{\"agent_name\":\"payment_worker\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tmockInnerSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(innerTransferMsg, nil).Times(1)\n\n\tinnerFinalMsg := schema.AssistantMessage(\"Payment has been processed and approved.\", nil)\n\tmockInnerSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(innerFinalMsg, nil).AnyTimes()\n\n\tinnerSupervisorChatAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"payment_supervisor\",\n\t\tDescription: \"the supervisor agent responsible for payment operations\",\n\t\tInstruction: \"You are a payment supervisor. Delegate payment tasks to payment_worker.\",\n\t\tModel:       mockInnerSupervisorModel,\n\t\tExit:        &adk.ExitTool{},\n\t})\n\tassert.NoError(t, err)\n\n\tinnerSupervisorAgent, err := New(ctx, &Config{\n\t\tSupervisor: innerSupervisorChatAgent,\n\t\tSubAgents:  []adk.Agent{workerAgent},\n\t})\n\tassert.NoError(t, err)\n\n\tinnerSupervisorWrapped := &namedAgent{\n\t\tResumableAgent: innerSupervisorAgent,\n\t\tname:           \"payment_department\",\n\t\tdescription:    \"the department responsible for all payment-related operations\",\n\t}\n\n\tmockOuterSupervisorModel.EXPECT().WithTools(gomock.Any()).Return(mockOuterSupervisorModel, nil).AnyTimes()\n\n\touterTransferMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"outer_transfer_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"transfer_to_agent\",\n\t\t\t\tArguments: `{\"agent_name\":\"payment_department\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tmockOuterSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(outerTransferMsg, nil).Times(1)\n\n\touterFinalMsg := schema.AssistantMessage(\"The payment request has been fully processed by the payment department.\", nil)\n\tmockOuterSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(outerFinalMsg, nil).AnyTimes()\n\n\touterSupervisorChatAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"company_coordinator\",\n\t\tDescription: \"the top-level coordinator for company operations\",\n\t\tInstruction: \"You are the company coordinator. Route payment requests to payment_department.\",\n\t\tModel:       mockOuterSupervisorModel,\n\t\tExit:        &adk.ExitTool{},\n\t})\n\tassert.NoError(t, err)\n\n\touterSupervisorAgent, err := New(ctx, &Config{\n\t\tSupervisor: outerSupervisorChatAgent,\n\t\tSubAgents:  []adk.Agent{innerSupervisorWrapped},\n\t})\n\tassert.NoError(t, err)\n\n\touterSupervisorWrapped := &namedAgent{\n\t\tResumableAgent: outerSupervisorAgent,\n\t\tname:           \"headquarters\",\n\t\tdescription:    \"the company headquarters that coordinates all departments\",\n\t}\n\n\tstore := newCheckpointStore()\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{\n\t\tAgent:           outerSupervisorWrapped,\n\t\tCheckPointStore: store,\n\t})\n\n\tt.Log(\"========================================\")\n\tt.Log(\"Starting Nested Supervisor Integration Test (with namedAgent wrappers)\")\n\tt.Log(\"Hierarchy: headquarters(wrapper) -> company_coordinator -> payment_department(wrapper) -> payment_supervisor -> payment_worker -> process_payment tool\")\n\tt.Log(\"========================================\")\n\n\tcheckpointID := \"test-nested-supervisor-1\"\n\titer := runner.Run(ctx, userInput, adk.WithCheckPointID(checkpointID))\n\n\tvar interruptEvent *adk.AgentEvent\n\teventCount := 0\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\teventCount++\n\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tinterruptEvent = event\n\t\t\tt.Log(\"INTERRUPT DETECTED - Deep interrupt from tool within nested supervisor\")\n\t\t\tbreak\n\t\t}\n\t}\n\n\tif interruptEvent == nil {\n\t\tt.Fatal(\"Expected an interrupt event from the process_payment tool, but none was received\")\n\t}\n\n\tassert.NotNil(t, interruptEvent.Action.Interrupted, \"Should have interrupt info\")\n\tassert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, \"Should have interrupt contexts\")\n\n\tvar toolInterruptID string\n\tfor _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif intCtx.IsRootCause {\n\t\t\ttoolInterruptID = intCtx.ID\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.NotEmpty(t, toolInterruptID, \"Should have a root cause interrupt ID\")\n\n\tt.Logf(\"Resuming with approval for interrupt ID: %s\", toolInterruptID)\n\n\tresumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\ttoolInterruptID: &approvalResult{Approved: true},\n\t\t},\n\t})\n\tassert.NoError(t, err, \"Resume should not error\")\n\tassert.NotNil(t, resumeIter, \"Resume iterator should not be nil\")\n\n\tvar resumeEvents []*adk.AgentEvent\n\tfor {\n\t\tevent, ok := resumeIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tresumeEvents = append(resumeEvents, event)\n\t}\n\n\tassert.NotEmpty(t, resumeEvents, \"Should have resume events after approval\")\n\n\tfor _, event := range resumeEvents {\n\t\tassert.NoError(t, event.Err, \"Resume event should not have error\")\n\t}\n\n\tvar hasToolResponse, hasTransferBack bool\n\tfor _, event := range resumeEvents {\n\t\tif event.Output != nil && event.Output.MessageOutput != nil {\n\t\t\tmsg := event.Output.MessageOutput.Message\n\t\t\tif msg != nil && msg.Role == \"tool\" && strings.Contains(msg.Content, \"executed successfully\") {\n\t\t\t\thasToolResponse = true\n\t\t\t}\n\t\t}\n\t\tif event.Action != nil && event.Action.TransferToAgent != nil {\n\t\t\tif event.Action.TransferToAgent.DestAgentName == \"company_coordinator\" {\n\t\t\t\thasTransferBack = true\n\t\t\t}\n\t\t}\n\t}\n\n\tassert.True(t, hasToolResponse, \"Should have tool response indicating successful payment processing\")\n\tassert.True(t, hasTransferBack, \"Should have transfer back to outer supervisor indicating completion\")\n}\n\nfunc TestSupervisorExit(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tsupervisorAgent := mockAdk.NewMockAgent(ctrl)\n\tsubAgent := mockAdk.NewMockAgent(ctrl)\n\n\tsupervisorAgent.EXPECT().Name(gomock.Any()).Return(\"Supervisor\").AnyTimes()\n\tsupervisorAgent.EXPECT().Description(gomock.Any()).Return(\"Supervisor description\").AnyTimes()\n\tsubAgent.EXPECT().Name(gomock.Any()).Return(\"SubAgent\").AnyTimes()\n\n\t// 1. Supervisor transfers to SubAgent\n\taMsg, tMsg := adk.GenTransferMessages(ctx, \"SubAgent\")\n\ti1, g1 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tg1.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, \"\"))\n\tevent1 := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)\n\tevent1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: \"SubAgent\"}}\n\tg1.Send(event1)\n\tg1.Close()\n\tsupervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i1).Times(1)\n\n\t// 2. SubAgent emits Exit action\n\ti2, g2 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\texitEvent := &adk.AgentEvent{\n\t\tAgentName: \"SubAgent\",\n\t\tAction:    &adk.AgentAction{Exit: true},\n\t\tOutput: &adk.AgentOutput{\n\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tMessage: schema.AssistantMessage(\"Exiting...\", nil),\n\t\t\t},\n\t\t},\n\t}\n\tg2.Send(exitEvent)\n\tg2.Close()\n\tsubAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i2).Times(1)\n\n\tconf := &Config{\n\t\tSupervisor: supervisorAgent,\n\t\tSubAgents:  []adk.Agent{subAgent},\n\t}\n\n\tmultiAgent, err := New(ctx, conf)\n\tassert.NoError(t, err)\n\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent})\n\taIter := runner.Run(ctx, []adk.Message{schema.UserMessage(\"test\")})\n\n\t// Collect events\n\tvar events []*adk.AgentEvent\n\tfor {\n\t\tevent, ok := aIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\tfoundExit := false\n\tfoundTransferBack := false\n\n\tfor _, e := range events {\n\t\tif e.Action != nil {\n\t\t\tif e.Action.Exit {\n\t\t\t\tfoundExit = true\n\t\t\t}\n\t\t\tif e.Action.TransferToAgent != nil && e.Action.TransferToAgent.DestAgentName == \"Supervisor\" {\n\t\t\t\tfoundTransferBack = true\n\t\t\t}\n\t\t}\n\t}\n\n\tassert.True(t, foundExit, \"Should have found Exit action\")\n\tassert.False(t, foundTransferBack, \"Should NOT have found Transfer back to Supervisor after Exit\")\n}\n\nfunc TestNestedSupervisorExit(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\ttopSupervisor := mockAdk.NewMockAgent(ctrl)\n\tmidSupervisor := mockAdk.NewMockAgent(ctrl)\n\tworker := mockAdk.NewMockAgent(ctrl)\n\n\ttopSupervisor.EXPECT().Name(gomock.Any()).Return(\"TopSupervisor\").AnyTimes()\n\ttopSupervisor.EXPECT().Description(gomock.Any()).Return(\"Top supervisor description\").AnyTimes()\n\tmidSupervisor.EXPECT().Name(gomock.Any()).Return(\"MidSupervisor\").AnyTimes()\n\tmidSupervisor.EXPECT().Description(gomock.Any()).Return(\"Mid supervisor description\").AnyTimes()\n\tworker.EXPECT().Name(gomock.Any()).Return(\"Worker\").AnyTimes()\n\n\t// 1. TopSupervisor transfers to MidSupervisor\n\taMsg1, tMsg1 := adk.GenTransferMessages(ctx, \"MidSupervisor\")\n\ti1, g1 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tg1.Send(adk.EventFromMessage(aMsg1, nil, schema.Assistant, \"\"))\n\tevent1 := adk.EventFromMessage(tMsg1, nil, schema.Tool, tMsg1.ToolName)\n\tevent1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: \"MidSupervisor\"}}\n\tg1.Send(event1)\n\tg1.Close()\n\ttopSupervisor.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i1).AnyTimes()\n\n\t// 2. MidSupervisor transfers to Worker\n\taMsg2, tMsg2 := adk.GenTransferMessages(ctx, \"Worker\")\n\ti2, g2 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\tg2.Send(adk.EventFromMessage(aMsg2, nil, schema.Assistant, \"\"))\n\tevent2 := adk.EventFromMessage(tMsg2, nil, schema.Tool, tMsg2.ToolName)\n\tevent2.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: \"Worker\"}}\n\tg2.Send(event2)\n\tg2.Close()\n\tmidSupervisor.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i2).AnyTimes()\n\n\t// 3. Worker emits Exit action\n\ti3, g3 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\texitEvent := &adk.AgentEvent{\n\t\tAgentName: \"Worker\",\n\t\tAction:    &adk.AgentAction{Exit: true},\n\t\tOutput: &adk.AgentOutput{\n\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tMessage: schema.AssistantMessage(\"Worker Exiting...\", nil),\n\t\t\t},\n\t\t},\n\t}\n\tg3.Send(exitEvent)\n\tg3.Close()\n\tworker.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i3).Times(1)\n\n\t// Build Nested System\n\t// Mid System: MidSupervisor -> [Worker]\n\tmidSystem, err := New(ctx, &Config{\n\t\tSupervisor: midSupervisor,\n\t\tSubAgents:  []adk.Agent{worker},\n\t})\n\tassert.NoError(t, err)\n\t// We need to give the midSystem the name \"MidSupervisor\" so TopSupervisor can find it\n\t// supervisor.New returns a ResumableAgent that delegates Name() to the supervisor agent.\n\t// So midSystem.Name() should already be \"MidSupervisor\" because midSupervisor.Name() is \"MidSupervisor\".\n\n\t// Top System: TopSupervisor -> [midSystem]\n\ttopSystem, err := New(ctx, &Config{\n\t\tSupervisor: topSupervisor,\n\t\tSubAgents:  []adk.Agent{midSystem},\n\t})\n\tassert.NoError(t, err)\n\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: topSystem})\n\taIter := runner.Run(ctx, []adk.Message{schema.UserMessage(\"test nested exit\")})\n\n\t// Collect events\n\tvar events []*adk.AgentEvent\n\tfor {\n\t\tevent, ok := aIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\tfoundExit := false\n\tfoundTransferBackToMidAfterExit := false\n\tfoundTransferBackToTopAfterExit := false\n\n\tfor _, e := range events {\n\t\tif e.Action != nil {\n\t\t\tif e.Action.Exit {\n\t\t\t\tfoundExit = true\n\t\t\t}\n\t\t\tif foundExit && e.Action.TransferToAgent != nil {\n\t\t\t\tif e.Action.TransferToAgent.DestAgentName == \"MidSupervisor\" {\n\t\t\t\t\tfoundTransferBackToMidAfterExit = true\n\t\t\t\t}\n\t\t\t\tif e.Action.TransferToAgent.DestAgentName == \"TopSupervisor\" {\n\t\t\t\t\tfoundTransferBackToTopAfterExit = true\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tassert.True(t, foundExit, \"Should have found Exit action\")\n\tassert.False(t, foundTransferBackToMidAfterExit, \"Should NOT have found Transfer back to MidSupervisor after Exit\")\n\tassert.False(t, foundTransferBackToTopAfterExit, \"Should NOT have found Transfer back to TopSupervisor after Exit\")\n}\n\nfunc TestChatModelAgentInternalEventsExit(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tsupervisorAgent := mockAdk.NewMockAgent(ctrl)\n\tworkerModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tinnerAgent := mockAdk.NewMockAgent(ctrl)\n\n\tsupervisorAgent.EXPECT().Name(gomock.Any()).Return(\"Supervisor\").AnyTimes()\n\tsupervisorAgent.EXPECT().Description(gomock.Any()).Return(\"Supervisor description\").AnyTimes()\n\tinnerAgent.EXPECT().Name(gomock.Any()).Return(\"InnerAgent\").AnyTimes()\n\tinnerAgent.EXPECT().Description(gomock.Any()).Return(\"Inner Agent Description\").AnyTimes()\n\n\t// 1. Supervisor transfers to Worker (only once, then exits when worker transfers back)\n\tsupervisorRunCount := 0\n\tsupervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\t\t\tsupervisorRunCount++\n\t\t\titer, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tif supervisorRunCount == 1 {\n\t\t\t\t\taMsg, tMsg := adk.GenTransferMessages(ctx, \"Worker\")\n\t\t\t\t\tgen.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, \"\"))\n\t\t\t\t\tevent1 := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)\n\t\t\t\t\tevent1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: \"Worker\"}}\n\t\t\t\t\tgen.Send(event1)\n\t\t\t\t} else {\n\t\t\t\t\texitEvent := &adk.AgentEvent{\n\t\t\t\t\t\tAgentName: \"Supervisor\",\n\t\t\t\t\t\tAction:    &adk.AgentAction{Exit: true},\n\t\t\t\t\t\tOutput: &adk.AgentOutput{\n\t\t\t\t\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\t\t\t\tMessage: schema.AssistantMessage(\"Supervisor done\", nil),\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t}\n\t\t\t\t\tgen.Send(exitEvent)\n\t\t\t\t}\n\t\t\t}()\n\t\t\treturn iter\n\t\t}).AnyTimes()\n\n\t// 2. Worker runs, calls AgentTool (InnerAgent)\n\t// Mock WorkerModel behavior\n\tworkerModel.EXPECT().WithTools(gomock.Any()).Return(workerModel, nil).AnyTimes()\n\n\t// 2.1 Worker generates tool call\n\ttoolCallMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"call_inner_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"InnerAgent\",\n\t\t\t\tArguments: `{\"request\": \"do exit\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tworkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(toolCallMsg, nil).Times(1)\n\n\t// 2.2 InnerAgent runs and emits Exit\n\tinnerAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\t\t\titer, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer gen.Close()\n\t\t\t\tinnerExitEvent := &adk.AgentEvent{\n\t\t\t\t\tAgentName: \"InnerAgent\",\n\t\t\t\t\tAction:    &adk.AgentAction{Exit: true},\n\t\t\t\t\tRunPath:   []adk.RunStep{},\n\t\t\t\t\tOutput: &adk.AgentOutput{\n\t\t\t\t\t\tMessageOutput: &adk.MessageVariant{\n\t\t\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(\"Inner Exiting...\", nil),\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}\n\t\t\t\tgen.Send(innerExitEvent)\n\t\t\t}()\n\t\t\treturn iter\n\t\t}).AnyTimes()\n\n\t// 2.3 Worker receives tool result (empty string or whatever AgentTool returns on exit/interrupt)\n\t// AgentTool implementation details: if Exit action is present, it returns whatever output is there.\n\t// The Exit action itself is passed as internal event.\n\n\t// 2.4 Worker generates final response\n\tfinalMsg := schema.AssistantMessage(\"Worker Finished\", nil)\n\tworkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(finalMsg, nil).AnyTimes()\n\n\t// Build Worker Agent\n\tagentTool := adk.NewAgentTool(ctx, innerAgent)\n\tworkerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"Worker\",\n\t\tDescription: \"Worker Agent\",\n\t\tModel:       workerModel,\n\t\tToolsConfig: adk.ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{agentTool},\n\t\t\t},\n\t\t\tEmitInternalEvents: true, // Key configuration\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\t// Build System\n\tsys, err := New(ctx, &Config{\n\t\tSupervisor: supervisorAgent,\n\t\tSubAgents:  []adk.Agent{workerAgent},\n\t})\n\tassert.NoError(t, err)\n\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: sys})\n\taIter := runner.Run(ctx, []adk.Message{schema.UserMessage(\"start\")})\n\n\t// Collect events\n\tvar events []*adk.AgentEvent\n\tfor {\n\t\tevent, ok := aIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\tfoundInnerExit := false\n\tfoundTransferBack := false\n\n\tfor _, e := range events {\n\t\t// Check for InnerAgent exit event (propagated as internal event)\n\t\tif e.AgentName == \"InnerAgent\" && e.Action != nil && e.Action.Exit {\n\t\t\tfoundInnerExit = true\n\t\t}\n\n\t\t// Check for transfer back to Supervisor\n\t\tif e.AgentName == \"Worker\" && e.Action != nil && e.Action.TransferToAgent != nil &&\n\t\t\te.Action.TransferToAgent.DestAgentName == \"Supervisor\" {\n\t\t\tfoundTransferBack = true\n\t\t}\n\t}\n\n\tassert.True(t, foundInnerExit, \"Should have captured InnerAgent Exit event\")\n\tassert.True(t, foundTransferBack, \"Should have found Transfer back to Supervisor (Worker should NOT be considered exited)\")\n}\n\nfunc TestSupervisorContainerUnifiedTracing(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tsupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tsubAgentModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tsupervisorModel.EXPECT().WithTools(gomock.Any()).Return(supervisorModel, nil).AnyTimes()\n\tsubAgentModel.EXPECT().WithTools(gomock.Any()).Return(subAgentModel, nil).AnyTimes()\n\n\ttransferMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"transfer_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"transfer_to_agent\",\n\t\t\t\tArguments: `{\"agent_name\":\"SubAgent\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tsupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(transferMsg, nil).Times(1)\n\n\tsubAgentResponse := schema.AssistantMessage(\"SubAgent response\", nil)\n\tsubAgentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(subAgentResponse, nil).Times(1)\n\n\tfinalResponse := schema.AssistantMessage(\"Final response\", nil)\n\tsupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(finalResponse, nil).Times(1)\n\n\tsupervisorAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"SupervisorAgent\",\n\t\tDescription: \"Supervisor agent\",\n\t\tInstruction: \"You are a supervisor\",\n\t\tModel:       supervisorModel,\n\t})\n\tassert.NoError(t, err)\n\n\tsubAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"SubAgent\",\n\t\tDescription: \"Sub agent\",\n\t\tInstruction: \"You are a sub agent\",\n\t\tModel:       subAgentModel,\n\t})\n\tassert.NoError(t, err)\n\n\tmultiAgent, err := New(ctx, &Config{\n\t\tSupervisor: supervisorAgent,\n\t\tSubAgents:  []adk.Agent{subAgent},\n\t})\n\tassert.NoError(t, err)\n\n\tassert.Equal(t, \"SupervisorAgent\", multiAgent.Name(ctx))\n\n\ttyper, ok := multiAgent.(components.Typer)\n\tassert.True(t, ok, \"Should implement components.Typer\")\n\tassert.Equal(t, \"Supervisor\", typer.GetType())\n\n\tvar mu sync.Mutex\n\tvar onStartCalls []string\n\tvar onEndCalls []string\n\n\thandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != adk.ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tonStartCalls = append(onStartCalls, info.Name+\":\"+info.Type)\n\t\t\tmu.Unlock()\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != adk.ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tonEndCalls = append(onEndCalls, info.Name+\":\"+info.Type)\n\t\t\tmu.Unlock()\n\t\t\tif agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent})\n\titer := runner.Query(ctx, \"hello\", adk.WithCallbacks(handler))\n\n\tfor {\n\t\t_, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tmu.Lock()\n\tdefer mu.Unlock()\n\n\tassert.NotEmpty(t, onStartCalls, \"Should have OnStart calls\")\n\tassert.Contains(t, onStartCalls, \"SupervisorAgent:Supervisor\", \"Should have supervisor container OnStart with type 'Supervisor'\")\n}\n\ntype traceContextKey struct{}\n\nfunc TestSupervisorContainerUnifiedTracingOnResume(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\tdefer ctrl.Finish()\n\n\tsupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tworkerModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\tsupervisorModel.EXPECT().WithTools(gomock.Any()).Return(supervisorModel, nil).AnyTimes()\n\tworkerModel.EXPECT().WithTools(gomock.Any()).Return(workerModel, nil).AnyTimes()\n\n\tpaymentTool := &approvableTool{name: \"process_payment\", t: t}\n\n\tworkerToolCallMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"call_payment_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"process_payment\",\n\t\t\t\tArguments: `{\"action\": \"process $1000 payment\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tworkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(workerToolCallMsg, nil).Times(1)\n\n\tworkerCompletionMsg := schema.AssistantMessage(\"Payment processed successfully\", nil)\n\tworkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(workerCompletionMsg, nil).AnyTimes()\n\n\tworkerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"Worker\",\n\t\tDescription: \"Worker agent\",\n\t\tInstruction: \"You are a worker\",\n\t\tModel:       workerModel,\n\t\tToolsConfig: adk.ToolsConfig{\n\t\t\tToolsNodeConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{paymentTool},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\ttransferMsg := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{\n\t\t\tID:   \"transfer_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"transfer_to_agent\",\n\t\t\t\tArguments: `{\"agent_name\":\"Worker\"}`,\n\t\t\t},\n\t\t},\n\t})\n\tsupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(transferMsg, nil).Times(1)\n\n\tfinalResponse := schema.AssistantMessage(\"Final response\", nil)\n\tsupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tReturn(finalResponse, nil).AnyTimes()\n\n\tsupervisorAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{\n\t\tName:        \"SupervisorAgent\",\n\t\tDescription: \"Supervisor agent\",\n\t\tInstruction: \"You are a supervisor\",\n\t\tModel:       supervisorModel,\n\t\tExit:        &adk.ExitTool{},\n\t})\n\tassert.NoError(t, err)\n\n\tmultiAgent, err := New(ctx, &Config{\n\t\tSupervisor: supervisorAgent,\n\t\tSubAgents:  []adk.Agent{workerAgent},\n\t})\n\tassert.NoError(t, err)\n\n\tstore := newCheckpointStore()\n\trunner := adk.NewRunner(ctx, adk.RunnerConfig{\n\t\tAgent:           multiAgent,\n\t\tCheckPointStore: store,\n\t})\n\n\tvar mu sync.Mutex\n\tvar runOnStartCalls []string\n\tvar resumeOnStartCalls []string\n\tvar resumeParentTraceIDs []string\n\n\trunHandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != adk.ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\trunOnStartCalls = append(runOnStartCalls, info.Name+\":\"+info.Type)\n\t\t\tmu.Unlock()\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != adk.ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\tcheckpointID := \"test-unified-tracing-resume\"\n\titer := runner.Run(ctx, []adk.Message{schema.UserMessage(\"Process payment\")}, adk.WithCallbacks(runHandler), adk.WithCheckPointID(checkpointID))\n\n\tvar interruptEvent *adk.AgentEvent\n\tfor {\n\t\tevent, ok := iter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tinterruptEvent = event\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassert.NotNil(t, interruptEvent, \"Should have interrupt event\")\n\n\tvar toolInterruptID string\n\tfor _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif intCtx.IsRootCause {\n\t\t\ttoolInterruptID = intCtx.ID\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.NotEmpty(t, toolInterruptID, \"Should have a root cause interrupt ID\")\n\n\tmu.Lock()\n\tt.Logf(\"Run OnStart calls: %v\", runOnStartCalls)\n\tassert.Contains(t, runOnStartCalls, \"SupervisorAgent:Supervisor\", \"Run should have supervisor container OnStart\")\n\tmu.Unlock()\n\n\tresumeHandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component != adk.ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tresumeOnStartCalls = append(resumeOnStartCalls, info.Name+\":\"+info.Type)\n\t\t\tparentID, _ := ctx.Value(traceContextKey{}).(string)\n\t\t\tresumeParentTraceIDs = append(resumeParentTraceIDs, info.Name+\":parent=\"+parentID)\n\t\t\tmu.Unlock()\n\t\t\tif info.Type == \"Supervisor\" {\n\t\t\t\treturn context.WithValue(ctx, traceContextKey{}, \"supervisor-trace-id\")\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tif info.Component != adk.ComponentOfAgent {\n\t\t\t\treturn ctx\n\t\t\t}\n\t\t\tif agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {\n\t\t\t\tgo func() {\n\t\t\t\t\tfor {\n\t\t\t\t\t\t_, ok := agentOutput.Events.Next()\n\t\t\t\t\t\tif !ok {\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\tresumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\ttoolInterruptID: &approvalResult{Approved: true},\n\t\t},\n\t}, adk.WithCallbacks(resumeHandler))\n\tassert.NoError(t, err)\n\n\tfor {\n\t\tevent, ok := resumeIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, event.Err)\n\t}\n\n\tmu.Lock()\n\tdefer mu.Unlock()\n\n\tt.Logf(\"Resume OnStart calls: %v\", resumeOnStartCalls)\n\tt.Logf(\"Resume parent trace IDs: %v\", resumeParentTraceIDs)\n\tassert.NotEmpty(t, resumeOnStartCalls, \"Should have OnStart calls during resume\")\n\tassert.Contains(t, resumeOnStartCalls, \"SupervisorAgent:Supervisor\", \"Resume should have supervisor container OnStart with type 'Supervisor'\")\n\n\tfoundInnerSupervisorWithParent := false\n\tfor _, entry := range resumeParentTraceIDs {\n\t\tif strings.Contains(entry, \"SupervisorAgent\") && !strings.Contains(entry, \"parent=supervisor-trace-id\") && entry != \"SupervisorAgent:parent=\" {\n\t\t\tif strings.Contains(resumeOnStartCalls[0], \"Supervisor\") {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t}\n\t\tif strings.Contains(entry, \"parent=supervisor-trace-id\") {\n\t\t\tfoundInnerSupervisorWithParent = true\n\t\t}\n\t}\n\tassert.True(t, foundInnerSupervisorWithParent,\n\t\t\"Inner agents should have parent trace from Supervisor container during Resume. Got: %v\", resumeParentTraceIDs)\n}\n"
  },
  {
    "path": "adk/react.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/gob\"\n\t\"errors\"\n\t\"io\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// ErrExceedMaxIterations indicates the agent reached the maximum iterations limit.\nvar ErrExceedMaxIterations = errors.New(\"exceeds max iterations\")\n\n// State holds agent runtime state including messages and user-extensible storage.\n//\n// Deprecated: This type will be unexported in v1.0.0. Use ChatModelAgentState\n// in HandlerMiddleware and AgentMiddleware callbacks instead. Direct use of\n// compose.ProcessState[*State] is discouraged and will stop working in v1.0.0;\n// use the handler APIs instead.\ntype State struct {\n\tMessages []Message\n\tExtra    map[string]any\n\n\t// Internal fields below - do not access directly.\n\t// Kept exported for backward compatibility with existing checkpoints.\n\tHasReturnDirectly        bool\n\tReturnDirectlyToolCallID string\n\tToolGenActions           map[string]*AgentAction\n\tAgentName                string\n\tRemainingIterations      int\n\tReturnDirectlyEvent      *AgentEvent\n\tRetryAttempt             int\n}\n\nconst (\n\tstateGobNameV07 = \"_eino_adk_react_state\"\n\n\t// stateGobNameV080 is a v0.8.0-v0.8.3-only alias used after byte-patching\n\t// raw checkpoint bytes in preprocessADKCheckpoint.\n\t// It must stay the same byte length as stateGobNameV07 so the length-prefixed\n\t// gob string in the stream remains valid.\n\tstateGobNameV080 = \"_eino_adk_state_v080_\"\n)\n\nfunc init() {\n\t// Checkpoint compatibility notes:\n\t// - ADK/compose checkpoints are gob-encoded and may store state behind `any`, so gob relies on\n\t//   an on-wire type name to choose a local Go type.\n\t// - Gob allows only one local Go type per name, and it treats \"struct wire\" and \"GobEncoder wire\"\n\t//   as incompatible even if the name matches.\n\t//\n\t// This file maintains 2 epochs of *State decoding:\n\t// - v0.7.* and current: \"_eino_adk_react_state\" + struct wire → decode into *State directly.\n\t//   State's exported fields are a superset of v0.7, so gob handles missing fields gracefully.\n\t// - v0.8.0-v0.8.3: \"_eino_adk_react_state\" + GobEncoder wire → byte-patched to stateGobNameV080,\n\t//   decode into stateV080 and migrate.\n\tschema.RegisterName[*State](stateGobNameV07)\n\tschema.RegisterName[*stateV080](stateGobNameV080)\n\n\t// the following two lines of registration mainly for backward compatibility\n\t// when decoding checkpoints created by v0.8.0 - v0.8.3\n\tgob.Register(&AgentEvent{})\n\tgob.Register(int(0))\n}\n\nfunc (s *State) getReturnDirectlyEvent() *AgentEvent {\n\treturn s.ReturnDirectlyEvent\n}\n\nfunc (s *State) setReturnDirectlyEvent(event *AgentEvent) {\n\ts.ReturnDirectlyEvent = event\n}\n\nfunc (s *State) getRetryAttempt() int {\n\treturn s.RetryAttempt\n}\n\nfunc (s *State) setRetryAttempt(attempt int) {\n\ts.RetryAttempt = attempt\n}\n\nfunc (s *State) getReturnDirectlyToolCallID() string {\n\treturn s.ReturnDirectlyToolCallID\n}\n\nfunc (s *State) setReturnDirectlyToolCallID(id string) {\n\ts.ReturnDirectlyToolCallID = id\n\ts.HasReturnDirectly = id != \"\"\n}\n\nfunc (s *State) getToolGenActions() map[string]*AgentAction {\n\treturn s.ToolGenActions\n}\n\nfunc (s *State) setToolGenAction(key string, action *AgentAction) {\n\tif s.ToolGenActions == nil {\n\t\ts.ToolGenActions = make(map[string]*AgentAction)\n\t}\n\ts.ToolGenActions[key] = action\n}\n\nfunc (s *State) popToolGenAction(key string) *AgentAction {\n\tif s.ToolGenActions == nil {\n\t\treturn nil\n\t}\n\taction := s.ToolGenActions[key]\n\tdelete(s.ToolGenActions, key)\n\treturn action\n}\n\nfunc (s *State) getRemainingIterations() int {\n\treturn s.RemainingIterations\n}\n\nfunc (s *State) setRemainingIterations(iterations int) {\n\ts.RemainingIterations = iterations\n}\n\nfunc (s *State) decrementRemainingIterations() {\n\tcurrent := s.getRemainingIterations()\n\ts.RemainingIterations = current - 1\n}\n\n// stateV080 handles the v0.8.0-v0.8.3 checkpoint format.\n// In those versions, *State implemented GobEncoder and was registered under\n// \"_eino_adk_react_state\". GobEncode serialized a stateSerialization struct\n// into opaque bytes. This type's GobDecode reads that format.\n// It is registered under \"_eino_adk_state_v080_\" — a same-length alias used\n// only after byte-patching the checkpoint data in preprocessADKCheckpoint.\ntype stateV080 struct {\n\tMessages                 []Message\n\tHasReturnDirectly        bool\n\tReturnDirectlyToolCallID string\n\tToolGenActions           map[string]*AgentAction\n\tAgentName                string\n\tRemainingIterations      int\n\tRetryAttempt             int\n\tReturnDirectlyEvent      *AgentEvent\n\tExtra                    map[string]any\n\tInternals                map[string]any\n}\n\n// stateV080Serialization is the on-wire format that v0.8.0-v0.8.3 GobEncode produced.\n// It is only used by stateV080.GobDecode to parse those legacy opaque bytes.\ntype stateV080Serialization stateV080\n\nfunc (sc *stateV080) GobDecode(b []byte) error {\n\tss := &stateV080Serialization{}\n\tif err := gob.NewDecoder(bytes.NewReader(b)).Decode(ss); err != nil {\n\t\treturn err\n\t}\n\tsc.Messages = ss.Messages\n\tsc.HasReturnDirectly = ss.HasReturnDirectly\n\tsc.ReturnDirectlyToolCallID = ss.ReturnDirectlyToolCallID\n\tsc.ToolGenActions = ss.ToolGenActions\n\tsc.AgentName = ss.AgentName\n\tsc.RemainingIterations = ss.RemainingIterations\n\tsc.Extra = ss.Extra\n\tsc.Internals = ss.Internals\n\treturn nil\n}\n\n// stateV080ToState converts a legacy *stateV080 (v0.8.0-v0.8.3) to a current *State.\nfunc stateV080ToState(sc *stateV080) *State {\n\ts := &State{\n\t\tMessages:                 sc.Messages,\n\t\tHasReturnDirectly:        sc.HasReturnDirectly,\n\t\tReturnDirectlyToolCallID: sc.ReturnDirectlyToolCallID,\n\t\tToolGenActions:           sc.ToolGenActions,\n\t\tAgentName:                sc.AgentName,\n\t\tRemainingIterations:      sc.RemainingIterations,\n\t\tExtra:                    sc.Extra,\n\t}\n\tif sc.ReturnDirectlyToolCallID != \"\" {\n\t\ts.setReturnDirectlyToolCallID(sc.ReturnDirectlyToolCallID)\n\t}\n\tif sc.Internals != nil && s.RetryAttempt == 0 {\n\t\tif v, ok := sc.Internals[\"_retryAttempt\"].(int); ok {\n\t\t\ts.RetryAttempt = v\n\t\t}\n\t}\n\tif sc.Internals != nil && s.ReturnDirectlyEvent == nil {\n\t\tif v, ok := sc.Internals[\"_returnDirectlyEvent\"].(*AgentEvent); ok {\n\t\t\ts.ReturnDirectlyEvent = v\n\t\t}\n\t}\n\treturn s\n}\n\n// SendToolGenAction attaches an AgentAction to the next tool event emitted for the\n// current tool execution.\n//\n// Where/when to use:\n//   - Invoke within a tool's Run (Invokable/Streamable) implementation to include\n//     an action alongside that tool's output event.\n//   - The action is scoped by the current tool call context: if a ToolCallID is\n//     available, it is used as the key to support concurrent calls of the same\n//     tool with different parameters; otherwise, the provided toolName is used.\n//   - The stored action is ephemeral and will be popped and attached to the tool\n//     event when the tool finishes (including streaming completion).\n//\n// Limitation:\n//   - This function is intended for use within ChatModelAgent runs only. It relies\n//     on ChatModelAgent's internal State to store and pop actions, which is not\n//     available in other agent types.\nfunc SendToolGenAction(ctx context.Context, toolName string, action *AgentAction) error {\n\tkey := toolName\n\ttoolCallID := compose.GetToolCallID(ctx)\n\tif len(toolCallID) > 0 {\n\t\tkey = toolCallID\n\t}\n\n\treturn compose.ProcessState(ctx, func(ctx context.Context, st *State) error {\n\t\tst.setToolGenAction(key, action)\n\t\treturn nil\n\t})\n}\n\ntype reactInput struct {\n\tmessages []Message\n}\n\ntype reactConfig struct {\n\t// model is the chat model used by the react graph.\n\t// Tools are configured via model.WithTools call option, not the WithTools method.\n\tmodel model.BaseChatModel\n\n\ttoolsConfig      *compose.ToolsNodeConfig\n\tmodelWrapperConf *modelWrapperConfig\n\n\ttoolsReturnDirectly map[string]bool\n\n\tagentName string\n\n\tmaxIterations int\n}\n\nfunc genToolInfos(ctx context.Context, config *compose.ToolsNodeConfig) ([]*schema.ToolInfo, error) {\n\ttoolInfos := make([]*schema.ToolInfo, 0, len(config.Tools))\n\tfor _, t := range config.Tools {\n\t\ttl, err := t.Info(ctx)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\ttoolInfos = append(toolInfos, tl)\n\t}\n\n\treturn toolInfos, nil\n}\n\ntype reactGraph = *compose.Graph[*reactInput, Message]\ntype sToolNodeOutput = *schema.StreamReader[[]Message]\ntype sGraphOutput = MessageStream\n\nfunc getReturnDirectlyToolCallID(ctx context.Context) (string, bool) {\n\tvar toolCallID string\n\thandler := func(_ context.Context, st *State) error {\n\t\ttoolCallID = st.getReturnDirectlyToolCallID()\n\t\treturn nil\n\t}\n\n\t_ = compose.ProcessState(ctx, handler)\n\n\treturn toolCallID, toolCallID != \"\"\n}\n\nfunc genReactState(config *reactConfig) func(ctx context.Context) *State {\n\treturn func(ctx context.Context) *State {\n\t\tst := &State{\n\t\t\tAgentName: config.agentName,\n\t\t}\n\t\tmaxIter := 20\n\t\tif config.maxIterations > 0 {\n\t\t\tmaxIter = config.maxIterations\n\t\t}\n\t\tst.setRemainingIterations(maxIter)\n\t\treturn st\n\t}\n}\n\nfunc newReact(ctx context.Context, config *reactConfig) (reactGraph, error) {\n\tconst (\n\t\tinitNode_  = \"Init\"\n\t\tchatModel_ = \"ChatModel\"\n\t\ttoolNode_  = \"ToolNode\"\n\t)\n\n\tg := compose.NewGraph[*reactInput, Message](compose.WithGenLocalState(genReactState(config)))\n\n\tinitLambda := func(ctx context.Context, input *reactInput) ([]Message, error) {\n\t\treturn input.messages, nil\n\t}\n\t_ = g.AddLambdaNode(initNode_, compose.InvokableLambda(initLambda), compose.WithNodeName(initNode_))\n\n\tvar wrappedModel model.BaseChatModel = config.model\n\tif config.modelWrapperConf != nil {\n\t\twrappedModel = buildModelWrappers(config.model, config.modelWrapperConf)\n\t}\n\n\ttoolsNode, err := compose.NewToolNode(ctx, config.toolsConfig)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tmodelPreHandle := func(ctx context.Context, input []Message, st *State) ([]Message, error) {\n\t\tif st.getRemainingIterations() <= 0 {\n\t\t\treturn nil, ErrExceedMaxIterations\n\t\t}\n\t\tst.decrementRemainingIterations()\n\t\treturn input, nil\n\t}\n\t_ = g.AddChatModelNode(chatModel_, wrappedModel,\n\t\tcompose.WithStatePreHandler(modelPreHandle), compose.WithNodeName(chatModel_))\n\n\ttoolPreHandle := func(ctx context.Context, _ Message, st *State) (Message, error) {\n\t\tinput := st.Messages[len(st.Messages)-1]\n\n\t\treturnDirectly := config.toolsReturnDirectly\n\t\tif execCtx := getChatModelAgentExecCtx(ctx); execCtx != nil && len(execCtx.runtimeReturnDirectly) > 0 {\n\t\t\treturnDirectly = execCtx.runtimeReturnDirectly\n\t\t}\n\n\t\tif len(returnDirectly) > 0 {\n\t\t\tfor i := range input.ToolCalls {\n\t\t\t\ttoolName := input.ToolCalls[i].Function.Name\n\t\t\t\tif _, ok := returnDirectly[toolName]; ok {\n\t\t\t\t\tst.setReturnDirectlyToolCallID(input.ToolCalls[i].ID)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\treturn input, nil\n\t}\n\n\ttoolPostHandle := func(ctx context.Context, out *schema.StreamReader[[]*schema.Message], st *State) (*schema.StreamReader[[]*schema.Message], error) {\n\t\tif event := st.getReturnDirectlyEvent(); event != nil {\n\t\t\tgetChatModelAgentExecCtx(ctx).send(event)\n\t\t\tst.setReturnDirectlyEvent(nil)\n\t\t}\n\t\treturn out, nil\n\t}\n\n\t_ = g.AddToolsNode(toolNode_, toolsNode,\n\t\tcompose.WithStatePreHandler(toolPreHandle),\n\t\tcompose.WithStreamStatePostHandler(toolPostHandle),\n\t\tcompose.WithNodeName(toolNode_))\n\n\t_ = g.AddEdge(compose.START, initNode_)\n\t_ = g.AddEdge(initNode_, chatModel_)\n\n\ttoolCallCheck := func(ctx context.Context, sMsg MessageStream) (string, error) {\n\t\tdefer sMsg.Close()\n\t\tfor {\n\t\t\tchunk, err_ := sMsg.Recv()\n\t\t\tif err_ != nil {\n\t\t\t\tif err_ == io.EOF {\n\t\t\t\t\treturn compose.END, nil\n\t\t\t\t}\n\n\t\t\t\treturn \"\", err_\n\t\t\t}\n\n\t\t\tif len(chunk.ToolCalls) > 0 {\n\t\t\t\treturn toolNode_, nil\n\t\t\t}\n\t\t}\n\t}\n\tbranch := compose.NewStreamGraphBranch(toolCallCheck, map[string]bool{compose.END: true, toolNode_: true})\n\t_ = g.AddBranch(chatModel_, branch)\n\n\tif len(config.toolsReturnDirectly) > 0 {\n\t\tconst (\n\t\t\ttoolNodeToEndConverter = \"ToolNodeToEndConverter\"\n\t\t)\n\n\t\tcvt := func(ctx context.Context, sToolCallMessages sToolNodeOutput) (sGraphOutput, error) {\n\t\t\tid, _ := getReturnDirectlyToolCallID(ctx)\n\n\t\t\treturn schema.StreamReaderWithConvert(sToolCallMessages,\n\t\t\t\tfunc(in []Message) (Message, error) {\n\n\t\t\t\t\tfor _, chunk := range in {\n\t\t\t\t\t\tif chunk != nil && chunk.ToolCallID == id {\n\t\t\t\t\t\t\treturn chunk, nil\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\n\t\t\t\t\treturn nil, schema.ErrNoValue\n\t\t\t\t}), nil\n\t\t}\n\n\t\t_ = g.AddLambdaNode(toolNodeToEndConverter, compose.TransformableLambda(cvt),\n\t\t\tcompose.WithNodeName(toolNodeToEndConverter))\n\t\t_ = g.AddEdge(toolNodeToEndConverter, compose.END)\n\n\t\tcheckReturnDirect := func(ctx context.Context,\n\t\t\tsToolCallMessages sToolNodeOutput) (string, error) {\n\n\t\t\t_, ok := getReturnDirectlyToolCallID(ctx)\n\n\t\t\tif ok {\n\t\t\t\treturn toolNodeToEndConverter, nil\n\t\t\t}\n\n\t\t\treturn chatModel_, nil\n\t\t}\n\n\t\tbranch = compose.NewStreamGraphBranch(checkReturnDirect,\n\t\t\tmap[string]bool{toolNodeToEndConverter: true, chatModel_: true})\n\t\t_ = g.AddBranch(toolNode_, branch)\n\t} else {\n\t\t_ = g.AddEdge(toolNode_, chatModel_)\n\t}\n\n\treturn g, nil\n}\n"
  },
  {
    "path": "adk/react_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/gob\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"math/rand\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype testModelWrapper struct {\n\tinner model.ToolCallingChatModel\n}\n\nfunc TestStateCompatConversions_V080(t *testing.T) {\n\tt.Run(\"stateV080GobDecodeAndToState\", func(t *testing.T) {\n\t\tss := &stateV080Serialization{\n\t\t\tReturnDirectlyToolCallID: \"tcid\",\n\t\t\tRemainingIterations:      2,\n\t\t\tInternals: map[string]any{\n\t\t\t\t\"_retryAttempt\":        9,\n\t\t\t\t\"_returnDirectlyEvent\": &AgentEvent{AgentName: \"agent\"},\n\t\t\t},\n\t\t}\n\n\t\tvar buf bytes.Buffer\n\t\tassert.NoError(t, gob.NewEncoder(&buf).Encode(ss))\n\n\t\tvar legacy stateV080\n\t\tassert.NoError(t, legacy.GobDecode(buf.Bytes()))\n\n\t\ts := stateV080ToState(&legacy)\n\t\tassert.Equal(t, \"tcid\", s.ReturnDirectlyToolCallID)\n\t\tassert.True(t, s.HasReturnDirectly)\n\t\tassert.Equal(t, 2, s.RemainingIterations)\n\t\tassert.Equal(t, 9, s.RetryAttempt)\n\t\tassert.NotNil(t, s.ReturnDirectlyEvent)\n\t\tassert.Equal(t, \"agent\", s.ReturnDirectlyEvent.AgentName)\n\t})\n}\n\nfunc TestStateGetToolGenActions(t *testing.T) {\n\tst := &State{\n\t\tToolGenActions: map[string]*AgentAction{\n\t\t\t\"k\": {},\n\t\t},\n\t}\n\tassert.NotNil(t, st.getToolGenActions())\n\tassert.Contains(t, st.getToolGenActions(), \"k\")\n}\n\nfunc (w *testModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\treturn (&stateModelWrapper{inner: w.inner, original: w.inner}).Generate(ctx, input, opts...)\n}\n\nfunc (w *testModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\treturn (&stateModelWrapper{inner: w.inner, original: w.inner}).Stream(ctx, input, opts...)\n}\n\nfunc (w *testModelWrapper) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\tnewInner, err := w.inner.WithTools(tools)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn &testModelWrapper{inner: newInner}, nil\n}\n\n// TestReact tests the newReact function with different scenarios\nfunc TestReact(t *testing.T) {\n\t// Basic test for newReact function\n\tt.Run(\"Invoke\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a fake tool for testing\n\t\tfakeTool := &fakeToolForTest{\n\t\t\ttarCount: 3,\n\t\t}\n\n\t\tinfo, err := fakeTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\ttimes := 0\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) (Message, error) {\n\t\t\t\ttimes++\n\t\t\t\tif times <= 2 {\n\t\t\t\t\treturn schema.AssistantMessage(\"hello test\",\n\t\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID: randStrForTest(),\n\t\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"123\"}`, randStrForTest()),\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil\n\t\t\t\t}\n\n\t\t\t\treturn schema.AssistantMessage(\"bye\", nil), nil\n\t\t\t}).AnyTimes()\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create a reactConfig\n\t\tconfig := &reactConfig{\n\t\t\tmodel: &testModelWrapper{inner: cm},\n\t\t\ttoolsConfig: &compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t\ttoolsReturnDirectly: map[string]bool{},\n\t\t}\n\n\t\tgraph, err := newReact(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, graph)\n\n\t\tcompiled, err := graph.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, compiled)\n\n\t\t// Test with a user message\n\t\tresult, err := compiled.Invoke(ctx, &reactInput{messages: []Message{\n\t\t\t{\n\t\t\t\tRole:    schema.User,\n\t\t\t\tContent: \"Use the test tool to say hello\",\n\t\t\t},\n\t\t}})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, result)\n\t})\n\n\t// Test with toolsReturnDirectly\n\tt.Run(\"ToolsReturnDirectly\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a fake tool for testing\n\t\tfakeTool := &fakeToolForTest{\n\t\t\ttarCount: 3,\n\t\t}\n\n\t\tinfo, err := fakeTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\ttimes := 0\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) (Message, error) {\n\t\t\t\ttimes++\n\t\t\t\tif times <= 2 {\n\t\t\t\t\treturn schema.AssistantMessage(\"hello test\",\n\t\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID: randStrForTest(),\n\t\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"123\"}`, randStrForTest()),\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil\n\t\t\t\t}\n\n\t\t\t\treturn schema.AssistantMessage(\"bye\", nil), nil\n\t\t\t}).AnyTimes()\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create a reactConfig with toolsReturnDirectly\n\t\tconfig := &reactConfig{\n\t\t\tmodel: &testModelWrapper{inner: cm},\n\t\t\ttoolsConfig: &compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t\ttoolsReturnDirectly: map[string]bool{info.Name: true},\n\t\t}\n\n\t\tgraph, err := newReact(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, graph)\n\n\t\tcompiled, err := graph.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, compiled)\n\n\t\t// Test with a user message when tool returns directly\n\t\tresult, err := compiled.Invoke(ctx, &reactInput{messages: []Message{\n\t\t\t{\n\t\t\t\tRole:    schema.User,\n\t\t\t\tContent: \"Use the test tool to say hello\",\n\t\t\t},\n\t\t}})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, result)\n\n\t\tassert.Equal(t, result.Role, schema.Tool)\n\t})\n\n\t// Test streaming functionality\n\tt.Run(\"Stream\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a fake tool for testing\n\t\tfakeTool := &fakeToolForTest{\n\t\t\ttarCount: 3,\n\t\t}\n\n\t\tfakeStreamTool := &fakeStreamToolForTest{\n\t\t\ttarCount: 3,\n\t\t}\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\ttimes := 0\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) (\n\t\t\t\tMessageStream, error) {\n\t\t\t\tsr, sw := schema.Pipe[Message](1)\n\t\t\t\tdefer sw.Close()\n\n\t\t\t\tinfo, _ := fakeTool.Info(ctx)\n\t\t\t\tstreamInfo, _ := fakeStreamTool.Info(ctx)\n\n\t\t\t\ttimes++\n\t\t\t\tif times <= 1 {\n\t\t\t\t\tsw.Send(schema.AssistantMessage(\"hello test\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStrForTest(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"tool\"}`, randStrForTest()),\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil)\n\t\t\t\t\treturn sr, nil\n\t\t\t\t} else if times == 2 {\n\t\t\t\t\tsw.Send(schema.AssistantMessage(\"hello stream\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStrForTest(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      streamInfo.Name,\n\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"stream tool\"}`, randStrForTest()),\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil)\n\t\t\t\t\treturn sr, nil\n\t\t\t\t}\n\n\t\t\t\tsw.Send(schema.AssistantMessage(\"bye\", nil), nil)\n\t\t\t\treturn sr, nil\n\t\t\t}).AnyTimes()\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create a reactConfig\n\t\tconfig := &reactConfig{\n\t\t\tmodel: &testModelWrapper{inner: cm},\n\t\t\ttoolsConfig: &compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool, fakeStreamTool},\n\t\t\t},\n\t\t\ttoolsReturnDirectly: map[string]bool{},\n\t\t}\n\n\t\tgraph, err := newReact(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, graph)\n\n\t\tcompiled, err := graph.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, compiled)\n\n\t\t// Test streaming with a user message\n\t\toutStream, err := compiled.Stream(ctx, &reactInput{messages: []Message{\n\t\t\t{\n\t\t\t\tRole:    schema.User,\n\t\t\t\tContent: \"Use the test tool to say hello\",\n\t\t\t},\n\t\t}})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, outStream)\n\n\t\tdefer outStream.Close()\n\n\t\tmsgs := make([]Message, 0)\n\t\tfor {\n\t\t\tmsg, err_ := outStream.Recv()\n\t\t\tif err_ != nil {\n\t\t\t\tif errors.Is(err_, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tt.Fatal(err_)\n\t\t\t}\n\n\t\t\tmsgs = append(msgs, msg)\n\t\t}\n\n\t\tassert.NotEmpty(t, msgs)\n\t})\n\n\t// Test streaming with toolsReturnDirectly\n\tt.Run(\"StreamWithToolsReturnDirectly\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a fake tool for testing\n\t\tfakeTool := &fakeToolForTest{\n\t\t\ttarCount: 3,\n\t\t}\n\n\t\tfakeStreamTool := &fakeStreamToolForTest{\n\t\t\ttarCount: 3,\n\t\t}\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\ttimes := 0\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) (\n\t\t\t\tMessageStream, error) {\n\t\t\t\tsr, sw := schema.Pipe[Message](1)\n\t\t\t\tdefer sw.Close()\n\n\t\t\t\tinfo, _ := fakeTool.Info(ctx)\n\t\t\t\tstreamInfo, _ := fakeStreamTool.Info(ctx)\n\n\t\t\t\ttimes++\n\t\t\t\tif times <= 1 {\n\t\t\t\t\tsw.Send(schema.AssistantMessage(\"hello test\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStrForTest(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"tool\"}`, randStrForTest()),\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil)\n\t\t\t\t\treturn sr, nil\n\t\t\t\t} else if times == 2 {\n\t\t\t\t\tsw.Send(schema.AssistantMessage(\"hello stream\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStrForTest(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      streamInfo.Name,\n\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"stream tool\"}`, randStrForTest()),\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil)\n\t\t\t\t\treturn sr, nil\n\t\t\t\t}\n\n\t\t\t\tsw.Send(schema.AssistantMessage(\"bye\", nil), nil)\n\t\t\t\treturn sr, nil\n\t\t\t}).AnyTimes()\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\tstreamInfo, err := fakeStreamTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Create a reactConfig with toolsReturnDirectly\n\t\tconfig := &reactConfig{\n\t\t\tmodel: &testModelWrapper{inner: cm},\n\t\t\ttoolsConfig: &compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool, fakeStreamTool},\n\t\t\t},\n\t\t\ttoolsReturnDirectly: map[string]bool{streamInfo.Name: true},\n\t\t}\n\n\t\tgraph, err := newReact(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, graph)\n\n\t\tcompiled, err := graph.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, compiled)\n\n\t\t// Reset times counter\n\t\ttimes = 0\n\n\t\t// Test streaming with a user message when tool returns directly\n\t\toutStream, err := compiled.Stream(ctx, &reactInput{messages: []Message{\n\t\t\t{\n\t\t\t\tRole:    schema.User,\n\t\t\t\tContent: \"Use the test tool to say hello\",\n\t\t\t},\n\t\t}})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, outStream)\n\n\t\tmsgs := make([]Message, 0)\n\t\tfor {\n\t\t\tmsg, err_ := outStream.Recv()\n\t\t\tif err_ != nil {\n\t\t\t\tif errors.Is(err_, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\n\t\t\tassert.Equal(t, msg.Role, schema.Tool)\n\n\t\t\tmsgs = append(msgs, msg)\n\t\t}\n\n\t\toutStream.Close()\n\n\t\tassert.NotEmpty(t, msgs)\n\t})\n\n\tt.Run(\"MaxIterations\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a fake tool for testing\n\t\tfakeTool := &fakeToolForTest{\n\t\t\ttarCount: 3,\n\t\t}\n\n\t\tinfo, err := fakeTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Create a mock chat model\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\t// Set up expectations for the mock model\n\t\ttimes := 0\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) (Message, error) {\n\t\t\t\ttimes++\n\t\t\t\tif times <= 5 {\n\t\t\t\t\treturn schema.AssistantMessage(\"hello test\",\n\t\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID: randStrForTest(),\n\t\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"123\"}`, randStrForTest()),\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil\n\t\t\t\t}\n\n\t\t\t\treturn schema.AssistantMessage(\"bye\", nil), nil\n\t\t\t}).AnyTimes()\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// don't exceed max iterations\n\t\tconfig := &reactConfig{\n\t\t\tmodel: &testModelWrapper{inner: cm},\n\t\t\ttoolsConfig: &compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t\ttoolsReturnDirectly: map[string]bool{},\n\t\t\tmaxIterations:       6,\n\t\t}\n\n\t\tgraph, err := newReact(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, graph)\n\n\t\tcompiled, err := graph.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, compiled)\n\n\t\t// Test with a user message\n\t\tresult, err := compiled.Invoke(ctx, &reactInput{messages: []Message{\n\t\t\t{\n\t\t\t\tRole:    schema.User,\n\t\t\t\tContent: \"Use the test tool to say hello\",\n\t\t\t},\n\t\t}})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, result.Content, \"bye\")\n\n\t\t// reset chat model times counter\n\t\ttimes = 0\n\t\t// exceed max iterations\n\t\tconfig = &reactConfig{\n\t\t\tmodel: &testModelWrapper{inner: cm},\n\t\t\ttoolsConfig: &compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t\ttoolsReturnDirectly: map[string]bool{},\n\t\t\tmaxIterations:       5,\n\t\t}\n\n\t\tgraph, err = newReact(ctx, config)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, graph)\n\n\t\tcompiled, err = graph.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, compiled)\n\n\t\t// Test with a user message\n\t\tresult, err = compiled.Invoke(ctx, &reactInput{messages: []Message{\n\t\t\t{\n\t\t\t\tRole:    schema.User,\n\t\t\t\tContent: \"Use the test tool to say hello\",\n\t\t\t},\n\t\t}})\n\t\tassert.Error(t, err)\n\t\tt.Logf(\"actual error: %v\", err.Error())\n\t\tassert.ErrorIs(t, err, ErrExceedMaxIterations)\n\n\t\tassert.Contains(t, err.Error(), ErrExceedMaxIterations.Error())\n\t})\n}\n\n// Helper types and functions for testing\n\ntype fakeStreamToolForTest struct {\n\ttarCount int\n\tcurCount int\n}\n\nfunc (t *fakeStreamToolForTest) StreamableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (\n\t*schema.StreamReader[string], error) {\n\tp := &fakeToolInputForTest{}\n\terr := sonic.UnmarshalString(argumentsInJSON, p)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tif t.curCount >= t.tarCount {\n\t\ts := schema.StreamReaderFromArray([]string{`{\"say\": \"bye\"}`})\n\t\treturn s, nil\n\t}\n\tt.curCount++\n\ts := schema.StreamReaderFromArray([]string{fmt.Sprintf(`{\"say\": \"hello %v\"}`, p.Name)})\n\treturn s, nil\n}\n\ntype fakeToolForTest struct {\n\ttarCount int\n\tcurCount int\n}\n\nfunc (t *fakeToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"test_tool\",\n\t\tDesc: \"test tool for unit testing\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"name\": {\n\t\t\t\t\tDesc:     \"user name for testing\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t\tType:     schema.String,\n\t\t\t\t},\n\t\t\t}),\n\t}, nil\n}\n\nfunc (t *fakeStreamToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"test_stream_tool\",\n\t\tDesc: \"test stream tool for unit testing\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"name\": {\n\t\t\t\t\tDesc:     \"user name for testing\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t\tType:     schema.String,\n\t\t\t\t},\n\t\t\t}),\n\t}, nil\n}\n\nfunc (t *fakeToolForTest) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\tp := &fakeToolInputForTest{}\n\terr := sonic.UnmarshalString(argumentsInJSON, p)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tif t.curCount >= t.tarCount {\n\t\treturn `{\"say\": \"bye\"}`, nil\n\t}\n\n\tt.curCount++\n\treturn fmt.Sprintf(`{\"say\": \"hello %v\"}`, p.Name), nil\n}\n\ntype fakeToolInputForTest struct {\n\tName string `json:\"name\"`\n}\n\nfunc randStrForTest() string {\n\tseeds := []rune(\"test seed\")\n\tb := make([]rune, 8)\n\tfor i := range b {\n\t\tb[i] = seeds[rand.Intn(len(seeds))]\n\t}\n\treturn string(b)\n}\n"
  },
  {
    "path": "adk/retry_chatmodel.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"log\"\n\t\"math/rand\"\n\t\"time\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nvar (\n\t// ErrExceedMaxRetries is returned when the maximum number of retries has been exceeded.\n\t// Use errors.Is to check if an error is due to max retries being exceeded:\n\t//\n\t//   if errors.Is(err, adk.ErrExceedMaxRetries) {\n\t//       // handle max retries exceeded\n\t//   }\n\t//\n\t// Use errors.As to extract the underlying RetryExhaustedError for the last error details:\n\t//\n\t//   var retryErr *adk.RetryExhaustedError\n\t//   if errors.As(err, &retryErr) {\n\t//       fmt.Printf(\"last error was: %v\\n\", retryErr.LastErr)\n\t//   }\n\tErrExceedMaxRetries = errors.New(\"exceeds max retries\")\n)\n\n// RetryExhaustedError is returned when all retry attempts have been exhausted.\n// It wraps the last error that occurred during retry attempts.\ntype RetryExhaustedError struct {\n\tLastErr      error\n\tTotalRetries int\n}\n\nfunc (e *RetryExhaustedError) Error() string {\n\tif e.LastErr != nil {\n\t\treturn fmt.Sprintf(\"exceeds max retries: last error: %v\", e.LastErr)\n\t}\n\treturn \"exceeds max retries\"\n}\n\nfunc (e *RetryExhaustedError) Unwrap() error {\n\treturn ErrExceedMaxRetries\n}\n\n// WillRetryError is emitted when a retryable error occurs and a retry will be attempted.\n// It allows end-users to observe retry events in real-time via AgentEvent.\n//\n// Field design rationale:\n//   - ErrStr (exported): Stores the error message string for Gob serialization during checkpointing.\n//     This ensures the error message is preserved after checkpoint restore.\n//   - err (unexported): Stores the original error for Unwrap() support at runtime.\n//     This field is intentionally unexported because Gob serialization would fail for unregistered\n//     concrete error types. Since end-users only need the original error when the AgentEvent first\n//     occurs (not after restoring from checkpoint), skipping serialization is acceptable.\n//     After checkpoint restore, err will be nil and Unwrap() returns nil.\ntype WillRetryError struct {\n\tErrStr       string\n\tRetryAttempt int\n\terr          error\n}\n\nfunc (e *WillRetryError) Error() string {\n\treturn e.ErrStr\n}\n\nfunc (e *WillRetryError) Unwrap() error {\n\treturn e.err\n}\n\nfunc init() {\n\tschema.RegisterName[*WillRetryError](\"eino_adk_chatmodel_will_retry_error\")\n}\n\n// ModelRetryConfig configures retry behavior for the ChatModel node.\n// It defines how the agent should handle transient failures when calling the ChatModel.\ntype ModelRetryConfig struct {\n\t// MaxRetries specifies the maximum number of retry attempts.\n\t// A value of 0 means no retries will be attempted.\n\t// A value of 3 means up to 3 retry attempts (4 total calls including the initial attempt).\n\tMaxRetries int\n\n\t// IsRetryAble is a function that determines whether an error should trigger a retry.\n\t// If nil, all errors are considered retry-able.\n\t// Return true if the error is transient and the operation should be retried.\n\t// Return false if the error is permanent and should be propagated immediately.\n\tIsRetryAble func(ctx context.Context, err error) bool\n\n\t// BackoffFunc calculates the delay before the next retry attempt.\n\t// The attempt parameter starts at 1 for the first retry.\n\t// If nil, a default exponential backoff with jitter is used:\n\t// base delay 100ms, exponentially increasing up to 10s max,\n\t// with random jitter (0-50% of delay) to prevent thundering herd.\n\tBackoffFunc func(ctx context.Context, attempt int) time.Duration\n}\n\nfunc defaultIsRetryAble(_ context.Context, err error) bool {\n\treturn err != nil\n}\n\nfunc defaultBackoff(_ context.Context, attempt int) time.Duration {\n\tbaseDelay := 100 * time.Millisecond\n\tmaxDelay := 10 * time.Second\n\n\tif attempt <= 0 {\n\t\treturn baseDelay\n\t}\n\n\tif attempt > 7 {\n\t\treturn maxDelay + time.Duration(rand.Int63n(int64(maxDelay/2)))\n\t}\n\n\tdelay := baseDelay * time.Duration(1<<uint(attempt-1))\n\tif delay > maxDelay {\n\t\tdelay = maxDelay\n\t}\n\n\tjitter := time.Duration(rand.Int63n(int64(delay / 2)))\n\treturn delay + jitter\n}\n\nfunc genErrWrapper(ctx context.Context, maxRetries, attempt int, isRetryAbleFunc func(ctx context.Context, err error) bool) func(error) error {\n\treturn func(err error) error {\n\t\tisRetryAble := isRetryAbleFunc == nil || isRetryAbleFunc(ctx, err)\n\t\thasRetriesLeft := attempt < maxRetries\n\n\t\tif isRetryAble && hasRetriesLeft {\n\t\t\treturn &WillRetryError{ErrStr: err.Error(), RetryAttempt: attempt, err: err}\n\t\t}\n\t\treturn err\n\t}\n}\n\nfunc consumeStreamForError(stream *schema.StreamReader[*schema.Message]) error {\n\tdefer stream.Close()\n\tfor {\n\t\t_, err := stream.Recv()\n\t\tif err == io.EOF {\n\t\t\treturn nil\n\t\t}\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n}\n\n// retryModelWrapper wraps a BaseChatModel with retry logic.\n// This is used inside the model wrapper chain, positioned between eventSenderModelWrapper\n// and stateModelWrapper, so that retry only affects the inner chain (event sending, user wrappers,\n// callback injection) without re-running state management (BeforeModelRewriteState/AfterModelRewriteState).\ntype retryModelWrapper struct {\n\tinner  model.BaseChatModel\n\tconfig *ModelRetryConfig\n}\n\nfunc newRetryModelWrapper(inner model.BaseChatModel, config *ModelRetryConfig) *retryModelWrapper {\n\treturn &retryModelWrapper{inner: inner, config: config}\n}\n\nfunc (r *retryModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tisRetryAble := r.config.IsRetryAble\n\tif isRetryAble == nil {\n\t\tisRetryAble = defaultIsRetryAble\n\t}\n\tbackoffFunc := r.config.BackoffFunc\n\tif backoffFunc == nil {\n\t\tbackoffFunc = defaultBackoff\n\t}\n\n\tvar lastErr error\n\tfor attempt := 0; attempt <= r.config.MaxRetries; attempt++ {\n\t\tout, err := r.inner.Generate(ctx, input, opts...)\n\t\tif err == nil {\n\t\t\treturn out, nil\n\t\t}\n\n\t\tif !isRetryAble(ctx, err) {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tlastErr = err\n\t\tif attempt < r.config.MaxRetries {\n\t\t\tlog.Printf(\"retrying ChatModel.Generate (attempt %d/%d): %v\", attempt+1, r.config.MaxRetries, err)\n\t\t\ttime.Sleep(backoffFunc(ctx, attempt+1))\n\t\t}\n\t}\n\n\treturn nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}\n}\n\nfunc (r *retryModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (\n\t*schema.StreamReader[*schema.Message], error) {\n\n\tisRetryAble := r.config.IsRetryAble\n\tif isRetryAble == nil {\n\t\tisRetryAble = defaultIsRetryAble\n\t}\n\tbackoffFunc := r.config.BackoffFunc\n\tif backoffFunc == nil {\n\t\tbackoffFunc = defaultBackoff\n\t}\n\n\tdefer func() {\n\t\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\t\tst.setRetryAttempt(0)\n\t\t\treturn nil\n\t\t})\n\t}()\n\n\tvar lastErr error\n\tfor attempt := 0; attempt <= r.config.MaxRetries; attempt++ {\n\t\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\t\tst.setRetryAttempt(attempt)\n\t\t\treturn nil\n\t\t})\n\n\t\tstream, err := r.inner.Stream(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\tif !isRetryAble(ctx, err) {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tlastErr = err\n\t\t\tif attempt < r.config.MaxRetries {\n\t\t\t\tlog.Printf(\"retrying ChatModel.Stream (attempt %d/%d): %v\", attempt+1, r.config.MaxRetries, err)\n\t\t\t\ttime.Sleep(backoffFunc(ctx, attempt+1))\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\n\t\tcopies := stream.Copy(2)\n\t\tcheckCopy := copies[0]\n\t\treturnCopy := copies[1]\n\n\t\tstreamErr := consumeStreamForError(checkCopy)\n\t\tif streamErr == nil {\n\t\t\treturn returnCopy, nil\n\t\t}\n\n\t\treturnCopy.Close()\n\t\tif !isRetryAble(ctx, streamErr) {\n\t\t\treturn nil, streamErr\n\t\t}\n\n\t\tlastErr = streamErr\n\t\tif attempt < r.config.MaxRetries {\n\t\t\tlog.Printf(\"retrying ChatModel.Stream (attempt %d/%d): %v\", attempt+1, r.config.MaxRetries, streamErr)\n\t\t\ttime.Sleep(backoffFunc(ctx, attempt+1))\n\t\t}\n\t}\n\n\treturn nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}\n}\n"
  },
  {
    "path": "adk/runctx.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/gob\"\n\t\"fmt\"\n\t\"sort\"\n\t\"sync\"\n\t\"time\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// runSession CheckpointSchema: persisted via serialization.RunCtx (gob).\ntype runSession struct {\n\tValues    map[string]any\n\tvaluesMtx *sync.Mutex\n\n\tEvents     []*agentEventWrapper\n\tLaneEvents *laneEvents\n\tmtx        sync.Mutex\n}\n\n// laneEvents CheckpointSchema: persisted via serialization.RunCtx (gob).\ntype laneEvents struct {\n\tEvents []*agentEventWrapper\n\tParent *laneEvents\n}\n\n// agentEventWrapper CheckpointSchema: persisted via serialization.RunCtx (gob).\ntype agentEventWrapper struct {\n\t*AgentEvent\n\tmu                  sync.Mutex\n\tconcatenatedMessage Message\n\t// TS is the timestamp (in nanoseconds) when this event was created.\n\t// It is primarily used by the laneEvents mechanism to order events\n\t// from different agents in a multi-agent flow.\n\tTS int64\n\t// StreamErr stores the error message if the MessageStream contained an error.\n\t// This field guards against multiple calls to getMessageFromWrappedEvent\n\t// when the stream has already been consumed and errored.\n\t// Normally when StreamErr happens, the Agent will return with the error,\n\t// unless retry is configured for the agent generating this stream, in which case\n\t// this StreamErr will be of type WillRetryError (indicating retry is pending).\n\tStreamErr error\n}\n\ntype otherAgentEventWrapperForEncode agentEventWrapper\n\nfunc (a *agentEventWrapper) GobEncode() ([]byte, error) {\n\tif a.concatenatedMessage != nil && a.Output != nil && a.Output.MessageOutput != nil && a.Output.MessageOutput.IsStreaming {\n\t\ta.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray([]Message{a.concatenatedMessage})\n\t}\n\n\tbuf := &bytes.Buffer{}\n\terr := gob.NewEncoder(buf).Encode((*otherAgentEventWrapperForEncode)(a))\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to gob encode agent event wrapper: %w\", err)\n\t}\n\treturn buf.Bytes(), nil\n}\n\nfunc (a *agentEventWrapper) GobDecode(b []byte) error {\n\treturn gob.NewDecoder(bytes.NewReader(b)).Decode((*otherAgentEventWrapperForEncode)(a))\n}\n\nfunc newRunSession() *runSession {\n\treturn &runSession{\n\t\tValues:    make(map[string]any),\n\t\tvaluesMtx: &sync.Mutex{},\n\t}\n}\n\n// GetSessionValues returns all session key-value pairs for the current run.\nfunc GetSessionValues(ctx context.Context) map[string]any {\n\tsession := getSession(ctx)\n\tif session == nil {\n\t\treturn map[string]any{}\n\t}\n\n\treturn session.getValues()\n}\n\n// AddSessionValue sets a single session key-value pair for the current run.\nfunc AddSessionValue(ctx context.Context, key string, value any) {\n\tsession := getSession(ctx)\n\tif session == nil {\n\t\treturn\n\t}\n\n\tsession.addValue(key, value)\n}\n\n// AddSessionValues sets multiple session key-value pairs for the current run.\nfunc AddSessionValues(ctx context.Context, kvs map[string]any) {\n\tsession := getSession(ctx)\n\tif session == nil {\n\t\treturn\n\t}\n\n\tsession.addValues(kvs)\n}\n\n// GetSessionValue retrieves a session value by key and reports whether it exists.\nfunc GetSessionValue(ctx context.Context, key string) (any, bool) {\n\tsession := getSession(ctx)\n\tif session == nil {\n\t\treturn nil, false\n\t}\n\n\treturn session.getValue(key)\n}\n\nfunc (rs *runSession) addEvent(event *AgentEvent) {\n\twrapper := &agentEventWrapper{AgentEvent: event, TS: time.Now().UnixNano()}\n\t// If LaneEvents is not nil, we are in a parallel lane.\n\t// Append to the lane's local event slice (lock-free).\n\tif rs.LaneEvents != nil {\n\t\trs.LaneEvents.Events = append(rs.LaneEvents.Events, wrapper)\n\t\treturn\n\t}\n\n\t// Otherwise, we are on the main path. Append to the shared Events slice (with lock).\n\trs.mtx.Lock()\n\trs.Events = append(rs.Events, wrapper)\n\trs.mtx.Unlock()\n}\n\nfunc (rs *runSession) getEvents() []*agentEventWrapper {\n\t// If there are no in-flight lane events, we can return the main slice directly.\n\tif rs.LaneEvents == nil {\n\t\trs.mtx.Lock()\n\t\tevents := rs.Events\n\t\trs.mtx.Unlock()\n\t\treturn events\n\t}\n\n\t// If there are in-flight events, we must construct the full view.\n\t// First, get the committed history from the main Events slice.\n\trs.mtx.Lock()\n\tcommittedEvents := make([]*agentEventWrapper, len(rs.Events))\n\tcopy(committedEvents, rs.Events)\n\trs.mtx.Unlock()\n\n\t// Then, assemble the in-flight events by traversing the linked list.\n\t// Reading the .Parent pointer is safe without a lock because the parent of a lane is immutable after creation.\n\tvar laneSlices [][]*agentEventWrapper\n\ttotalLaneSize := 0\n\tfor lane := rs.LaneEvents; lane != nil; lane = lane.Parent {\n\t\tif len(lane.Events) > 0 {\n\t\t\tlaneSlices = append(laneSlices, lane.Events)\n\t\t\ttotalLaneSize += len(lane.Events)\n\t\t}\n\t}\n\n\t// Combine committed and in-flight history.\n\tfinalEvents := make([]*agentEventWrapper, 0, len(committedEvents)+totalLaneSize)\n\tfinalEvents = append(finalEvents, committedEvents...)\n\tfor i := len(laneSlices) - 1; i >= 0; i-- {\n\t\tfinalEvents = append(finalEvents, laneSlices[i]...)\n\t}\n\n\treturn finalEvents\n}\n\nfunc (rs *runSession) getValues() map[string]any {\n\trs.valuesMtx.Lock()\n\tvalues := make(map[string]any, len(rs.Values))\n\tfor k, v := range rs.Values {\n\t\tvalues[k] = v\n\t}\n\trs.valuesMtx.Unlock()\n\n\treturn values\n}\n\nfunc (rs *runSession) addValue(key string, value any) {\n\trs.valuesMtx.Lock()\n\trs.Values[key] = value\n\trs.valuesMtx.Unlock()\n}\n\nfunc (rs *runSession) addValues(kvs map[string]any) {\n\trs.valuesMtx.Lock()\n\tfor k, v := range kvs {\n\t\trs.Values[k] = v\n\t}\n\trs.valuesMtx.Unlock()\n}\n\nfunc (rs *runSession) getValue(key string) (any, bool) {\n\trs.valuesMtx.Lock()\n\tvalue, ok := rs.Values[key]\n\trs.valuesMtx.Unlock()\n\n\treturn value, ok\n}\n\ntype runContext struct {\n\tRootInput *AgentInput\n\tRunPath   []RunStep\n\n\tSession *runSession\n}\n\nfunc (rc *runContext) isRoot() bool {\n\treturn len(rc.RunPath) == 1\n}\n\nfunc (rc *runContext) deepCopy() *runContext {\n\tcopied := &runContext{\n\t\tRootInput: rc.RootInput,\n\t\tRunPath:   make([]RunStep, len(rc.RunPath)),\n\t\tSession:   rc.Session,\n\t}\n\n\tcopy(copied.RunPath, rc.RunPath)\n\n\treturn copied\n}\n\ntype runCtxKey struct{}\n\nfunc getRunCtx(ctx context.Context) *runContext {\n\trunCtx, ok := ctx.Value(runCtxKey{}).(*runContext)\n\tif !ok {\n\t\treturn nil\n\t}\n\treturn runCtx\n}\n\nfunc setRunCtx(ctx context.Context, runCtx *runContext) context.Context {\n\treturn context.WithValue(ctx, runCtxKey{}, runCtx)\n}\n\nfunc initRunCtx(ctx context.Context, agentName string, input *AgentInput) (context.Context, *runContext) {\n\trunCtx := getRunCtx(ctx)\n\tif runCtx != nil {\n\t\trunCtx = runCtx.deepCopy()\n\t} else {\n\t\trunCtx = &runContext{Session: newRunSession()}\n\t}\n\n\trunCtx.RunPath = append(runCtx.RunPath, RunStep{agentName: agentName})\n\tif runCtx.isRoot() && input != nil {\n\t\trunCtx.RootInput = input\n\t}\n\n\treturn setRunCtx(ctx, runCtx), runCtx\n}\n\nfunc joinRunCtxs(parentCtx context.Context, childCtxs ...context.Context) {\n\tswitch len(childCtxs) {\n\tcase 0:\n\t\treturn\n\tcase 1:\n\t\t// Optimization for the common case of a single branch.\n\t\tnewEvents := unwindLaneEvents(childCtxs...)\n\t\tcommitEvents(parentCtx, newEvents)\n\t\treturn\n\t}\n\n\t// 1. Collect all new events from the leaf nodes of each context's lane.\n\tnewEvents := unwindLaneEvents(childCtxs...)\n\n\t// 2. Sort the collected events by their creation timestamp for chronological order.\n\tsort.Slice(newEvents, func(i, j int) bool {\n\t\treturn newEvents[i].TS < newEvents[j].TS\n\t})\n\n\t// 3. Commit the sorted events to the parent.\n\tcommitEvents(parentCtx, newEvents)\n}\n\n// commitEvents appends a slice of new events to the correct parent lane or main event log.\nfunc commitEvents(ctx context.Context, newEvents []*agentEventWrapper) {\n\trunCtx := getRunCtx(ctx)\n\tif runCtx == nil || runCtx.Session == nil {\n\t\t// Should not happen, but handle defensively.\n\t\treturn\n\t}\n\n\t// If the context we are committing to is itself a lane, append to its event slice.\n\tif runCtx.Session.LaneEvents != nil {\n\t\trunCtx.Session.LaneEvents.Events = append(runCtx.Session.LaneEvents.Events, newEvents...)\n\t} else {\n\t\t// Otherwise, commit to the main, shared Events slice with a lock.\n\t\trunCtx.Session.mtx.Lock()\n\t\trunCtx.Session.Events = append(runCtx.Session.Events, newEvents...)\n\t\trunCtx.Session.mtx.Unlock()\n\t}\n}\n\n// unwindLaneEvents traverses the LaneEvents of the given contexts and collects\n// all events from the leaf nodes.\nfunc unwindLaneEvents(ctxs ...context.Context) []*agentEventWrapper {\n\tvar allNewEvents []*agentEventWrapper\n\tfor _, ctx := range ctxs {\n\t\trunCtx := getRunCtx(ctx)\n\t\tif runCtx != nil && runCtx.Session != nil && runCtx.Session.LaneEvents != nil {\n\t\t\tallNewEvents = append(allNewEvents, runCtx.Session.LaneEvents.Events...)\n\t\t}\n\t}\n\treturn allNewEvents\n}\n\nfunc forkRunCtx(ctx context.Context) context.Context {\n\tparentRunCtx := getRunCtx(ctx)\n\tif parentRunCtx == nil || parentRunCtx.Session == nil {\n\t\t// Should not happen in a parallel workflow, but handle defensively.\n\t\treturn ctx\n\t}\n\n\t// Create a new session for the child lane by manually copying the parent's session fields.\n\t// This is crucial to ensure a new mutex is created and that the LaneEvents pointer is unique.\n\tchildSession := &runSession{\n\t\tEvents:    parentRunCtx.Session.Events, // Share the committed history\n\t\tValues:    parentRunCtx.Session.Values, // Share the values map\n\t\tvaluesMtx: parentRunCtx.Session.valuesMtx,\n\t}\n\n\t// Fork the lane events within the new session struct.\n\tchildSession.LaneEvents = &laneEvents{\n\t\tParent: parentRunCtx.Session.LaneEvents,\n\t\tEvents: make([]*agentEventWrapper, 0),\n\t}\n\n\t// Create a new runContext for the child lane, pointing to the new session.\n\tchildRunCtx := &runContext{\n\t\tRootInput: parentRunCtx.RootInput,\n\t\tRunPath:   make([]RunStep, len(parentRunCtx.RunPath)),\n\t\tSession:   childSession,\n\t}\n\tcopy(childRunCtx.RunPath, parentRunCtx.RunPath)\n\n\treturn setRunCtx(ctx, childRunCtx)\n}\n\n// updateRunPathOnly creates a new context with an updated RunPath, but does NOT modify the Address.\n// This is used by sequential workflows to accumulate execution history for LLM context,\n// without incorrectly chaining the static addresses of peer agents.\nfunc updateRunPathOnly(ctx context.Context, agentNames ...string) context.Context {\n\trunCtx := getRunCtx(ctx)\n\tif runCtx == nil {\n\t\t// This should not happen in a sequential workflow context, but handle defensively.\n\t\trunCtx = &runContext{Session: newRunSession()}\n\t} else {\n\t\trunCtx = runCtx.deepCopy()\n\t}\n\n\tfor _, agentName := range agentNames {\n\t\trunCtx.RunPath = append(runCtx.RunPath, RunStep{agentName: agentName})\n\t}\n\n\treturn setRunCtx(ctx, runCtx)\n}\n\n// ClearRunCtx clears the run context of the multi-agents. This is particularly useful\n// when a customized agent with a multi-agents inside it is set as a subagent of another\n// multi-agents. In such cases, it's not expected to pass the outside run context to the\n// inside multi-agents, so this function helps isolate the contexts properly.\nfunc ClearRunCtx(ctx context.Context) context.Context {\n\treturn context.WithValue(ctx, runCtxKey{}, nil)\n}\n\nfunc ctxWithNewRunCtx(ctx context.Context, input *AgentInput, sharedParentSession bool) context.Context {\n\tvar session *runSession\n\tif sharedParentSession {\n\t\tif parentSession := getSession(ctx); parentSession != nil {\n\t\t\tsession = &runSession{\n\t\t\t\tValues:    parentSession.Values,\n\t\t\t\tvaluesMtx: parentSession.valuesMtx,\n\t\t\t}\n\t\t}\n\t}\n\tif session == nil {\n\t\tsession = newRunSession()\n\t}\n\treturn setRunCtx(ctx, &runContext{Session: session, RootInput: input})\n}\n\nfunc getSession(ctx context.Context) *runSession {\n\trunCtx := getRunCtx(ctx)\n\tif runCtx != nil {\n\t\treturn runCtx.Session\n\t}\n\n\treturn nil\n}\n"
  },
  {
    "path": "adk/runctx_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestSessionValues(t *testing.T) {\n\t// Test Case 1: Basic AddSessionValues and GetSessionValues\n\tt.Run(\"BasicSessionValues\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a run session\n\t\tsession := newRunSession()\n\t\trunCtx := &runContext{Session: session}\n\t\tctx = setRunCtx(ctx, runCtx)\n\n\t\t// Add values to the session\n\t\tvalues := map[string]any{\n\t\t\t\"key1\": \"value1\",\n\t\t\t\"key2\": 42,\n\t\t\t\"key3\": true,\n\t\t}\n\t\tAddSessionValues(ctx, values)\n\n\t\t// Get all values from the session\n\t\tretrievedValues := GetSessionValues(ctx)\n\n\t\t// Verify the values were added correctly\n\t\tassert.Equal(t, \"value1\", retrievedValues[\"key1\"])\n\t\tassert.Equal(t, 42, retrievedValues[\"key2\"])\n\t\tassert.Equal(t, true, retrievedValues[\"key3\"])\n\t\tassert.Len(t, retrievedValues, 3)\n\t})\n\n\t// Test Case 2: AddSessionValues with empty context\n\tt.Run(\"AddSessionValuesEmptyContext\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Add values to a context without a run session\n\t\tvalues := map[string]any{\n\t\t\t\"key1\": \"value1\",\n\t\t}\n\t\tAddSessionValues(ctx, values)\n\n\t\t// Get values should return empty map\n\t\tretrievedValues := GetSessionValues(ctx)\n\t\tassert.Empty(t, retrievedValues)\n\t})\n\n\t// Test Case 3: GetSessionValues with empty context\n\tt.Run(\"GetSessionValuesEmptyContext\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Get values from a context without a run session\n\t\tretrievedValues := GetSessionValues(ctx)\n\t\tassert.Empty(t, retrievedValues)\n\t})\n\n\t// Test Case 4: AddSessionValues with nil values\n\tt.Run(\"AddSessionValuesNilValues\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a run session\n\t\tsession := newRunSession()\n\t\trunCtx := &runContext{Session: session}\n\t\tctx = setRunCtx(ctx, runCtx)\n\n\t\t// Add nil values map\n\t\tAddSessionValues(ctx, nil)\n\n\t\t// Get values should still be empty\n\t\tretrievedValues := GetSessionValues(ctx)\n\t\tassert.Empty(t, retrievedValues)\n\t})\n\n\t// Test Case 5: AddSessionValues with empty values\n\tt.Run(\"AddSessionValuesEmptyValues\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a run session\n\t\tsession := newRunSession()\n\t\trunCtx := &runContext{Session: session}\n\t\tctx = setRunCtx(ctx, runCtx)\n\n\t\t// Add empty values map\n\t\tAddSessionValues(ctx, map[string]any{})\n\n\t\t// Get values should be empty\n\t\tretrievedValues := GetSessionValues(ctx)\n\t\tassert.Empty(t, retrievedValues)\n\t})\n\n\t// Test Case 6: AddSessionValues with complex data types\n\tt.Run(\"AddSessionValuesComplexTypes\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a run session\n\t\tsession := newRunSession()\n\t\trunCtx := &runContext{Session: session}\n\t\tctx = setRunCtx(ctx, runCtx)\n\n\t\t// Add complex values to the session\n\t\tvalues := map[string]any{\n\t\t\t\"string\": \"hello world\",\n\t\t\t\"int\":    123,\n\t\t\t\"float\":  45.67,\n\t\t\t\"bool\":   true,\n\t\t\t\"slice\":  []string{\"a\", \"b\", \"c\"},\n\t\t\t\"map\":    map[string]int{\"x\": 1, \"y\": 2},\n\t\t\t\"struct\": struct{ Name string }{Name: \"test\"},\n\t\t}\n\t\tAddSessionValues(ctx, values)\n\n\t\t// Get all values from the session\n\t\tretrievedValues := GetSessionValues(ctx)\n\n\t\t// Verify the complex values were added correctly\n\t\tassert.Equal(t, \"hello world\", retrievedValues[\"string\"])\n\t\tassert.Equal(t, 123, retrievedValues[\"int\"])\n\t\tassert.Equal(t, 45.67, retrievedValues[\"float\"])\n\t\tassert.Equal(t, true, retrievedValues[\"bool\"])\n\t\tassert.Equal(t, []string{\"a\", \"b\", \"c\"}, retrievedValues[\"slice\"])\n\t\tassert.Equal(t, map[string]int{\"x\": 1, \"y\": 2}, retrievedValues[\"map\"])\n\t\tassert.Equal(t, struct{ Name string }{Name: \"test\"}, retrievedValues[\"struct\"])\n\t\tassert.Len(t, retrievedValues, 7)\n\t})\n\n\t// Test Case 7: AddSessionValues overwrites existing values\n\tt.Run(\"AddSessionValuesOverwrite\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a run session\n\t\tsession := newRunSession()\n\t\trunCtx := &runContext{Session: session}\n\t\tctx = setRunCtx(ctx, runCtx)\n\n\t\t// Add initial values\n\t\tinitialValues := map[string]any{\n\t\t\t\"key1\": \"initial1\",\n\t\t\t\"key2\": \"initial2\",\n\t\t}\n\t\tAddSessionValues(ctx, initialValues)\n\n\t\t// Add values that overwrite some keys\n\t\toverwriteValues := map[string]any{\n\t\t\t\"key1\": \"overwritten1\",\n\t\t\t\"key3\": \"new3\",\n\t\t}\n\t\tAddSessionValues(ctx, overwriteValues)\n\n\t\t// Get all values from the session\n\t\tretrievedValues := GetSessionValues(ctx)\n\n\t\t// Verify the values were overwritten correctly\n\t\tassert.Equal(t, \"overwritten1\", retrievedValues[\"key1\"]) // overwritten\n\t\tassert.Equal(t, \"initial2\", retrievedValues[\"key2\"])     // unchanged\n\t\tassert.Equal(t, \"new3\", retrievedValues[\"key3\"])         // new\n\t\tassert.Len(t, retrievedValues, 3)\n\t})\n\n\t// Test Case 8: Concurrent access to session values\n\tt.Run(\"ConcurrentSessionValues\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a run session\n\t\tsession := newRunSession()\n\t\trunCtx := &runContext{Session: session}\n\t\tctx = setRunCtx(ctx, runCtx)\n\n\t\t// Add initial values\n\t\tinitialValues := map[string]any{\n\t\t\t\"counter\": 0,\n\t\t}\n\t\tAddSessionValues(ctx, initialValues)\n\n\t\t// Simulate concurrent access\n\t\tdone := make(chan bool)\n\n\t\t// Goroutine 1: Add values\n\t\tgo func() {\n\t\t\tfor i := 0; i < 100; i++ {\n\t\t\t\tvalues := map[string]any{\n\t\t\t\t\t\"goroutine1\": i,\n\t\t\t\t}\n\t\t\t\tAddSessionValues(ctx, values)\n\t\t\t}\n\t\t\tdone <- true\n\t\t}()\n\n\t\t// Goroutine 2: Add different values\n\t\tgo func() {\n\t\t\tfor i := 0; i < 100; i++ {\n\t\t\t\tvalues := map[string]any{\n\t\t\t\t\t\"goroutine2\": i,\n\t\t\t\t}\n\t\t\t\tAddSessionValues(ctx, values)\n\t\t\t}\n\t\t\tdone <- true\n\t\t}()\n\n\t\t// Wait for both goroutines to complete\n\t\t<-done\n\t\t<-done\n\n\t\t// Verify that both values were set (last write wins)\n\t\tretrievedValues := GetSessionValues(ctx)\n\t\tassert.Equal(t, 0, retrievedValues[\"counter\"])\n\t\tassert.Equal(t, 99, retrievedValues[\"goroutine1\"])\n\t\tassert.Equal(t, 99, retrievedValues[\"goroutine2\"])\n\t})\n\n\t// Test Case 9: GetSessionValue individual value\n\tt.Run(\"GetSessionValueIndividual\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a run session\n\t\tsession := newRunSession()\n\t\trunCtx := &runContext{Session: session}\n\t\tctx = setRunCtx(ctx, runCtx)\n\n\t\t// Add values to the session\n\t\tvalues := map[string]any{\n\t\t\t\"key1\": \"value1\",\n\t\t\t\"key2\": 42,\n\t\t}\n\t\tAddSessionValues(ctx, values)\n\n\t\t// Get individual values\n\t\tvalue1, exists1 := GetSessionValue(ctx, \"key1\")\n\t\tvalue2, exists2 := GetSessionValue(ctx, \"key2\")\n\t\tvalue3, exists3 := GetSessionValue(ctx, \"nonexistent\")\n\n\t\t// Verify individual values\n\t\tassert.True(t, exists1)\n\t\tassert.Equal(t, \"value1\", value1)\n\n\t\tassert.True(t, exists2)\n\t\tassert.Equal(t, 42, value2)\n\n\t\tassert.False(t, exists3)\n\t\tassert.Nil(t, value3)\n\t})\n\n\t// Test Case 10: AddSessionValue individual value\n\tt.Run(\"AddSessionValueIndividual\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a run session\n\t\tsession := newRunSession()\n\t\trunCtx := &runContext{Session: session}\n\t\tctx = setRunCtx(ctx, runCtx)\n\n\t\t// Add individual values\n\t\tAddSessionValue(ctx, \"key1\", \"value1\")\n\t\tAddSessionValue(ctx, \"key2\", 42)\n\n\t\t// Get all values\n\t\tretrievedValues := GetSessionValues(ctx)\n\n\t\t// Verify the values were added correctly\n\t\tassert.Equal(t, \"value1\", retrievedValues[\"key1\"])\n\t\tassert.Equal(t, 42, retrievedValues[\"key2\"])\n\t\tassert.Len(t, retrievedValues, 2)\n\t})\n\n\t// Test Case 11: AddSessionValue with empty context\n\tt.Run(\"AddSessionValueEmptyContext\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Add individual value to a context without a run session\n\t\tAddSessionValue(ctx, \"key1\", \"value1\")\n\n\t\t// Get individual value should return false\n\t\tvalue, exists := GetSessionValue(ctx, \"key1\")\n\t\tassert.False(t, exists)\n\t\tassert.Nil(t, value)\n\n\t\t// Get all values should return empty map\n\t\tretrievedValues := GetSessionValues(ctx)\n\t\tassert.Empty(t, retrievedValues)\n\t})\n\n\t// Test Case 12: Integration with run context initialization\n\tt.Run(\"IntegrationWithRunContext\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Initialize a run context with an agent\n\t\tinput := &AgentInput{\n\t\t\tMessages: []Message{\n\t\t\t\tschema.UserMessage(\"test input\"),\n\t\t\t},\n\t\t}\n\t\tctx, runCtx := initRunCtx(ctx, \"test-agent\", input)\n\n\t\t// Verify the run context was created\n\t\tassert.NotNil(t, runCtx)\n\t\tassert.NotNil(t, runCtx.Session)\n\n\t\t// Add values to the session\n\t\tvalues := map[string]any{\n\t\t\t\"integration_key\": \"integration_value\",\n\t\t}\n\t\tAddSessionValues(ctx, values)\n\n\t\t// Get values from the session\n\t\tretrievedValues := GetSessionValues(ctx)\n\t\tassert.Equal(t, \"integration_value\", retrievedValues[\"integration_key\"])\n\n\t\t// Verify the run path was set correctly\n\t\tassert.Len(t, runCtx.RunPath, 1)\n\t\tassert.Equal(t, \"test-agent\", runCtx.RunPath[0].agentName)\n\t})\n}\n\nfunc TestForkJoinRunCtx(t *testing.T) {\n\t// Helper to create a named event\n\tnewEvent := func(name string) *AgentEvent {\n\t\t// Add a small sleep to ensure timestamps are distinct\n\t\ttime.Sleep(1 * time.Millisecond)\n\t\treturn &AgentEvent{AgentName: name}\n\t}\n\n\t// Helper to get event names from a slice of wrappers\n\tgetEventNames := func(wrappers []*agentEventWrapper) []string {\n\t\tnames := make([]string, len(wrappers))\n\t\tfor i, w := range wrappers {\n\t\t\tnames[i] = w.AgentName\n\t\t}\n\t\treturn names\n\t}\n\n\t// 1. Setup: Create an initial runContext for the main execution path.\n\tmainCtx, mainRunCtx := initRunCtx(context.Background(), \"Main\", nil)\n\n\t// 2. Run Agent A\n\teventA := newEvent(\"A\")\n\tmainRunCtx.Session.addEvent(eventA)\n\tassert.Equal(t, []string{\"A\"}, getEventNames(mainRunCtx.Session.getEvents()), \"After A\")\n\n\t// 3. Fork for Par(B, C)\n\tctxB := forkRunCtx(mainCtx)\n\tctxC := forkRunCtx(mainCtx)\n\n\t// Assertions for Fork\n\trunCtxB := getRunCtx(ctxB)\n\trunCtxC := getRunCtx(ctxC)\n\tassert.NotSame(t, mainRunCtx.Session, runCtxB.Session, \"Session B should be a new struct\")\n\tassert.NotSame(t, mainRunCtx.Session, runCtxC.Session, \"Session C should be a new struct\")\n\tassert.NotSame(t, runCtxB.Session, runCtxC.Session, \"Sessions B and C should be different\")\n\tassert.Nil(t, mainRunCtx.Session.LaneEvents, \"Main session should have no lane events yet\")\n\tassert.NotNil(t, runCtxB.Session.LaneEvents, \"Session B should have lane events\")\n\tassert.NotNil(t, runCtxC.Session.LaneEvents, \"Session C should have lane events\")\n\tassert.Nil(t, runCtxB.Session.LaneEvents.Parent, \"Lane B's parent should be the main (nil) lane\")\n\tassert.Nil(t, runCtxC.Session.LaneEvents.Parent, \"Lane C's parent should be the main (nil) lane\")\n\n\t// 4. Run Agent B\n\teventB := newEvent(\"B\")\n\trunCtxB.Session.addEvent(eventB)\n\tassert.Equal(t, []string{\"A\", \"B\"}, getEventNames(runCtxB.Session.getEvents()), \"After B\")\n\n\t// 5. Run Agent C (and Nested Fork for Par(D, E))\n\teventC1 := newEvent(\"C1\")\n\trunCtxC.Session.addEvent(eventC1)\n\tassert.Equal(t, []string{\"A\", \"C1\"}, getEventNames(runCtxC.Session.getEvents()), \"After C1\")\n\n\tctxD := forkRunCtx(ctxC)\n\tctxE := forkRunCtx(ctxC)\n\n\t// Assertions for Nested Fork\n\trunCtxD := getRunCtx(ctxD)\n\trunCtxE := getRunCtx(ctxE)\n\tassert.NotNil(t, runCtxD.Session.LaneEvents.Parent, \"Lane D's parent should be Lane C\")\n\tassert.Same(t, runCtxC.Session.LaneEvents, runCtxD.Session.LaneEvents.Parent, \"Lane D's parent must be Lane C's node\")\n\tassert.Same(t, runCtxC.Session.LaneEvents, runCtxE.Session.LaneEvents.Parent, \"Lane E's parent must be Lane C's node\")\n\n\t// 6. Run Agents D and E\n\teventD := newEvent(\"D\")\n\trunCtxD.Session.addEvent(eventD)\n\teventE := newEvent(\"E\")\n\trunCtxE.Session.addEvent(eventE)\n\n\tassert.Equal(t, []string{\"A\", \"C1\", \"D\"}, getEventNames(runCtxD.Session.getEvents()), \"After D\")\n\tassert.Equal(t, []string{\"A\", \"C1\", \"E\"}, getEventNames(runCtxE.Session.getEvents()), \"After E\")\n\n\t// 7. Join Par(D, E)\n\tjoinRunCtxs(ctxC, ctxD, ctxE)\n\n\t// Assertions for Nested Join\n\t// The events should now be committed to Lane C's event slice.\n\tassert.Equal(t, []string{\"A\", \"C1\", \"D\", \"E\"}, getEventNames(runCtxC.Session.getEvents()), \"After joining D and E\")\n\n\t// 8. Join Par(B, C)\n\tjoinRunCtxs(mainCtx, ctxB, ctxC)\n\n\t// Assertions for Top-Level Join\n\t// The events should now be committed to the main session's Events slice.\n\tassert.Equal(t, []string{\"A\", \"B\", \"C1\", \"D\", \"E\"}, getEventNames(mainRunCtx.Session.getEvents()), \"After joining B and C\")\n\n\t// 9. Run Agent F\n\teventF := newEvent(\"F\")\n\tmainRunCtx.Session.addEvent(eventF)\n\tassert.Equal(t, []string{\"A\", \"B\", \"C1\", \"D\", \"E\", \"F\"}, getEventNames(mainRunCtx.Session.getEvents()), \"After F\")\n}\n"
  },
  {
    "path": "adk/runner.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"runtime/debug\"\n\t\"sync\"\n\n\t\"github.com/cloudwego/eino/internal/core\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Runner is the primary entry point for executing an Agent.\n// It manages the agent's lifecycle, including starting, resuming, and checkpointing.\ntype Runner struct {\n\t// a is the agent to be executed.\n\ta Agent\n\t// enableStreaming dictates whether the execution should be in streaming mode.\n\tenableStreaming bool\n\t// store is the checkpoint store used to persist agent state upon interruption.\n\t// If nil, checkpointing is disabled.\n\tstore CheckPointStore\n}\n\ntype CheckPointStore = core.CheckPointStore\n\ntype RunnerConfig struct {\n\tAgent           Agent\n\tEnableStreaming bool\n\n\tCheckPointStore CheckPointStore\n}\n\n// ResumeParams contains all parameters needed to resume an execution.\n// This struct provides an extensible way to pass resume parameters without\n// requiring breaking changes to method signatures.\ntype ResumeParams struct {\n\t// Targets contains the addresses of components to be resumed as keys,\n\t// with their corresponding resume data as values\n\tTargets map[string]any\n\t// Future extensible fields can be added here without breaking changes\n}\n\n// NewRunner creates a Runner that executes an Agent with optional streaming\n// and checkpoint persistence.\nfunc NewRunner(_ context.Context, conf RunnerConfig) *Runner {\n\treturn &Runner{\n\t\tenableStreaming: conf.EnableStreaming,\n\t\ta:               conf.Agent,\n\t\tstore:           conf.CheckPointStore,\n\t}\n}\n\n// Run starts a new execution of the agent with a given set of messages.\n// It returns an iterator that yields agent events as they occur.\n// If the Runner was configured with a CheckPointStore, it will automatically save the agent's state\n// upon interruption.\nfunc (r *Runner) Run(ctx context.Context, messages []Message,\n\topts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\to := getCommonOptions(nil, opts...)\n\n\tfa := toFlowAgent(ctx, r.a)\n\n\tinput := &AgentInput{\n\t\tMessages:        messages,\n\t\tEnableStreaming: r.enableStreaming,\n\t}\n\n\tctx = ctxWithNewRunCtx(ctx, input, o.sharedParentSession)\n\n\tAddSessionValues(ctx, o.sessionValues)\n\n\titer := fa.Run(ctx, input, opts...)\n\tif r.store == nil {\n\t\treturn iter\n\t}\n\n\tniter, gen := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo r.handleIter(ctx, iter, gen, o.checkPointID)\n\treturn niter\n}\n\n// Query is a convenience method that starts a new execution with a single user query string.\nfunc (r *Runner) Query(ctx context.Context,\n\tquery string, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\n\treturn r.Run(ctx, []Message{schema.UserMessage(query)}, opts...)\n}\n\n// Resume continues an interrupted execution from a checkpoint, using an \"Implicit Resume All\" strategy.\n// This method is best for simpler use cases where the act of resuming implies that all previously\n// interrupted points should proceed without specific data.\n//\n// When using this method, all interrupted agents will receive `isResumeFlow = false` when they\n// call `GetResumeContext`, as no specific agent was targeted. This is suitable for the \"Simple Confirmation\"\n// pattern where an agent only needs to know `wasInterrupted` is true to continue.\nfunc (r *Runner) Resume(ctx context.Context, checkPointID string, opts ...AgentRunOption) (\n\t*AsyncIterator[*AgentEvent], error) {\n\treturn r.resume(ctx, checkPointID, nil, opts...)\n}\n\n// ResumeWithParams continues an interrupted execution from a checkpoint with specific parameters.\n// This is the most common and powerful way to resume, allowing you to target specific interrupt points\n// (identified by their address/ID) and provide them with data.\n//\n// The params.Targets map should contain the addresses of the components to be resumed as keys. These addresses\n// can point to any interruptible component in the entire execution graph, including ADK agents, compose\n// graph nodes, or tools. The value can be the resume data for that component, or `nil` if no data is needed.\n//\n// When using this method:\n//   - Components whose addresses are in the params.Targets map will receive `isResumeFlow = true` when they\n//     call `GetResumeContext`.\n//   - Interrupted components whose addresses are NOT in the params.Targets map must decide how to proceed:\n//     -- \"Leaf\" components (the actual root causes of the original interrupt) MUST re-interrupt themselves\n//     to preserve their state.\n//     -- \"Composite\" agents (like SequentialAgent or ChatModelAgent) should generally proceed with their\n//     execution. They act as conduits, allowing the resume signal to flow to their children. They will\n//     naturally re-interrupt if one of their interrupted children re-interrupts, as they receive the\n//     new `CompositeInterrupt` signal from them.\nfunc (r *Runner) ResumeWithParams(ctx context.Context, checkPointID string, params *ResumeParams, opts ...AgentRunOption) (*AsyncIterator[*AgentEvent], error) {\n\treturn r.resume(ctx, checkPointID, params.Targets, opts...)\n}\n\n// resume is the internal implementation for both Resume and ResumeWithParams.\nfunc (r *Runner) resume(ctx context.Context, checkPointID string, resumeData map[string]any,\n\topts ...AgentRunOption) (*AsyncIterator[*AgentEvent], error) {\n\tif r.store == nil {\n\t\treturn nil, fmt.Errorf(\"failed to resume: store is nil\")\n\t}\n\n\tctx, runCtx, resumeInfo, err := r.loadCheckPoint(ctx, checkPointID)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to load from checkpoint: %w\", err)\n\t}\n\n\to := getCommonOptions(nil, opts...)\n\tif o.sharedParentSession {\n\t\tparentSession := getSession(ctx)\n\t\tif parentSession != nil {\n\t\t\trunCtx.Session.Values = parentSession.Values\n\t\t\trunCtx.Session.valuesMtx = parentSession.valuesMtx\n\t\t}\n\t}\n\tif runCtx.Session.valuesMtx == nil {\n\t\trunCtx.Session.valuesMtx = &sync.Mutex{}\n\t}\n\tif runCtx.Session.Values == nil {\n\t\trunCtx.Session.Values = make(map[string]any)\n\t}\n\n\tctx = setRunCtx(ctx, runCtx)\n\n\tAddSessionValues(ctx, o.sessionValues)\n\n\tif len(resumeData) > 0 {\n\t\tctx = core.BatchResumeWithData(ctx, resumeData)\n\t}\n\n\tfa := toFlowAgent(ctx, r.a)\n\taIter := fa.Resume(ctx, resumeInfo, opts...)\n\tif r.store == nil {\n\t\treturn aIter, nil\n\t}\n\n\tniter, gen := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo r.handleIter(ctx, aIter, gen, &checkPointID)\n\treturn niter, nil\n}\n\nfunc (r *Runner) handleIter(ctx context.Context, aIter *AsyncIterator[*AgentEvent],\n\tgen *AsyncGenerator[*AgentEvent], checkPointID *string) {\n\tdefer func() {\n\t\tpanicErr := recover()\n\t\tif panicErr != nil {\n\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\tgen.Send(&AgentEvent{Err: e})\n\t\t}\n\n\t\tgen.Close()\n\t}()\n\tvar (\n\t\tinterruptSignal *core.InterruptSignal\n\t\tlegacyData      any\n\t)\n\tfor {\n\t\tevent, ok := aIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\n\t\tif event.Action != nil && event.Action.internalInterrupted != nil {\n\t\t\tif interruptSignal != nil {\n\t\t\t\t// even if multiple interrupt happens, they should be merged into one\n\t\t\t\t// action by CompositeInterrupt, so here in Runner we must assume at most\n\t\t\t\t// one interrupt action happens\n\t\t\t\tpanic(\"multiple interrupt actions should not happen in Runner\")\n\t\t\t}\n\t\t\tinterruptSignal = event.Action.internalInterrupted\n\t\t\tinterruptContexts := core.ToInterruptContexts(interruptSignal, allowedAddressSegmentTypes)\n\t\t\tevent = &AgentEvent{\n\t\t\t\tAgentName: event.AgentName,\n\t\t\t\tRunPath:   event.RunPath,\n\t\t\t\tOutput:    event.Output,\n\t\t\t\tAction: &AgentAction{\n\t\t\t\t\tInterrupted: &InterruptInfo{\n\t\t\t\t\t\tData:              event.Action.Interrupted.Data,\n\t\t\t\t\t\tInterruptContexts: interruptContexts,\n\t\t\t\t\t},\n\t\t\t\t\tinternalInterrupted: interruptSignal,\n\t\t\t\t},\n\t\t\t}\n\t\t\tlegacyData = event.Action.Interrupted.Data\n\n\t\t\tif checkPointID != nil {\n\t\t\t\t// save checkpoint first before sending interrupt event,\n\t\t\t\t// so when end-user receives interrupt event, they can resume from this checkpoint\n\t\t\t\terr := r.saveCheckPoint(ctx, *checkPointID, &InterruptInfo{\n\t\t\t\t\tData: legacyData,\n\t\t\t\t}, interruptSignal)\n\t\t\t\tif err != nil {\n\t\t\t\t\tgen.Send(&AgentEvent{Err: fmt.Errorf(\"failed to save checkpoint: %w\", err)})\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tgen.Send(event)\n\t}\n}\n"
  },
  {
    "path": "adk/runner_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// mockRunnerAgent is a simple implementation of the Agent interface for testing Runner\ntype mockRunnerAgent struct {\n\tname        string\n\tdescription string\n\tresponses   []*AgentEvent\n\t// Track calls to verify correct parameters were passed\n\tcallCount       int\n\tlastInput       *AgentInput\n\tenableStreaming bool\n}\n\nfunc (a *mockRunnerAgent) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *mockRunnerAgent) Description(_ context.Context) string {\n\treturn a.description\n}\n\nfunc (a *mockRunnerAgent) Run(_ context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t// Record the call details for verification\n\ta.callCount++\n\ta.lastInput = input\n\ta.enableStreaming = input.EnableStreaming\n\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo func() {\n\t\tdefer generator.Close()\n\n\t\tfor _, event := range a.responses {\n\t\t\tgenerator.Send(event)\n\n\t\t\t// If the event has an Exit action, stop sending events\n\t\t\tif event.Action != nil && event.Action.Exit {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}()\n\n\treturn iterator\n}\n\nfunc newMockRunnerAgent(name, description string, responses []*AgentEvent) *mockRunnerAgent {\n\treturn &mockRunnerAgent{\n\t\tname:        name,\n\t\tdescription: description,\n\t\tresponses:   responses,\n\t}\n}\n\nfunc TestNewRunner(t *testing.T) {\n\tctx := context.Background()\n\tconfig := RunnerConfig{}\n\n\trunner := NewRunner(ctx, config)\n\n\t// Verify that a non-nil runner is returned\n\tassert.NotNil(t, runner)\n}\n\nfunc TestRunner_Run(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock agent with predefined responses\n\tmockAgent_ := newMockRunnerAgent(\"TestAgent\", \"Test agent for Runner\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"TestAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Response from test agent\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t}},\n\t})\n\n\t// Create a runner\n\trunner := NewRunner(ctx, RunnerConfig{Agent: mockAgent_})\n\n\t// Create test messages\n\tmsgs := []Message{\n\t\tschema.UserMessage(\"Hello, agent!\"),\n\t}\n\n\t// Test Run method without streaming\n\titerator := runner.Run(ctx, msgs)\n\n\t// Verify that the agent's Run method was called with the correct parameters\n\tassert.Equal(t, 1, mockAgent_.callCount)\n\tassert.Equal(t, msgs, mockAgent_.lastInput.Messages)\n\tassert.False(t, mockAgent_.enableStreaming)\n\n\t// Verify that we can get the expected response from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"TestAgent\", event.AgentName)\n\tassert.NotNil(t, event.Output)\n\tassert.NotNil(t, event.Output.MessageOutput)\n\tassert.NotNil(t, event.Output.MessageOutput.Message)\n\tassert.Equal(t, \"Response from test agent\", event.Output.MessageOutput.Message.Content)\n\n\t// Verify that the iterator is now closed\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestRunner_Run_WithStreaming(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock agent with predefined responses\n\tmockAgent_ := newMockRunnerAgent(\"TestAgent\", \"Test agent for Runner\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"TestAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessage:       nil,\n\t\t\t\t\tMessageStream: schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(\"Streaming response\", nil)}),\n\t\t\t\t\tRole:          schema.Assistant,\n\t\t\t\t},\n\t\t\t}},\n\t})\n\n\t// Create a runner\n\trunner := NewRunner(ctx, RunnerConfig{EnableStreaming: true, Agent: mockAgent_})\n\n\t// Create test messages\n\tmsgs := []Message{\n\t\tschema.UserMessage(\"Hello, agent!\"),\n\t}\n\n\t// Test Run method with streaming enabled\n\titerator := runner.Run(ctx, msgs)\n\n\t// Verify that the agent's Run method was called with the correct parameters\n\tassert.Equal(t, 1, mockAgent_.callCount)\n\tassert.Equal(t, msgs, mockAgent_.lastInput.Messages)\n\tassert.True(t, mockAgent_.enableStreaming)\n\n\t// Verify that we can get the expected response from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"TestAgent\", event.AgentName)\n\tassert.NotNil(t, event.Output)\n\tassert.NotNil(t, event.Output.MessageOutput)\n\tassert.True(t, event.Output.MessageOutput.IsStreaming)\n\n\t// Verify that the iterator is now closed\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestRunner_Query(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock agent with predefined responses\n\tmockAgent_ := newMockRunnerAgent(\"TestAgent\", \"Test agent for Runner\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"TestAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Response to query\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t}},\n\t})\n\n\t// Create a runner\n\trunner := NewRunner(ctx, RunnerConfig{Agent: mockAgent_})\n\n\t// Test Query method\n\titerator := runner.Query(ctx, \"Test query\")\n\n\t// Verify that the agent's Run method was called with the correct parameters\n\tassert.Equal(t, 1, mockAgent_.callCount)\n\tassert.Equal(t, 1, len(mockAgent_.lastInput.Messages))\n\tassert.Equal(t, \"Test query\", mockAgent_.lastInput.Messages[0].Content)\n\tassert.False(t, mockAgent_.enableStreaming)\n\n\t// Verify that we can get the expected response from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"TestAgent\", event.AgentName)\n\tassert.NotNil(t, event.Output)\n\tassert.NotNil(t, event.Output.MessageOutput)\n\tassert.NotNil(t, event.Output.MessageOutput.Message)\n\tassert.Equal(t, \"Response to query\", event.Output.MessageOutput.Message.Content)\n\n\t// Verify that the iterator is now closed\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestRunner_Query_WithStreaming(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock agent with predefined responses\n\tmockAgent_ := newMockRunnerAgent(\"TestAgent\", \"Test agent for Runner\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"TestAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessage:       nil,\n\t\t\t\t\tMessageStream: schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(\"Streaming query response\", nil)}),\n\t\t\t\t\tRole:          schema.Assistant,\n\t\t\t\t},\n\t\t\t}},\n\t})\n\n\t// Create a runner\n\trunner := NewRunner(ctx, RunnerConfig{EnableStreaming: true, Agent: mockAgent_})\n\n\t// Test Query method with streaming enabled\n\titerator := runner.Query(ctx, \"Test query\")\n\n\t// Verify that the agent's Run method was called with the correct parameters\n\tassert.Equal(t, 1, mockAgent_.callCount)\n\tassert.Equal(t, 1, len(mockAgent_.lastInput.Messages))\n\tassert.Equal(t, \"Test query\", mockAgent_.lastInput.Messages[0].Content)\n\tassert.True(t, mockAgent_.enableStreaming)\n\n\t// Verify that we can get the expected response from the iterator\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"TestAgent\", event.AgentName)\n\tassert.NotNil(t, event.Output)\n\tassert.NotNil(t, event.Output.MessageOutput)\n\tassert.True(t, event.Output.MessageOutput.IsStreaming)\n\n\t// Verify that the iterator is now closed\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n"
  },
  {
    "path": "adk/utils.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"io\"\n\t\"strings\"\n\n\t\"github.com/google/uuid\"\n\n\t\"github.com/cloudwego/eino/internal\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype AsyncIterator[T any] struct {\n\tch *internal.UnboundedChan[T]\n}\n\nfunc (ai *AsyncIterator[T]) Next() (T, bool) {\n\treturn ai.ch.Receive()\n}\n\ntype AsyncGenerator[T any] struct {\n\tch *internal.UnboundedChan[T]\n}\n\nfunc (ag *AsyncGenerator[T]) Send(v T) {\n\tag.ch.Send(v)\n}\n\nfunc (ag *AsyncGenerator[T]) Close() {\n\tag.ch.Close()\n}\n\n// NewAsyncIteratorPair returns a paired async iterator and generator\n// that share the same underlying channel.\nfunc NewAsyncIteratorPair[T any]() (*AsyncIterator[T], *AsyncGenerator[T]) {\n\tch := internal.NewUnboundedChan[T]()\n\treturn &AsyncIterator[T]{ch}, &AsyncGenerator[T]{ch}\n}\n\nfunc copyMap[K comparable, V any](m map[K]V) map[K]V {\n\tres := make(map[K]V, len(m))\n\tfor k, v := range m {\n\t\tres[k] = v\n\t}\n\treturn res\n}\n\nfunc cloneSlice[T any](s []T) []T {\n\tif s == nil {\n\t\treturn nil\n\t}\n\tres := make([]T, len(s))\n\tcopy(res, s)\n\treturn res\n}\n\nfunc concatInstructions(instructions ...string) string {\n\tvar sb strings.Builder\n\tsb.WriteString(instructions[0])\n\tfor i := 1; i < len(instructions); i++ {\n\t\tsb.WriteString(\"\\n\\n\")\n\t\tsb.WriteString(instructions[i])\n\t}\n\n\treturn sb.String()\n}\n\n// GenTransferMessages generates assistant and tool messages to instruct a\n// transfer-to-agent tool call targeting the destination agent.\nfunc GenTransferMessages(_ context.Context, destAgentName string) (Message, Message) {\n\ttoolCallID := uuid.NewString()\n\ttooCall := schema.ToolCall{ID: toolCallID, Function: schema.FunctionCall{Name: TransferToAgentToolName, Arguments: destAgentName}}\n\tassistantMessage := schema.AssistantMessage(\"\", []schema.ToolCall{tooCall})\n\tmsg := transferToAgentToolOutput(destAgentName)\n\ttoolMessage := schema.ToolMessage(msg, toolCallID, schema.WithToolName(TransferToAgentToolName))\n\treturn assistantMessage, toolMessage\n}\n\n// set automatic close for event's message stream\nfunc setAutomaticClose(e *AgentEvent) {\n\tif e.Output == nil || e.Output.MessageOutput == nil || !e.Output.MessageOutput.IsStreaming {\n\t\treturn\n\t}\n\n\te.Output.MessageOutput.MessageStream.SetAutomaticClose()\n}\n\n// getMessageFromWrappedEvent extracts the message from an AgentEvent.\n// If the stream contains an error chunk, this function returns (nil, err) and\n// sets StreamErr to prevent re-consumption. The nil message ensures that\n// failed stream responses are not included in subsequent agents' context windows.\nfunc getMessageFromWrappedEvent(e *agentEventWrapper) (Message, error) {\n\tif e.AgentEvent.Output == nil || e.AgentEvent.Output.MessageOutput == nil {\n\t\treturn nil, nil\n\t}\n\n\tif !e.AgentEvent.Output.MessageOutput.IsStreaming {\n\t\treturn e.AgentEvent.Output.MessageOutput.Message, nil\n\t}\n\n\tif e.concatenatedMessage != nil {\n\t\treturn e.concatenatedMessage, nil\n\t}\n\n\tif e.StreamErr != nil {\n\t\treturn nil, e.StreamErr\n\t}\n\n\te.mu.Lock()\n\tdefer e.mu.Unlock()\n\tif e.concatenatedMessage != nil {\n\t\treturn e.concatenatedMessage, nil\n\t}\n\n\tvar (\n\t\tmsgs []Message\n\t\ts    = e.AgentEvent.Output.MessageOutput.MessageStream\n\t)\n\n\tdefer s.Close()\n\tfor {\n\t\tmsg, err := s.Recv()\n\t\tif err != nil {\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\te.StreamErr = err\n\t\t\t// Replace the stream with successfully received messages only (no error at the end).\n\t\t\t// The error is preserved in StreamErr for users to check.\n\t\t\t// We intentionally exclude the error from the new stream to ensure gob encoding\n\t\t\t// compatibility, as the stream may be consumed during serialization.\n\t\t\te.AgentEvent.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs)\n\t\t\treturn nil, err\n\t\t}\n\n\t\tmsgs = append(msgs, msg)\n\t}\n\n\tif len(msgs) == 0 {\n\t\treturn nil, errors.New(\"no messages in MessageVariant.MessageStream\")\n\t}\n\n\tif len(msgs) == 1 {\n\t\te.concatenatedMessage = msgs[0]\n\t} else {\n\t\tvar err error\n\t\te.concatenatedMessage, err = schema.ConcatMessages(msgs)\n\t\tif err != nil {\n\t\t\te.StreamErr = err\n\t\t\te.AgentEvent.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs)\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\treturn e.concatenatedMessage, nil\n}\n\n// copyAgentEvent copies an AgentEvent.\n// If the MessageVariant is streaming, the MessageStream will be copied.\n// RunPath will be deep copied.\n// The result of Copy will be a new AgentEvent that is:\n// - safe to set fields of AgentEvent\n// - safe to extend RunPath\n// - safe to receive from MessageStream\n// NOTE: even if the AgentEvent is copied, it's still not recommended to modify\n// the Message itself or Chunks of the MessageStream, as they are not copied.\n// NOTE: if you have CustomizedOutput or CustomizedAction, they are NOT copied.\nfunc copyAgentEvent(ae *AgentEvent) *AgentEvent {\n\trp := make([]RunStep, len(ae.RunPath))\n\tcopy(rp, ae.RunPath)\n\n\tcopied := &AgentEvent{\n\t\tAgentName: ae.AgentName,\n\t\tRunPath:   rp,\n\t\tAction:    ae.Action,\n\t\tErr:       ae.Err,\n\t}\n\n\tif ae.Output == nil {\n\t\treturn copied\n\t}\n\n\tcopied.Output = &AgentOutput{\n\t\tCustomizedOutput: ae.Output.CustomizedOutput,\n\t}\n\n\tmv := ae.Output.MessageOutput\n\tif mv == nil {\n\t\treturn copied\n\t}\n\n\tcopied.Output.MessageOutput = &MessageVariant{\n\t\tIsStreaming: mv.IsStreaming,\n\t\tRole:        mv.Role,\n\t\tToolName:    mv.ToolName,\n\t}\n\tif mv.IsStreaming {\n\t\tsts := ae.Output.MessageOutput.MessageStream.Copy(2)\n\t\tmv.MessageStream = sts[0]\n\t\tcopied.Output.MessageOutput.MessageStream = sts[1]\n\t} else {\n\t\tcopied.Output.MessageOutput.Message = mv.Message\n\t}\n\n\treturn copied\n}\n\n// GetMessage extracts the Message from an AgentEvent. For streaming output,\n// it duplicates the stream and concatenates it into a single Message.\nfunc GetMessage(e *AgentEvent) (Message, *AgentEvent, error) {\n\tif e.Output == nil || e.Output.MessageOutput == nil {\n\t\treturn nil, e, nil\n\t}\n\n\tmsgOutput := e.Output.MessageOutput\n\tif msgOutput.IsStreaming {\n\t\tss := msgOutput.MessageStream.Copy(2)\n\t\te.Output.MessageOutput.MessageStream = ss[0]\n\n\t\tmsg, err := schema.ConcatMessageStream(ss[1])\n\n\t\treturn msg, e, err\n\t}\n\n\treturn msgOutput.Message, e, nil\n}\n\nfunc genErrorIter(err error) *AsyncIterator[*AgentEvent] {\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\tgenerator.Send(&AgentEvent{Err: err})\n\tgenerator.Close()\n\treturn iterator\n}\n"
  },
  {
    "path": "adk/utils_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"bytes\"\n\t\"encoding/gob\"\n\t\"errors\"\n\t\"fmt\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestAsyncIteratorPair_Basic(t *testing.T) {\n\t// Create a new iterator-generator pair\n\titerator, generator := NewAsyncIteratorPair[string]()\n\n\t// Test sending and receiving a value\n\tgenerator.Send(\"test1\")\n\tval, ok := iterator.Next()\n\tif !ok {\n\t\tt.Error(\"receive should succeed\")\n\t}\n\tif val != \"test1\" {\n\t\tt.Errorf(\"expected 'test1', got '%s'\", val)\n\t}\n\n\t// Test sending and receiving multiple values\n\tgenerator.Send(\"test2\")\n\tgenerator.Send(\"test3\")\n\n\tval, ok = iterator.Next()\n\tif !ok {\n\t\tt.Error(\"receive should succeed\")\n\t}\n\tif val != \"test2\" {\n\t\tt.Errorf(\"expected 'test2', got '%s'\", val)\n\t}\n\n\tval, ok = iterator.Next()\n\tif !ok {\n\t\tt.Error(\"receive should succeed\")\n\t}\n\tif val != \"test3\" {\n\t\tt.Errorf(\"expected 'test3', got '%s'\", val)\n\t}\n}\n\nfunc TestAsyncIteratorPair_Close(t *testing.T) {\n\titerator, generator := NewAsyncIteratorPair[int]()\n\n\t// Send some values\n\tgenerator.Send(1)\n\tgenerator.Send(2)\n\n\t// Close the generator\n\tgenerator.Close()\n\n\t// Should still be able to read existing values\n\tval, ok := iterator.Next()\n\tif !ok {\n\t\tt.Error(\"receive should succeed\")\n\t}\n\tif val != 1 {\n\t\tt.Errorf(\"expected 1, got %d\", val)\n\t}\n\n\tval, ok = iterator.Next()\n\tif !ok {\n\t\tt.Error(\"receive should succeed\")\n\t}\n\tif val != 2 {\n\t\tt.Errorf(\"expected 2, got %d\", val)\n\t}\n\n\t// After consuming all values, Next should return false\n\t_, ok = iterator.Next()\n\tif ok {\n\t\tt.Error(\"receive from closed, empty channel should return ok=false\")\n\t}\n}\n\nfunc TestAsyncIteratorPair_Concurrency(t *testing.T) {\n\titerator, generator := NewAsyncIteratorPair[int]()\n\tconst numSenders = 5\n\tconst numReceivers = 3\n\tconst messagesPerSender = 100\n\n\tvar rwg, swg sync.WaitGroup\n\trwg.Add(numReceivers)\n\tswg.Add(numSenders)\n\n\t// Start senders\n\tfor i := 0; i < numSenders; i++ {\n\t\tgo func(id int) {\n\t\t\tdefer swg.Done()\n\t\t\tfor j := 0; j < messagesPerSender; j++ {\n\t\t\t\tgenerator.Send(id*messagesPerSender + j)\n\t\t\t\ttime.Sleep(time.Microsecond) // Small delay to increase concurrency chance\n\t\t\t}\n\t\t}(i)\n\t}\n\n\t// Start receivers\n\treceived := make([]int, 0, numSenders*messagesPerSender)\n\tvar mu sync.Mutex\n\n\tfor i := 0; i < numReceivers; i++ {\n\t\tgo func() {\n\t\t\tdefer rwg.Done()\n\t\t\tfor {\n\t\t\t\tval, ok := iterator.Next()\n\t\t\t\tif !ok {\n\t\t\t\t\treturn\n\t\t\t\t}\n\t\t\t\tmu.Lock()\n\t\t\t\treceived = append(received, val)\n\t\t\t\tmu.Unlock()\n\t\t\t}\n\t\t}()\n\t}\n\n\t// Wait for senders to finish\n\tswg.Wait()\n\tgenerator.Close()\n\n\t// Wait for all goroutines to finish\n\trwg.Wait()\n\n\t// Verify we received all messages\n\tif len(received) != numSenders*messagesPerSender {\n\t\tt.Errorf(\"expected %d messages, got %d\", numSenders*messagesPerSender, len(received))\n\t}\n\n\t// Create a map to check for duplicates and missing values\n\treceivedMap := make(map[int]bool)\n\tfor _, val := range received {\n\t\treceivedMap[val] = true\n\t}\n\n\tif len(receivedMap) != numSenders*messagesPerSender {\n\t\tt.Error(\"duplicate or missing messages detected\")\n\t}\n}\n\nfunc TestGenErrorIter(t *testing.T) {\n\titer := genErrorIter(fmt.Errorf(\"test\"))\n\te, ok := iter.Next()\n\tassert.True(t, ok)\n\tassert.Equal(t, \"test\", e.Err.Error())\n\t_, ok = iter.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestGetMessageFromWrappedEvent_StreamError_MultipleCallsGuard(t *testing.T) {\n\tstreamErr := errors.New(\"stream error\")\n\n\tsr, sw := schema.Pipe[Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(schema.AssistantMessage(\"chunk1\", nil), nil)\n\t\tsw.Send(schema.AssistantMessage(\"chunk2\", nil), nil)\n\t\tsw.Send(nil, streamErr)\n\t}()\n\n\twrapper := &agentEventWrapper{\n\t\tAgentEvent: &AgentEvent{\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessageStream: sr,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\n\tmsg1, err1 := getMessageFromWrappedEvent(wrapper)\n\tassert.Nil(t, msg1)\n\tassert.NotNil(t, err1)\n\tassert.Equal(t, \"stream error\", err1.Error())\n\n\tassert.NotEmpty(t, wrapper.StreamErr)\n\tassert.Equal(t, err1, wrapper.StreamErr)\n\n\tmsg2, err2 := getMessageFromWrappedEvent(wrapper)\n\tassert.Nil(t, msg2)\n\tassert.Equal(t, err1, err2)\n}\n\nfunc TestGetMessageFromWrappedEvent_StreamSuccess_MultipleCallsCached(t *testing.T) {\n\tsr, sw := schema.Pipe[Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(schema.AssistantMessage(\"chunk1\", nil), nil)\n\t\tsw.Send(schema.AssistantMessage(\"chunk2\", nil), nil)\n\t}()\n\n\twrapper := &agentEventWrapper{\n\t\tAgentEvent: &AgentEvent{\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessageStream: sr,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\n\tmsg1, err1 := getMessageFromWrappedEvent(wrapper)\n\tassert.NotNil(t, msg1)\n\tassert.Nil(t, err1)\n\tassert.Equal(t, \"chunk1chunk2\", msg1.Content)\n\n\tassert.NotNil(t, wrapper.concatenatedMessage)\n\n\tmsg2, err2 := getMessageFromWrappedEvent(wrapper)\n\tassert.NotNil(t, msg2)\n\tassert.Nil(t, err2)\n\tassert.Equal(t, \"chunk1chunk2\", msg2.Content)\n\tassert.Same(t, msg1, msg2)\n}\n\nfunc TestGetMessageFromWrappedEvent_StreamError_PartialMessagesPreserved(t *testing.T) {\n\tstreamErr := errors.New(\"stream error at chunk3\")\n\n\tsr, sw := schema.Pipe[Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(schema.AssistantMessage(\"chunk1\", nil), nil)\n\t\tsw.Send(schema.AssistantMessage(\"chunk2\", nil), nil)\n\t\tsw.Send(nil, streamErr)\n\t}()\n\n\twrapper := &agentEventWrapper{\n\t\tAgentEvent: &AgentEvent{\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessageStream: sr,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\n\t_, err := getMessageFromWrappedEvent(wrapper)\n\tassert.NotNil(t, err)\n\tassert.Equal(t, streamErr, wrapper.StreamErr)\n\n\tnewStream := wrapper.AgentEvent.Output.MessageOutput.MessageStream\n\tassert.NotNil(t, newStream)\n\n\tvar msgs []Message\n\tfor {\n\t\tmsg, err := newStream.Recv()\n\t\tif err != nil {\n\t\t\tbreak\n\t\t}\n\t\tmsgs = append(msgs, msg)\n\t}\n\n\tassert.Equal(t, 2, len(msgs))\n\tassert.Equal(t, \"chunk1\", msgs[0].Content)\n\tassert.Equal(t, \"chunk2\", msgs[1].Content)\n}\n\nfunc TestAgentEventWrapper_GobEncoding_WithWillRetryError(t *testing.T) {\n\tstreamErr := &WillRetryError{ErrStr: \"stream error\", RetryAttempt: 2}\n\n\tsr, sw := schema.Pipe[Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(schema.AssistantMessage(\"partial1\", nil), nil)\n\t\tsw.Send(schema.AssistantMessage(\"partial2\", nil), nil)\n\t\tsw.Send(nil, streamErr)\n\t}()\n\n\twrapper := &agentEventWrapper{\n\t\tAgentEvent: &AgentEvent{\n\t\t\tAgentName: \"TestAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessageStream: sr,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\tTS: 12345,\n\t}\n\n\t_, err := getMessageFromWrappedEvent(wrapper)\n\tassert.NotNil(t, err)\n\tvar wrapperErr *WillRetryError\n\tassert.True(t, errors.As(wrapper.StreamErr, &wrapperErr))\n\tassert.Equal(t, streamErr.ErrStr, wrapperErr.ErrStr)\n\tassert.Equal(t, streamErr.RetryAttempt, wrapperErr.RetryAttempt)\n\n\tvar buf bytes.Buffer\n\tenc := gob.NewEncoder(&buf)\n\terr = enc.Encode(wrapper)\n\tassert.NoError(t, err)\n\n\tvar decoded agentEventWrapper\n\tdec := gob.NewDecoder(&buf)\n\terr = dec.Decode(&decoded)\n\tassert.NoError(t, err)\n\n\tassert.Equal(t, \"TestAgent\", decoded.AgentName)\n\tassert.Equal(t, int64(12345), decoded.TS)\n\tvar decodedErr *WillRetryError\n\tassert.True(t, errors.As(decoded.StreamErr, &decodedErr))\n\tassert.Equal(t, streamErr.ErrStr, decodedErr.ErrStr)\n\tassert.Equal(t, streamErr.RetryAttempt, decodedErr.RetryAttempt)\n}\n\nfunc TestAgentEventWrapper_GobEncoding_WithUnregisteredError(t *testing.T) {\n\tstreamErr := errors.New(\"unregistered error type\")\n\n\tsr, sw := schema.Pipe[Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(schema.AssistantMessage(\"partial1\", nil), nil)\n\t\tsw.Send(nil, streamErr)\n\t}()\n\n\twrapper := &agentEventWrapper{\n\t\tAgentEvent: &AgentEvent{\n\t\t\tAgentName: \"TestAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessageStream: sr,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\tTS: 22222,\n\t}\n\n\t_, err := getMessageFromWrappedEvent(wrapper)\n\tassert.NotNil(t, err)\n\tassert.Equal(t, streamErr, wrapper.StreamErr)\n\n\tvar buf bytes.Buffer\n\tenc := gob.NewEncoder(&buf)\n\terr = enc.Encode(wrapper)\n\tassert.Error(t, err, \"gob encoding should fail for unregistered error types\")\n}\n\nfunc TestAgentEventWrapper_GobEncoding_WithStreamSuccess(t *testing.T) {\n\tsr, sw := schema.Pipe[Message](10)\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tsw.Send(schema.AssistantMessage(\"success1\", nil), nil)\n\t\tsw.Send(schema.AssistantMessage(\"success2\", nil), nil)\n\t}()\n\n\twrapper := &agentEventWrapper{\n\t\tAgentEvent: &AgentEvent{\n\t\t\tAgentName: \"TestAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming:   true,\n\t\t\t\t\tMessageStream: sr,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\tTS: 67890,\n\t}\n\n\tmsg, err := getMessageFromWrappedEvent(wrapper)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"success1success2\", msg.Content)\n\n\tvar buf bytes.Buffer\n\tenc := gob.NewEncoder(&buf)\n\terr = enc.Encode(wrapper)\n\tassert.NoError(t, err)\n\n\tvar decoded agentEventWrapper\n\tdec := gob.NewDecoder(&buf)\n\terr = dec.Decode(&decoded)\n\tassert.NoError(t, err)\n\n\tassert.Equal(t, \"TestAgent\", decoded.AgentName)\n\tassert.Equal(t, int64(67890), decoded.TS)\n\tassert.Empty(t, decoded.StreamErr)\n}\n"
  },
  {
    "path": "adk/workflow.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"runtime/debug\"\n\t\"sync\"\n\n\t\"github.com/cloudwego/eino/internal/core\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype workflowAgentMode int\n\nconst (\n\tworkflowAgentModeUnknown workflowAgentMode = iota\n\tworkflowAgentModeSequential\n\tworkflowAgentModeLoop\n\tworkflowAgentModeParallel\n)\n\ntype workflowAgent struct {\n\tname        string\n\tdescription string\n\tsubAgents   []*flowAgent\n\n\tmode workflowAgentMode\n\n\tmaxIterations int\n}\n\nfunc (a *workflowAgent) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *workflowAgent) Description(_ context.Context) string {\n\treturn a.description\n}\n\nfunc (a *workflowAgent) GetType() string {\n\tswitch a.mode {\n\tcase workflowAgentModeSequential:\n\t\treturn \"Sequential\"\n\tcase workflowAgentModeParallel:\n\t\treturn \"Parallel\"\n\tcase workflowAgentModeLoop:\n\t\treturn \"Loop\"\n\tdefault:\n\t\treturn \"WorkflowAgent\"\n\t}\n}\n\nfunc (a *workflowAgent) Run(ctx context.Context, _ *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo func() {\n\n\t\tvar err error\n\t\tdefer func() {\n\t\t\tpanicErr := recover()\n\t\t\tif panicErr != nil {\n\t\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\t\tgenerator.Send(&AgentEvent{Err: e})\n\t\t\t} else if err != nil {\n\t\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t\t}\n\n\t\t\tgenerator.Close()\n\t\t}()\n\n\t\t// Different workflow execution based on mode\n\t\tswitch a.mode {\n\t\tcase workflowAgentModeSequential:\n\t\t\terr = a.runSequential(ctx, generator, nil, nil, opts...)\n\t\tcase workflowAgentModeLoop:\n\t\t\terr = a.runLoop(ctx, generator, nil, nil, opts...)\n\t\tcase workflowAgentModeParallel:\n\t\t\terr = a.runParallel(ctx, generator, nil, nil, opts...)\n\t\tdefault:\n\t\t\terr = fmt.Errorf(\"unsupported workflow agent mode: %d\", a.mode)\n\t\t}\n\t}()\n\n\treturn iterator\n}\n\ntype sequentialWorkflowState struct {\n\tInterruptIndex int\n}\n\ntype parallelWorkflowState struct {\n\tSubAgentEvents map[int][]*agentEventWrapper\n}\n\ntype loopWorkflowState struct {\n\tLoopIterations int\n\tSubAgentIndex  int\n}\n\nfunc init() {\n\tschema.RegisterName[*sequentialWorkflowState](\"eino_adk_sequential_workflow_state\")\n\tschema.RegisterName[*parallelWorkflowState](\"eino_adk_parallel_workflow_state\")\n\tschema.RegisterName[*loopWorkflowState](\"eino_adk_loop_workflow_state\")\n}\n\nfunc (a *workflowAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo func() {\n\t\tvar err error\n\t\tdefer func() {\n\t\t\tpanicErr := recover()\n\t\t\tif panicErr != nil {\n\t\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\t\tgenerator.Send(&AgentEvent{Err: e})\n\t\t\t} else if err != nil {\n\t\t\t\tgenerator.Send(&AgentEvent{Err: err})\n\t\t\t}\n\n\t\t\tgenerator.Close()\n\t\t}()\n\n\t\tstate := info.InterruptState\n\t\tif state == nil {\n\t\t\tpanic(fmt.Sprintf(\"workflowAgent.Resume: agent '%s' was asked to resume but has no state\", a.Name(ctx)))\n\t\t}\n\n\t\t// Different workflow execution based on the type of our restored state.\n\t\tswitch s := state.(type) {\n\t\tcase *sequentialWorkflowState:\n\t\t\terr = a.runSequential(ctx, generator, s, info, opts...)\n\t\tcase *parallelWorkflowState:\n\t\t\terr = a.runParallel(ctx, generator, s, info, opts...)\n\t\tcase *loopWorkflowState:\n\t\t\terr = a.runLoop(ctx, generator, s, info, opts...)\n\t\tdefault:\n\t\t\terr = fmt.Errorf(\"unsupported workflow agent state type: %T\", s)\n\t\t}\n\t}()\n\treturn iterator\n}\n\n// WorkflowInterruptInfo CheckpointSchema: persisted via InterruptInfo.Data (gob).\ntype WorkflowInterruptInfo struct {\n\tOrigInput *AgentInput\n\n\tSequentialInterruptIndex int\n\tSequentialInterruptInfo  *InterruptInfo\n\n\tLoopIterations int\n\n\tParallelInterruptInfo map[int] /*index*/ *InterruptInfo\n}\n\nfunc (a *workflowAgent) runSequential(ctx context.Context,\n\tgenerator *AsyncGenerator[*AgentEvent], seqState *sequentialWorkflowState, info *ResumeInfo,\n\topts ...AgentRunOption) (err error) {\n\n\tstartIdx := 0\n\n\t// seqCtx tracks the accumulated RunPath across the sequence.\n\tseqCtx := ctx\n\n\t// If we are resuming, find which sub-agent to start from and prepare its context.\n\tif seqState != nil {\n\t\tstartIdx = seqState.InterruptIndex\n\n\t\tvar steps []string\n\t\tfor i := 0; i < startIdx; i++ {\n\t\t\tsteps = append(steps, a.subAgents[i].Name(seqCtx))\n\t\t}\n\n\t\tseqCtx = updateRunPathOnly(seqCtx, steps...)\n\t}\n\n\tfor i := startIdx; i < len(a.subAgents); i++ {\n\t\tsubAgent := a.subAgents[i]\n\n\t\tvar subIterator *AsyncIterator[*AgentEvent]\n\t\tif seqState != nil {\n\t\t\tsubIterator = subAgent.Resume(seqCtx, &ResumeInfo{\n\t\t\t\tEnableStreaming: info.EnableStreaming,\n\t\t\t\tInterruptInfo:   info.Data.(*WorkflowInterruptInfo).SequentialInterruptInfo,\n\t\t\t}, opts...)\n\t\t\tseqState = nil\n\t\t} else {\n\t\t\tsubIterator = subAgent.Run(seqCtx, nil, opts...)\n\t\t}\n\n\t\tseqCtx = updateRunPathOnly(seqCtx, subAgent.Name(seqCtx))\n\n\t\tvar lastActionEvent *AgentEvent\n\t\tfor {\n\t\t\tevent, ok := subIterator.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tif event.Err != nil {\n\t\t\t\t// exit if report error\n\t\t\t\tgenerator.Send(event)\n\t\t\t\treturn nil\n\t\t\t}\n\n\t\t\tif lastActionEvent != nil {\n\t\t\t\tgenerator.Send(lastActionEvent)\n\t\t\t\tlastActionEvent = nil\n\t\t\t}\n\n\t\t\tif event.Action != nil {\n\t\t\t\tlastActionEvent = event\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tgenerator.Send(event)\n\t\t}\n\n\t\tif lastActionEvent != nil {\n\t\t\tif lastActionEvent.Action.internalInterrupted != nil {\n\t\t\t\t// A sub-agent interrupted. Wrap it with our own state, including the index.\n\t\t\t\tstate := &sequentialWorkflowState{\n\t\t\t\t\tInterruptIndex: i,\n\t\t\t\t}\n\t\t\t\t// Use CompositeInterrupt to funnel the sub-interrupt and add our own state.\n\t\t\t\t// The context for the composite interrupt must be the one from *before* the sub-agent ran.\n\t\t\t\tevent := CompositeInterrupt(ctx, \"Sequential workflow interrupted\", state,\n\t\t\t\t\tlastActionEvent.Action.internalInterrupted)\n\n\t\t\t\t// For backward compatibility, populate the deprecated Data field.\n\t\t\t\tevent.Action.Interrupted.Data = &WorkflowInterruptInfo{\n\t\t\t\t\tOrigInput:                getRunCtx(ctx).RootInput,\n\t\t\t\t\tSequentialInterruptIndex: i,\n\t\t\t\t\tSequentialInterruptInfo:  lastActionEvent.Action.Interrupted,\n\t\t\t\t}\n\t\t\t\tevent.AgentName = lastActionEvent.AgentName\n\t\t\t\tevent.RunPath = lastActionEvent.RunPath\n\n\t\t\t\tgenerator.Send(event)\n\t\t\t\treturn nil\n\t\t\t}\n\n\t\t\tif lastActionEvent.Action.Exit {\n\t\t\t\t// Forward the event\n\t\t\t\tgenerator.Send(lastActionEvent)\n\t\t\t\treturn nil\n\t\t\t}\n\n\t\t\tgenerator.Send(lastActionEvent)\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// BreakLoopAction is a programmatic-only agent action used to prematurely\n// terminate the execution of a loop workflow agent.\n// When a loop workflow agent receives this action from a sub-agent, it will stop its\n// current iteration and will not proceed to the next one.\n// It will mark the BreakLoopAction as Done, signalling to any 'upper level' loop agent\n// that this action has been processed and should be ignored further up.\n// This action is not intended to be used by LLMs.\ntype BreakLoopAction struct {\n\t// From records the name of the agent that initiated the break loop action.\n\tFrom string\n\t// Done is a state flag that can be used by the framework to mark when the\n\t// action has been handled.\n\tDone bool\n\t// CurrentIterations is populated by the framework to record at which\n\t// iteration the loop was broken.\n\tCurrentIterations int\n}\n\n// NewBreakLoopAction creates a new BreakLoopAction, signaling a request\n// to terminate the current loop.\nfunc NewBreakLoopAction(agentName string) *AgentAction {\n\treturn &AgentAction{BreakLoop: &BreakLoopAction{\n\t\tFrom: agentName,\n\t}}\n}\n\nfunc (a *workflowAgent) runLoop(ctx context.Context, generator *AsyncGenerator[*AgentEvent],\n\tloopState *loopWorkflowState, resumeInfo *ResumeInfo, opts ...AgentRunOption) (err error) {\n\n\tif len(a.subAgents) == 0 {\n\t\treturn nil\n\t}\n\n\tstartIter := 0\n\tstartIdx := 0\n\n\t// loopCtx tracks the accumulated RunPath across the full sequence within a single iteration.\n\tloopCtx := ctx\n\n\tif loopState != nil {\n\t\t// We are resuming.\n\t\tstartIter = loopState.LoopIterations\n\t\tstartIdx = loopState.SubAgentIndex\n\n\t\t// Rebuild the loopCtx to have the correct RunPath up to the point of resumption.\n\t\tvar steps []string\n\t\tfor i := 0; i < startIter; i++ {\n\t\t\tfor _, subAgent := range a.subAgents {\n\t\t\t\tsteps = append(steps, subAgent.Name(loopCtx))\n\t\t\t}\n\t\t}\n\t\tfor i := 0; i < startIdx; i++ {\n\t\t\tsteps = append(steps, a.subAgents[i].Name(loopCtx))\n\t\t}\n\t\tloopCtx = updateRunPathOnly(loopCtx, steps...)\n\t}\n\n\tfor i := startIter; i < a.maxIterations || a.maxIterations == 0; i++ {\n\t\tfor j := startIdx; j < len(a.subAgents); j++ {\n\t\t\tsubAgent := a.subAgents[j]\n\n\t\t\tvar subIterator *AsyncIterator[*AgentEvent]\n\t\t\tif loopState != nil {\n\t\t\t\t// This is the agent we need to resume.\n\t\t\t\tsubIterator = subAgent.Resume(loopCtx, &ResumeInfo{\n\t\t\t\t\tEnableStreaming: resumeInfo.EnableStreaming,\n\t\t\t\t\tInterruptInfo:   resumeInfo.Data.(*WorkflowInterruptInfo).SequentialInterruptInfo,\n\t\t\t\t}, opts...)\n\t\t\t\tloopState = nil // Only resume the first time.\n\t\t\t} else {\n\t\t\t\tsubIterator = subAgent.Run(loopCtx, nil, opts...)\n\t\t\t}\n\n\t\t\tloopCtx = updateRunPathOnly(loopCtx, subAgent.Name(loopCtx))\n\n\t\t\tvar lastActionEvent *AgentEvent\n\t\t\tvar breakLoopEvent *AgentEvent\n\t\t\tfor {\n\t\t\t\tevent, ok := subIterator.Next()\n\t\t\t\tif !ok {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tif event.Err != nil {\n\t\t\t\t\tgenerator.Send(event)\n\t\t\t\t\treturn nil\n\t\t\t\t}\n\n\t\t\t\tif lastActionEvent != nil {\n\t\t\t\t\tif lastActionEvent.Action.BreakLoop != nil && !lastActionEvent.Action.BreakLoop.Done {\n\t\t\t\t\t\tlastActionEvent.Action.BreakLoop.Done = true\n\t\t\t\t\t\tlastActionEvent.Action.BreakLoop.CurrentIterations = i\n\t\t\t\t\t\tbreakLoopEvent = lastActionEvent\n\t\t\t\t\t}\n\t\t\t\t\tgenerator.Send(lastActionEvent)\n\t\t\t\t\tlastActionEvent = nil\n\t\t\t\t}\n\n\t\t\t\tif event.Action != nil {\n\t\t\t\t\tlastActionEvent = event\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\t\t\t\tgenerator.Send(event)\n\t\t\t}\n\n\t\t\tif lastActionEvent != nil {\n\t\t\t\tif lastActionEvent.Action.BreakLoop != nil && !lastActionEvent.Action.BreakLoop.Done {\n\t\t\t\t\tlastActionEvent.Action.BreakLoop.Done = true\n\t\t\t\t\tlastActionEvent.Action.BreakLoop.CurrentIterations = i\n\t\t\t\t\tbreakLoopEvent = lastActionEvent\n\t\t\t\t}\n\n\t\t\t\tif lastActionEvent.Action.internalInterrupted != nil {\n\t\t\t\t\t// A sub-agent interrupted. Wrap it with our own loop state.\n\t\t\t\t\tstate := &loopWorkflowState{\n\t\t\t\t\t\tLoopIterations: i,\n\t\t\t\t\t\tSubAgentIndex:  j,\n\t\t\t\t\t}\n\t\t\t\t\t// Use CompositeInterrupt to funnel the sub-interrupt and add our own state.\n\t\t\t\t\tevent := CompositeInterrupt(ctx, \"Loop workflow interrupted\", state,\n\t\t\t\t\t\tlastActionEvent.Action.internalInterrupted)\n\n\t\t\t\t\t// For backward compatibility, populate the deprecated Data field.\n\t\t\t\t\tevent.Action.Interrupted.Data = &WorkflowInterruptInfo{\n\t\t\t\t\t\tOrigInput:                getRunCtx(ctx).RootInput,\n\t\t\t\t\t\tLoopIterations:           i,\n\t\t\t\t\t\tSequentialInterruptIndex: j,\n\t\t\t\t\t\tSequentialInterruptInfo:  lastActionEvent.Action.Interrupted,\n\t\t\t\t\t}\n\t\t\t\t\tevent.AgentName = lastActionEvent.AgentName\n\t\t\t\t\tevent.RunPath = lastActionEvent.RunPath\n\n\t\t\t\t\tgenerator.Send(event)\n\t\t\t\t\treturn\n\t\t\t\t}\n\n\t\t\t\tif lastActionEvent.Action.Exit {\n\t\t\t\t\tgenerator.Send(lastActionEvent)\n\t\t\t\t\treturn\n\t\t\t\t}\n\n\t\t\t\tgenerator.Send(lastActionEvent)\n\t\t\t}\n\n\t\t\tif breakLoopEvent != nil {\n\t\t\t\treturn\n\t\t\t}\n\t\t}\n\n\t\t// Reset the sub-agent index for the next iteration of the outer loop.\n\t\tstartIdx = 0\n\t}\n\n\treturn nil\n}\n\nfunc (a *workflowAgent) runParallel(ctx context.Context, generator *AsyncGenerator[*AgentEvent],\n\tparState *parallelWorkflowState, resumeInfo *ResumeInfo, opts ...AgentRunOption) error {\n\n\tif len(a.subAgents) == 0 {\n\t\treturn nil\n\t}\n\n\tvar (\n\t\twg                  sync.WaitGroup\n\t\tsubInterruptSignals []*core.InterruptSignal\n\t\tdataMap             = make(map[int]*InterruptInfo)\n\t\tmu                  sync.Mutex\n\t\tagentNames          map[string]bool\n\t\terr                 error\n\t\tchildContexts       = make([]context.Context, len(a.subAgents))\n\t)\n\n\t// If resuming, get the scoped ResumeInfo for each child that needs to be resumed.\n\tif parState != nil {\n\t\tagentNames, err = getNextResumeAgents(ctx, resumeInfo)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\t// Fork contexts for each sub-agent\n\tfor i := range a.subAgents {\n\t\tchildContexts[i] = forkRunCtx(ctx)\n\n\t\t// If we're resuming and this agent has existing events, add them to the child context\n\t\tif parState != nil && parState.SubAgentEvents != nil {\n\t\t\tif existingEvents, ok := parState.SubAgentEvents[i]; ok {\n\t\t\t\t// Add existing events to the child's lane events\n\t\t\t\tchildRunCtx := getRunCtx(childContexts[i])\n\t\t\t\tif childRunCtx != nil && childRunCtx.Session != nil {\n\t\t\t\t\tif childRunCtx.Session.LaneEvents == nil {\n\t\t\t\t\t\tchildRunCtx.Session.LaneEvents = &laneEvents{}\n\t\t\t\t\t}\n\t\t\t\t\tchildRunCtx.Session.LaneEvents.Events = append(childRunCtx.Session.LaneEvents.Events, existingEvents...)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tfor i := range a.subAgents {\n\t\twg.Add(1)\n\t\tgo func(idx int, agent *flowAgent) {\n\t\t\tdefer func() {\n\t\t\t\tpanicErr := recover()\n\t\t\t\tif panicErr != nil {\n\t\t\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\t\t\tgenerator.Send(&AgentEvent{Err: e})\n\t\t\t\t}\n\t\t\t\twg.Done()\n\t\t\t}()\n\n\t\t\tvar iterator *AsyncIterator[*AgentEvent]\n\n\t\t\tif _, ok := agentNames[agent.Name(ctx)]; ok {\n\t\t\t\t// This branch was interrupted and needs to be resumed.\n\t\t\t\titerator = agent.Resume(childContexts[idx], &ResumeInfo{\n\t\t\t\t\tEnableStreaming: resumeInfo.EnableStreaming,\n\t\t\t\t\tInterruptInfo:   resumeInfo.Data.(*WorkflowInterruptInfo).ParallelInterruptInfo[idx],\n\t\t\t\t}, opts...)\n\t\t\t} else if parState != nil {\n\t\t\t\t// We are resuming, but this child is not in the next points map.\n\t\t\t\t// This means it finished successfully, so we don't run it.\n\t\t\t\treturn\n\t\t\t} else {\n\t\t\t\titerator = agent.Run(childContexts[idx], nil, opts...)\n\t\t\t}\n\n\t\t\tfor {\n\t\t\t\tevent, ok := iterator.Next()\n\t\t\t\tif !ok {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tif event.Action != nil && event.Action.internalInterrupted != nil {\n\t\t\t\t\tmu.Lock()\n\t\t\t\t\tsubInterruptSignals = append(subInterruptSignals, event.Action.internalInterrupted)\n\t\t\t\t\tdataMap[idx] = event.Action.Interrupted\n\t\t\t\t\tmu.Unlock()\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tgenerator.Send(event)\n\t\t\t}\n\t\t}(i, a.subAgents[i])\n\t}\n\n\twg.Wait()\n\n\tif len(subInterruptSignals) == 0 {\n\t\t// Join all child contexts back to the parent\n\t\tjoinRunCtxs(ctx, childContexts...)\n\t\treturn nil\n\t}\n\n\tif len(subInterruptSignals) > 0 {\n\t\t// Before interrupting, collect the current events from each child context\n\t\tsubAgentEvents := make(map[int][]*agentEventWrapper)\n\t\tfor i, childCtx := range childContexts {\n\t\t\tchildRunCtx := getRunCtx(childCtx)\n\t\t\tif childRunCtx != nil && childRunCtx.Session != nil && childRunCtx.Session.LaneEvents != nil {\n\t\t\t\tsubAgentEvents[i] = childRunCtx.Session.LaneEvents.Events\n\t\t\t}\n\t\t}\n\n\t\tstate := &parallelWorkflowState{\n\t\t\tSubAgentEvents: subAgentEvents,\n\t\t}\n\t\tevent := CompositeInterrupt(ctx, \"Parallel workflow interrupted\", state, subInterruptSignals...)\n\n\t\t// For backward compatibility, populate the deprecated Data field.\n\t\tevent.Action.Interrupted.Data = &WorkflowInterruptInfo{\n\t\t\tOrigInput:             getRunCtx(ctx).RootInput,\n\t\t\tParallelInterruptInfo: dataMap,\n\t\t}\n\t\tevent.AgentName = a.Name(ctx)\n\t\tevent.RunPath = getRunCtx(ctx).RunPath\n\n\t\tgenerator.Send(event)\n\t}\n\n\treturn nil\n}\n\ntype SequentialAgentConfig struct {\n\tName        string\n\tDescription string\n\tSubAgents   []Agent\n}\n\ntype ParallelAgentConfig struct {\n\tName        string\n\tDescription string\n\tSubAgents   []Agent\n}\n\ntype LoopAgentConfig struct {\n\tName        string\n\tDescription string\n\tSubAgents   []Agent\n\n\tMaxIterations int\n}\n\nfunc newWorkflowAgent(ctx context.Context, name, desc string,\n\tsubAgents []Agent, mode workflowAgentMode, maxIterations int) (*flowAgent, error) {\n\n\twa := &workflowAgent{\n\t\tname:        name,\n\t\tdescription: desc,\n\t\tmode:        mode,\n\n\t\tmaxIterations: maxIterations,\n\t}\n\n\tfas := make([]Agent, len(subAgents))\n\tfor i, subAgent := range subAgents {\n\t\tfas[i] = toFlowAgent(ctx, subAgent, WithDisallowTransferToParent())\n\t}\n\n\tfa, err := setSubAgents(ctx, wa, fas)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\twa.subAgents = fa.subAgents\n\n\treturn fa, nil\n}\n\n// NewSequentialAgent creates an agent that runs sub-agents sequentially.\nfunc NewSequentialAgent(ctx context.Context, config *SequentialAgentConfig) (ResumableAgent, error) {\n\treturn newWorkflowAgent(ctx, config.Name, config.Description, config.SubAgents, workflowAgentModeSequential, 0)\n}\n\n// NewParallelAgent creates an agent that runs sub-agents in parallel.\nfunc NewParallelAgent(ctx context.Context, config *ParallelAgentConfig) (ResumableAgent, error) {\n\treturn newWorkflowAgent(ctx, config.Name, config.Description, config.SubAgents, workflowAgentModeParallel, 0)\n}\n\n// NewLoopAgent creates an agent that loops over sub-agents with a max iteration limit.\nfunc NewLoopAgent(ctx context.Context, config *LoopAgentConfig) (ResumableAgent, error) {\n\treturn newWorkflowAgent(ctx, config.Name, config.Description, config.SubAgents, workflowAgentModeLoop, config.MaxIterations)\n}\n"
  },
  {
    "path": "adk/workflow_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// mockAgent is a simple implementation of the Agent interface for testing\ntype mockAgent struct {\n\tname        string\n\tdescription string\n\tresponses   []*AgentEvent\n}\n\nfunc (a *mockAgent) Name(_ context.Context) string {\n\treturn a.name\n}\n\nfunc (a *mockAgent) Description(_ context.Context) string {\n\treturn a.description\n}\n\nfunc (a *mockAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\titerator, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\tgo func() {\n\t\tdefer generator.Close()\n\n\t\tfor _, event := range a.responses {\n\t\t\tgenerator.Send(event)\n\n\t\t\t// If the event has an Exit action, stop sending events\n\t\t\tif event.Action != nil && event.Action.Exit {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}()\n\n\treturn iterator\n}\n\n// newMockAgent creates a new mock agent with the given name, description, and responses\nfunc newMockAgent(name, description string, responses []*AgentEvent) *mockAgent {\n\treturn &mockAgent{\n\t\tname:        name,\n\t\tdescription: description,\n\t\tresponses:   responses,\n\t}\n}\n\n// TestSequentialAgent tests the sequential workflow agent\nfunc TestSequentialAgent(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create mock agents with predefined responses\n\tagent1 := newMockAgent(\"Agent1\", \"First agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"Agent1\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Response from Agent1\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\tagent2 := newMockAgent(\"Agent2\", \"Second agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"Agent2\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Response from Agent2\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t}},\n\t})\n\n\t// Create a sequential agent with the mock agents\n\tconfig := &SequentialAgentConfig{\n\t\tName:        \"SequentialTestAgent\",\n\t\tDescription: \"Test sequential agent\",\n\t\tSubAgents:   []Agent{agent1, agent2},\n\t}\n\n\tsequentialAgent, err := NewSequentialAgent(ctx, config)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, sequentialAgent)\n\n\tassert.Equal(t, \"Test sequential agent\", sequentialAgent.Description(ctx))\n\n\t// Run the sequential agent\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Test input\"),\n\t\t},\n\t}\n\n\t// Initialize the run context\n\tctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input)\n\n\titerator := sequentialAgent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// First event should be from agent1\n\tevent1, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event1)\n\tassert.Nil(t, event1.Err)\n\tassert.NotNil(t, event1.Output)\n\tassert.NotNil(t, event1.Output.MessageOutput)\n\n\t// Get the message content from agent1\n\tmsg1 := event1.Output.MessageOutput.Message\n\tassert.NotNil(t, msg1)\n\tassert.Equal(t, \"Response from Agent1\", msg1.Content)\n\n\t// Second event should be from agent2\n\tevent2, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event2)\n\tassert.Nil(t, event2.Err)\n\tassert.NotNil(t, event2.Output)\n\tassert.NotNil(t, event2.Output.MessageOutput)\n\n\t// Get the message content from agent2\n\tmsg2 := event2.Output.MessageOutput.Message\n\tassert.NotNil(t, msg2)\n\tassert.Equal(t, \"Response from Agent2\", msg2.Content)\n\n\t// No more events\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\n// TestSequentialAgentWithExit tests the sequential workflow agent with an exit action\nfunc TestSequentialAgentWithExit(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create mock agents with predefined responses\n\tagent1 := newMockAgent(\"Agent1\", \"First agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"Agent1\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Response from Agent1\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t},\n\t\t\tAction: &AgentAction{\n\t\t\t\tExit: true,\n\t\t\t},\n\t\t},\n\t})\n\n\tagent2 := newMockAgent(\"Agent2\", \"Second agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"Agent2\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Response from Agent2\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\t// Create a sequential agent with the mock agents\n\tconfig := &SequentialAgentConfig{\n\t\tName:        \"SequentialTestAgent\",\n\t\tDescription: \"Test sequential agent\",\n\t\tSubAgents:   []Agent{agent1, agent2},\n\t}\n\n\tsequentialAgent, err := NewSequentialAgent(ctx, config)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, sequentialAgent)\n\n\t// Run the sequential agent\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Test input\"),\n\t\t},\n\t}\n\n\tctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input)\n\n\titerator := sequentialAgent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// First event should be from agent1 with exit action\n\tevent1, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event1)\n\tassert.Nil(t, event1.Err)\n\tassert.NotNil(t, event1.Output)\n\tassert.NotNil(t, event1.Output.MessageOutput)\n\tassert.NotNil(t, event1.Action)\n\tassert.True(t, event1.Action.Exit)\n\n\t// No more events due to exit action\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\n// TestParallelAgent tests the parallel workflow agent\nfunc TestParallelAgent(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create mock agents with predefined responses\n\tagent1 := newMockAgent(\"Agent1\", \"First agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"Agent1\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Response from Agent1\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\tagent2 := newMockAgent(\"Agent2\", \"Second agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"Agent2\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Response from Agent2\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\t// Create a parallel agent with the mock agents\n\tconfig := &ParallelAgentConfig{\n\t\tName:        \"ParallelTestAgent\",\n\t\tDescription: \"Test parallel agent\",\n\t\tSubAgents:   []Agent{agent1, agent2},\n\t}\n\n\tparallelAgent, err := NewParallelAgent(ctx, config)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, parallelAgent)\n\n\t// Run the parallel agent\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Test input\"),\n\t\t},\n\t}\n\n\tctx, _ = initRunCtx(ctx, parallelAgent.Name(ctx), input)\n\n\titerator := parallelAgent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// Collect all events\n\tvar events []*AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\t// Should have two events, one from each agent\n\tassert.Equal(t, 2, len(events))\n\n\t// Verify the events\n\tfor _, event := range events {\n\t\tassert.Nil(t, event.Err)\n\t\tassert.NotNil(t, event.Output)\n\t\tassert.NotNil(t, event.Output.MessageOutput)\n\n\t\tmsg := event.Output.MessageOutput.Message\n\t\tassert.NotNil(t, msg)\n\t\tassert.NoError(t, err)\n\n\t\t// Check the source agent name and message content\n\t\tif event.AgentName == \"Agent1\" {\n\t\t\tassert.Equal(t, \"Response from Agent1\", msg.Content)\n\t\t} else if event.AgentName == \"Agent2\" {\n\t\t\tassert.Equal(t, \"Response from Agent2\", msg.Content)\n\t\t} else {\n\t\t\tt.Fatalf(\"Unexpected source agent name: %s\", event.AgentName)\n\t\t}\n\t}\n}\n\n// TestLoopAgent tests the loop workflow agent\nfunc TestLoopAgent(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock agent that will be called multiple times\n\tagent := newMockAgent(\"LoopAgent\", \"Loop agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"LoopAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Loop iteration\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\t// Create a loop agent with the mock agent and max iterations set to 3\n\tconfig := &LoopAgentConfig{\n\t\tName:        \"LoopTestAgent\",\n\t\tDescription: \"Test loop agent\",\n\t\tSubAgents:   []Agent{agent},\n\n\t\tMaxIterations: 3,\n\t}\n\n\tloopAgent, err := NewLoopAgent(ctx, config)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, loopAgent)\n\n\t// Run the loop agent\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Test input\"),\n\t\t},\n\t}\n\n\tctx, _ = initRunCtx(ctx, loopAgent.Name(ctx), input)\n\n\titerator := loopAgent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// Collect all events\n\tvar events []*AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\t// Should have 3 events (one for each iteration)\n\tassert.Equal(t, 3, len(events))\n\n\t// Verify all events\n\tfor _, event := range events {\n\t\tassert.Nil(t, event.Err)\n\t\tassert.NotNil(t, event.Output)\n\t\tassert.NotNil(t, event.Output.MessageOutput)\n\n\t\tmsg := event.Output.MessageOutput.Message\n\t\tassert.NotNil(t, msg)\n\t\tassert.Equal(t, \"Loop iteration\", msg.Content)\n\t}\n}\n\n// TestLoopAgentWithBreakLoop tests the loop workflow agent with an break loop action\nfunc TestLoopAgentWithBreakLoop(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a mock agent that will break the loop after the first iteration\n\tagent := newMockAgent(\"LoopAgent\", \"Loop agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"LoopAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Loop iteration with break loop\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t},\n\t\t\tAction: NewBreakLoopAction(\"LoopAgent\"),\n\t\t},\n\t})\n\n\t// Create a loop agent with the mock agent and max iterations set to 3\n\tconfig := &LoopAgentConfig{\n\t\tName:          \"LoopTestAgent\",\n\t\tDescription:   \"Test loop agent\",\n\t\tSubAgents:     []Agent{agent},\n\t\tMaxIterations: 3,\n\t}\n\n\tloopAgent, err := NewLoopAgent(ctx, config)\n\tassert.NoError(t, err)\n\tassert.NotNil(t, loopAgent)\n\n\t// Run the loop agent\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Test input\"),\n\t\t},\n\t}\n\tctx, _ = initRunCtx(ctx, loopAgent.Name(ctx), input)\n\n\titerator := loopAgent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// Collect all events\n\tvar events []*AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\t// Should have only 1 event due to break loop action\n\tassert.Equal(t, 1, len(events))\n\n\t// Verify the event\n\tevent := events[0]\n\tassert.Nil(t, event.Err)\n\tassert.NotNil(t, event.Output)\n\tassert.NotNil(t, event.Output.MessageOutput)\n\tassert.NotNil(t, event.Action)\n\tassert.NotNil(t, event.Action.BreakLoop)\n\tassert.True(t, event.Action.BreakLoop.Done)\n\tassert.Equal(t, \"LoopAgent\", event.Action.BreakLoop.From)\n\tassert.Equal(t, 0, event.Action.BreakLoop.CurrentIterations)\n\n\tmsg := event.Output.MessageOutput.Message\n\tassert.NotNil(t, msg)\n\tassert.Equal(t, \"Loop iteration with break loop\", msg.Content)\n}\n\n// Add these test functions to the existing workflow_test.go file\n\n// Replace the existing TestWorkflowAgentPanicRecovery function\nfunc TestWorkflowAgentPanicRecovery(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a panic agent that panics in Run method\n\tpanicAgent := &panicMockAgent{\n\t\tmockAgent: mockAgent{\n\t\t\tname:        \"PanicAgent\",\n\t\t\tdescription: \"Agent that panics\",\n\t\t\tresponses:   []*AgentEvent{},\n\t\t},\n\t}\n\n\t// Create a sequential agent with the panic agent\n\tconfig := &SequentialAgentConfig{\n\t\tName:        \"PanicTestAgent\",\n\t\tDescription: \"Test agent with panic\",\n\t\tSubAgents:   []Agent{panicAgent},\n\t}\n\n\tsequentialAgent, err := NewSequentialAgent(ctx, config)\n\tassert.NoError(t, err)\n\n\t// Run the agent and expect panic recovery\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Test input\"),\n\t\t},\n\t}\n\n\tctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input)\n\titerator := sequentialAgent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// Should receive an error event due to panic recovery\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.NotNil(t, event.Err)\n\tassert.Contains(t, event.Err.Error(), \"panic\")\n\n\t// No more events\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\n// Add these new mock agent types that properly panic\ntype panicMockAgent struct {\n\tmockAgent\n}\n\nfunc (a *panicMockAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\tpanic(\"test panic in agent\")\n}\n\nfunc TestParallelWorkflowResumeWithEvents(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create interruptible agents\n\tsa1 := &myAgent{\n\t\tname: \"sa1\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\t// Send a normal message event first, called event1\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"sa1\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa1 normal message\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tintEvent := Interrupt(ctx, \"sa1 interrupt data\")\n\t\t\tgenerator.Send(intEvent)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.True(t, info.WasInterrupted)\n\t\t\tassert.Nil(t, info.InterruptState)\n\t\t\tassert.True(t, info.IsResumeTarget)\n\t\t\tassert.Equal(t, \"resume sa1\", info.ResumeData)\n\n\t\t\t// Get the events from session and verify visibility\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tassert.NotNil(t, runCtx.Session, \"sa1 resumer should have session\")\n\t\t\tallEvents := runCtx.Session.getEvents()\n\n\t\t\t// Assert that allEvents only have 1 event, that is event1\n\t\t\tassert.Equal(t, 1, len(allEvents), \"sa1 should only see its own event in session\")\n\t\t\tassert.Equal(t, \"sa1\", allEvents[0].AgentEvent.AgentName, \"sa1 should see its own event\")\n\t\t\tassert.Equal(t, \"sa1 normal message\", allEvents[0].AgentEvent.Output.MessageOutput.Message.Content, \"sa1 should see its own message content\")\n\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tsa2 := &myAgent{\n\t\tname: \"sa2\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\t// Send a normal message event first, called event2\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"sa2\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa2 normal message\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tintEvent := StatefulInterrupt(ctx, \"sa2 interrupt data\", \"sa2 interrupt\")\n\t\t\tgenerator.Send(intEvent)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.True(t, info.WasInterrupted)\n\t\t\tassert.NotNil(t, info.InterruptState)\n\t\t\tassert.Equal(t, \"sa2 interrupt\", info.InterruptState)\n\t\t\tassert.True(t, info.IsResumeTarget)\n\t\t\tassert.Equal(t, \"resume sa2\", info.ResumeData)\n\n\t\t\t// Get the events from session and verify visibility\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tassert.NotNil(t, runCtx.Session, \"sa2 resumer should have session\")\n\t\t\tallEvents := runCtx.Session.getEvents()\n\n\t\t\t// Assert that allEvents only have 1 event, that is event2\n\t\t\tassert.Equal(t, 1, len(allEvents), \"sa2 should only see its own event in session\")\n\t\t\tassert.Equal(t, \"sa2\", allEvents[0].AgentEvent.AgentName, \"sa2 should see its own event\")\n\t\t\tassert.Equal(t, \"sa2 normal message\", allEvents[0].AgentEvent.Output.MessageOutput.Message.Content, \"sa2 should see its own message content\")\n\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tsa3 := &myAgent{\n\t\tname: \"sa3\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"sa3\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa3 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tsa4 := &myAgent{\n\t\tname: \"sa4\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"sa4\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"sa4 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tt.Run(\"test parallel workflow agent\", func(t *testing.T) {\n\t\t// parallel\n\t\ta, err := NewParallelAgent(ctx, &ParallelAgentConfig{\n\t\t\tName:      \"parallel agent\",\n\t\t\tSubAgents: []Agent{sa1, sa2, sa3, sa4},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\trunner := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           a,\n\t\t\tCheckPointStore: newMyStore(),\n\t\t})\n\t\titer := runner.Query(ctx, \"hello world\", WithCheckPointID(\"1\"))\n\t\tvar (\n\t\t\tevents         []*AgentEvent\n\t\t\tinterruptEvent *AgentEvent\n\t\t)\n\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\t\tinterruptEvent = event\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tevents = append(events, event)\n\t\t}\n\t\tassert.Equal(t, 4, len(events), \"should have 4 events (2 normal messages + 2 completed agents)\")\n\n\t\t// Verify specific properties of each event\n\t\tvar sa3Event, sa4Event *AgentEvent\n\t\tfor _, event := range events {\n\t\t\tif event.AgentName == \"sa3\" {\n\t\t\t\tsa3Event = event\n\t\t\t} else if event.AgentName == \"sa4\" {\n\t\t\t\tsa4Event = event\n\t\t\t}\n\t\t}\n\n\t\t// Verify sa3 event properties\n\t\tassert.NotNil(t, sa3Event, \"should have event from sa3\")\n\t\tassert.Equal(t, \"sa3\", sa3Event.AgentName, \"sa3 event should have correct agent name\")\n\t\tassert.Equal(t, []RunStep{{\"parallel agent\"}, {\"sa3\"}}, sa3Event.RunPath, \"sa3 event should have correct run path\")\n\t\tassert.NotNil(t, sa3Event.Output, \"sa3 event should have output\")\n\t\tassert.NotNil(t, sa3Event.Output.MessageOutput, \"sa3 event should have message output\")\n\t\tassert.Equal(t, \"sa3 completed\", sa3Event.Output.MessageOutput.Message.Content, \"sa3 event should have correct message content\")\n\n\t\t// Verify sa4 event properties\n\t\tassert.NotNil(t, sa4Event, \"should have event from sa4\")\n\t\tassert.Equal(t, \"sa4\", sa4Event.AgentName, \"sa4 event should have correct agent name\")\n\t\tassert.Equal(t, []RunStep{{\"parallel agent\"}, {\"sa4\"}}, sa4Event.RunPath, \"sa4 event should have correct run path\")\n\t\tassert.NotNil(t, sa4Event.Output, \"sa4 event should have output\")\n\t\tassert.NotNil(t, sa4Event.Output.MessageOutput, \"sa4 event should have message output\")\n\t\tassert.Equal(t, \"sa4 completed\", sa4Event.Output.MessageOutput.Message.Content, \"sa4 event should have correct message content\")\n\n\t\tassert.NotNil(t, interruptEvent)\n\t\tassert.Equal(t, \"parallel agent\", interruptEvent.AgentName)\n\t\tassert.Equal(t, []RunStep{{\"parallel agent\"}}, interruptEvent.RunPath)\n\t\tassert.NotNil(t, interruptEvent.Action.Interrupted)\n\n\t\tvar sa1InfoFound, sa2InfoFound bool\n\t\tfor _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\t\tif ctx.Info == \"sa1 interrupt data\" {\n\t\t\t\tsa1InfoFound = true\n\t\t\t} else if ctx.Info == \"sa2 interrupt data\" {\n\t\t\t\tsa2InfoFound = true\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts))\n\t\tassert.True(t, sa1InfoFound)\n\t\tassert.True(t, sa2InfoFound)\n\n\t\tvar parallelInterruptID1, parallelInterruptID2 string\n\t\tfor _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\t\tif ctx.Info == \"sa1 interrupt data\" {\n\t\t\t\tparallelInterruptID1 = ctx.ID\n\t\t\t} else if ctx.Info == \"sa2 interrupt data\" {\n\t\t\t\tparallelInterruptID2 = ctx.ID\n\t\t\t}\n\t\t}\n\t\tassert.NotEmpty(t, parallelInterruptID1)\n\t\tassert.NotEmpty(t, parallelInterruptID2)\n\n\t\titer, err = runner.ResumeWithParams(ctx, \"1\", &ResumeParams{\n\t\t\tTargets: map[string]any{\n\t\t\t\tparallelInterruptID1: \"resume sa1\",\n\t\t\t\tparallelInterruptID2: \"resume sa2\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\t_, ok := iter.Next()\n\t\tassert.False(t, ok)\n\t})\n}\n\nfunc TestNestedParallelWorkflow(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create predecessor agent that runs before the parallel structure\n\tpredecessorAgent := &myAgent{\n\t\tname: \"predecessor\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"predecessor\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"predecessor completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\t// Create interruptible inner agents\n\tinnerAgent1 := &myAgent{\n\t\tname: \"inner1\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\t\t\t// Verify inner1 can see predecessor's event\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tallEvents := runCtx.Session.getEvents()\n\t\t\tassert.Equal(t, 1, len(allEvents), \"inner1 should see exactly 1 event (predecessor)\")\n\n\t\t\tassert.Equal(t, \"predecessor\", allEvents[0].AgentEvent.AgentName, \"inner1 should see predecessor event\")\n\t\t\tassert.Equal(t, \"predecessor completed\", allEvents[0].AgentEvent.Output.MessageOutput.Message.Content, \"inner1 should see predecessor message content\")\n\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"inner1\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"inner1 normal\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tintEvent := Interrupt(ctx, \"inner1 interrupt\")\n\t\t\tgenerator.Send(intEvent)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.True(t, info.WasInterrupted)\n\t\t\tassert.Equal(t, \"resume inner1\", info.ResumeData)\n\n\t\t\t// Verify inner1 can see predecessor's event during resume\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tallEvents := runCtx.Session.getEvents()\n\t\t\tassert.Equal(t, 2, len(allEvents), \"inner1 should see exactly 2 events (predecessor + own normal message) during resume\")\n\n\t\t\t// Find and verify predecessor event\n\t\t\tvar foundPredecessor bool\n\t\t\tfor _, event := range allEvents {\n\t\t\t\tif event.AgentEvent != nil && event.AgentEvent.AgentName == \"predecessor\" {\n\t\t\t\t\tfoundPredecessor = true\n\t\t\t\t\tassert.Equal(t, \"predecessor completed\", event.AgentEvent.Output.MessageOutput.Message.Content)\n\t\t\t\t}\n\t\t\t}\n\t\t\tassert.True(t, foundPredecessor, \"inner1 should see predecessor event during resume\")\n\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tinnerAgent2 := &myAgent{\n\t\tname: \"inner2\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\t\t\t// Verify inner2 can see predecessor's event\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tallEvents := runCtx.Session.getEvents()\n\t\t\tassert.Equal(t, 1, len(allEvents), \"inner2 should see exactly 1 event (predecessor)\")\n\n\t\t\tassert.Equal(t, \"predecessor\", allEvents[0].AgentEvent.AgentName, \"inner2 should see predecessor event\")\n\t\t\tassert.Equal(t, \"predecessor completed\", allEvents[0].AgentEvent.Output.MessageOutput.Message.Content, \"inner2 should see predecessor message content\")\n\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"inner2\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"inner2 normal\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tintEvent := StatefulInterrupt(ctx, \"inner2 interrupt\", \"inner2 state\")\n\t\t\tgenerator.Send(intEvent)\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\tassert.True(t, info.WasInterrupted)\n\t\t\tassert.Equal(t, \"inner2 state\", info.InterruptState)\n\t\t\tassert.Equal(t, \"resume inner2\", info.ResumeData)\n\n\t\t\t// Verify inner2 can see predecessor's event during resume\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tallEvents := runCtx.Session.getEvents()\n\t\t\tassert.Equal(t, 2, len(allEvents), \"inner2 should see exactly 2 events (predecessor + own normal message) during resume\")\n\n\t\t\t// Find and verify predecessor event\n\t\t\tvar foundPredecessor bool\n\t\t\tfor _, event := range allEvents {\n\t\t\t\tif event.AgentEvent != nil && event.AgentEvent.AgentName == \"predecessor\" {\n\t\t\t\t\tfoundPredecessor = true\n\t\t\t\t\tassert.Equal(t, \"predecessor completed\", event.AgentEvent.Output.MessageOutput.Message.Content)\n\t\t\t\t}\n\t\t\t}\n\t\t\tassert.True(t, foundPredecessor, \"inner2 should see predecessor event during resume\")\n\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\t// Create inner parallel workflow\n\tinnerParallel, err := NewParallelAgent(ctx, &ParallelAgentConfig{\n\t\tName:      \"inner parallel\",\n\t\tSubAgents: []Agent{innerAgent1, innerAgent2},\n\t})\n\tassert.NoError(t, err)\n\n\t// Create simple outer agents\n\touterAgent1 := &myAgent{\n\t\tname: \"outer1\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"outer1\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"outer1 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\touterAgent2 := &myAgent{\n\t\tname: \"outer2\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"outer2\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"outer2 completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\t// Create outer parallel workflow with nested parallel agent\n\touterParallel, err := NewParallelAgent(ctx, &ParallelAgentConfig{\n\t\tName:      \"outer parallel\",\n\t\tSubAgents: []Agent{outerAgent1, innerParallel, outerAgent2},\n\t})\n\tassert.NoError(t, err)\n\n\t// Create successor agent that runs after the parallel structure\n\tsuccessorAgent := &myAgent{\n\t\tname: \"successor\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\n\t\t\t// Verify successor can see all events from predecessor and parallel agents\n\t\t\trunCtx := getRunCtx(ctx)\n\t\t\tallEvents := runCtx.Session.getEvents()\n\t\t\tassert.GreaterOrEqual(t, len(allEvents), 5, \"successor should see all events\")\n\n\t\t\tvar foundPredecessor, foundOuter1, foundOuter2, foundInner1, foundInner2 bool\n\t\t\tfor _, event := range allEvents {\n\t\t\t\tif event.AgentEvent != nil {\n\t\t\t\t\tswitch event.AgentEvent.AgentName {\n\t\t\t\t\tcase \"predecessor\":\n\t\t\t\t\t\tfoundPredecessor = true\n\t\t\t\t\t\tassert.Equal(t, \"predecessor completed\", event.AgentEvent.Output.MessageOutput.Message.Content)\n\t\t\t\t\tcase \"outer1\":\n\t\t\t\t\t\tfoundOuter1 = true\n\t\t\t\t\tcase \"outer2\":\n\t\t\t\t\t\tfoundOuter2 = true\n\t\t\t\t\tcase \"inner1\":\n\t\t\t\t\t\tfoundInner1 = true\n\t\t\t\t\tcase \"inner2\":\n\t\t\t\t\t\tfoundInner2 = true\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tassert.True(t, foundPredecessor, \"successor should see predecessor event\")\n\t\t\tassert.True(t, foundOuter1, \"successor should see outer1 event\")\n\t\t\tassert.True(t, foundOuter2, \"successor should see outer2 event\")\n\t\t\tassert.True(t, foundInner1, \"successor should see inner1 event\")\n\t\t\tassert.True(t, foundInner2, \"successor should see inner2 event\")\n\n\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\tAgentName: \"successor\",\n\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\tMessage: schema.UserMessage(\"successor completed\"),\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t})\n\t\t\tgenerator.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\t// Create sequential workflow: predecessor -> parallel -> successor\n\tsequentialWorkflow, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\tName:      \"sequential workflow\",\n\t\tSubAgents: []Agent{predecessorAgent, outerParallel, successorAgent},\n\t})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           sequentialWorkflow,\n\t\tCheckPointStore: newMyStore(),\n\t})\n\n\titer := runner.Query(ctx, \"test nested parallel with predecessor and successor\", WithCheckPointID(\"nested-parallel-test\"))\n\n\tvar events []*AgentEvent\n\tvar interruptEvent *AgentEvent\n\tfor event, ok := iter.Next(); ok; event, ok = iter.Next() {\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tinterruptEvent = event\n\t\t\tcontinue\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\t// Should get events from predecessor, outer agents, and inner normal messages (successor doesn't run due to interruption)\n\tassert.Equal(t, 5, len(events), \"should have 5 events (predecessor + 2 outer + 2 inner)\")\n\tif interruptEvent == nil {\n\t\tt.Fatal(\"should have interrupt event\")\n\t}\n\n\t// Resume the inner parallel workflow\n\tvar innerInterruptID1, innerInterruptID2 string\n\tfor _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif ctx.Info == \"inner1 interrupt\" {\n\t\t\tinnerInterruptID1 = ctx.ID\n\t\t} else if ctx.Info == \"inner2 interrupt\" {\n\t\t\tinnerInterruptID2 = ctx.ID\n\t\t}\n\t}\n\n\titer, err = runner.ResumeWithParams(ctx, \"nested-parallel-test\", &ResumeParams{\n\t\tTargets: map[string]any{\n\t\t\tinnerInterruptID1: \"resume inner1\",\n\t\t\tinnerInterruptID2: \"resume inner2\",\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\t// Verify resume completes successfully and successor runs\n\tvar resumeEvents []*AgentEvent\n\tfor event, ok := iter.Next(); ok; event, ok = iter.Next() {\n\t\tresumeEvents = append(resumeEvents, event)\n\t}\n\n\t// Should get successor event after resume\n\tassert.Equal(t, 1, len(resumeEvents), \"should have successor event after resume\")\n\tassert.Equal(t, \"successor\", resumeEvents[0].AgentName)\n}\n\n// TestWorkflowAgentUnsupportedMode tests unsupported workflow mode error (lines 65-71)\nfunc TestWorkflowAgentUnsupportedMode(t *testing.T) {\n\tctx := context.Background()\n\n\t// Create a workflow agent with unsupported mode\n\tagent := &workflowAgent{\n\t\tname:        \"UnsupportedModeAgent\",\n\t\tdescription: \"Agent with unsupported mode\",\n\t\tsubAgents:   []*flowAgent{},\n\t\tmode:        workflowAgentMode(999), // Invalid mode\n\t}\n\n\t// Run the agent and expect error\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Test input\"),\n\t\t},\n\t}\n\n\tctx, _ = initRunCtx(ctx, agent.Name(ctx), input)\n\titerator := agent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\t// Should receive an error event due to unsupported mode\n\tevent, ok := iterator.Next()\n\tassert.True(t, ok)\n\tassert.NotNil(t, event)\n\tassert.NotNil(t, event.Err)\n\tassert.Contains(t, event.Err.Error(), \"unsupported workflow agent mode\")\n\n\t// No more events\n\t_, ok = iterator.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestFilterOptions(t *testing.T) {\n\ta1 := &myAgent{\n\t\tname: \"Agent1\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\to := GetImplSpecificOptions[myAgentOptions](nil, opts...)\n\t\t\tassert.Equal(t, \"Agent1\", o.value)\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgen.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\ta2 := &myAgent{\n\t\tname: \"Agent2\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\to := GetImplSpecificOptions[myAgentOptions](nil, opts...)\n\t\t\tassert.Equal(t, \"Agent2\", o.value)\n\t\t\titer, gen := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgen.Close()\n\t\t\treturn iter\n\t\t},\n\t}\n\tctx := context.Background()\n\t// sequential\n\tseqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{\n\t\tSubAgents: []Agent{a1, a2},\n\t})\n\tassert.NoError(t, err)\n\titer := seqAgent.Run(ctx, &AgentInput{}, withValue(\"Agent1\").DesignateAgent(\"Agent1\"), withValue(\"Agent2\").DesignateAgent(\"Agent2\"))\n\t_, ok := iter.Next()\n\tassert.False(t, ok)\n\n\t// parallel\n\tparAgent, err := NewParallelAgent(ctx, &ParallelAgentConfig{\n\t\tSubAgents: []Agent{a1, a2},\n\t})\n\tassert.NoError(t, err)\n\titer = parAgent.Run(ctx, &AgentInput{}, withValue(\"Agent1\").DesignateAgent(\"Agent1\"), withValue(\"Agent2\").DesignateAgent(\"Agent2\"))\n\t_, ok = iter.Next()\n\tassert.False(t, ok)\n}\n\nfunc TestLoopAgentWithError(t *testing.T) {\n\tctx := context.Background()\n\n\titerationCount := 0\n\tagent := &myAgent{\n\t\tname: \"ErrorAgent\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer generator.Close()\n\t\t\t\titerationCount++\n\t\t\t\tif iterationCount == 3 {\n\t\t\t\t\tgenerator.Send(&AgentEvent{Err: fmt.Errorf(\"error on iteration %d\", iterationCount)})\n\t\t\t\t\treturn\n\t\t\t\t}\n\t\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(fmt.Sprintf(\"iteration %d\", iterationCount), nil),\n\t\t\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tloopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{\n\t\tName:          \"LoopErrorTestAgent\",\n\t\tSubAgents:     []Agent{agent},\n\t\tMaxIterations: 10,\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &AgentInput{Messages: []Message{schema.UserMessage(\"test\")}}\n\tctx, _ = initRunCtx(ctx, loopAgent.Name(ctx), input)\n\titerator := loopAgent.Run(ctx, input)\n\n\tvar events []*AgentEvent\n\tvar errorEvent *AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif event.Err != nil {\n\t\t\terrorEvent = event\n\t\t} else {\n\t\t\tevents = append(events, event)\n\t\t}\n\t}\n\n\tassert.Equal(t, 2, len(events), \"should have 2 successful iterations before error\")\n\tassert.NotNil(t, errorEvent, \"should have received error event\")\n\tassert.Contains(t, errorEvent.Err.Error(), \"error on iteration 3\")\n\tassert.Equal(t, 3, iterationCount, \"loop should stop at iteration 3\")\n}\n\nfunc TestWorkflowCallbackHandlerNotDoubled(t *testing.T) {\n\tctx := context.Background()\n\tstore := newMyStore()\n\n\tvar globalCallbackCount int\n\tvar designatedCallbackCount int\n\tvar mu sync.Mutex\n\n\tglobalHandler := callbacks.NewHandlerBuilder().OnStartFn(\n\t\tfunc(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component == ComponentOfAgent && info.Name == \"SubSubAgent\" {\n\t\t\t\tmu.Lock()\n\t\t\t\tglobalCallbackCount++\n\t\t\t\tmu.Unlock()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).Build()\n\n\tdesignatedHandler := callbacks.NewHandlerBuilder().OnStartFn(\n\t\tfunc(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Component == ComponentOfAgent && info.Name == \"SubSubAgent\" {\n\t\t\t\tmu.Lock()\n\t\t\t\tdesignatedCallbackCount++\n\t\t\t\tmu.Unlock()\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).Build()\n\n\titerationCount := 0\n\tshouldInterrupt := true\n\tsubSubAgent := &myAgent{\n\t\tname: \"SubSubAgent\",\n\t\trunFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer generator.Close()\n\t\t\t\titerationCount++\n\t\t\t\tif shouldInterrupt && iterationCount == 2 {\n\t\t\t\t\tgenerator.Send(Interrupt(ctx, \"test_interrupt\"))\n\t\t\t\t\treturn\n\t\t\t\t}\n\t\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(fmt.Sprintf(\"iteration %d\", iterationCount), nil),\n\t\t\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t\tresumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {\n\t\t\titer, generator := NewAsyncIteratorPair[*AgentEvent]()\n\t\t\tgo func() {\n\t\t\t\tdefer generator.Close()\n\t\t\t\titerationCount++\n\t\t\t\tgenerator.Send(&AgentEvent{\n\t\t\t\t\tOutput: &AgentOutput{\n\t\t\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\t\t\tMessage: schema.AssistantMessage(fmt.Sprintf(\"resumed iteration %d\", iterationCount), nil),\n\t\t\t\t\t\t\tRole:    schema.Assistant,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t}()\n\t\t\treturn iter\n\t\t},\n\t}\n\n\tsubWorkflow, err := NewLoopAgent(ctx, &LoopAgentConfig{\n\t\tName:          \"SubWorkflow\",\n\t\tSubAgents:     []Agent{subSubAgent},\n\t\tMaxIterations: 2,\n\t})\n\tassert.NoError(t, err)\n\n\tparentWorkflow, err := NewLoopAgent(ctx, &LoopAgentConfig{\n\t\tName:          \"ParentWorkflow\",\n\t\tSubAgents:     []Agent{subWorkflow},\n\t\tMaxIterations: 2,\n\t})\n\tassert.NoError(t, err)\n\n\trunner := NewRunner(ctx, RunnerConfig{\n\t\tAgent:           parentWorkflow,\n\t\tCheckPointStore: store,\n\t})\n\n\topts := []AgentRunOption{\n\t\tWithCallbacks(globalHandler),\n\t\tWithCallbacks(designatedHandler).DesignateAgent(\"ParentWorkflow\", \"SubSubAgent\"),\n\t\tWithCheckPointID(\"cp1\"),\n\t}\n\n\titerator := runner.Run(ctx, []Message{schema.UserMessage(\"test\")}, opts...)\n\n\tvar interruptEvent *AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tif event.Action != nil && event.Action.Interrupted != nil {\n\t\t\tinterruptEvent = event\n\t\t}\n\t}\n\n\tassert.NotNil(t, interruptEvent)\n\tassert.Equal(t, 2, iterationCount)\n\tassert.Equal(t, 2, globalCallbackCount)\n\tassert.Equal(t, 2, designatedCallbackCount)\n\n\tshouldInterrupt = false\n\tvar rootCauseID string\n\tfor _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {\n\t\tif intCtx.IsRootCause {\n\t\t\trootCauseID = intCtx.ID\n\t\t\tbreak\n\t\t}\n\t}\n\n\tresumeIter, err := runner.ResumeWithParams(ctx, \"cp1\", &ResumeParams{\n\t\tTargets: map[string]any{rootCauseID: nil},\n\t}, opts...)\n\tassert.NoError(t, err)\n\n\tfor {\n\t\t_, ok := resumeIter.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassert.Equal(t, 5, iterationCount)\n\tassert.Equal(t, 5, globalCallbackCount)\n\tassert.Equal(t, 5, designatedCallbackCount)\n}\n\nfunc TestLoopAgentWithBreakLoopFollowedByMoreEvents(t *testing.T) {\n\tctx := context.Background()\n\n\tagent := newMockAgent(\"SubAgent\", \"Sub agent\", []*AgentEvent{\n\t\t{\n\t\t\tAgentName: \"SubAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.ToolMessage(\"tool result\", \"call_123\"),\n\t\t\t\t\tRole:        schema.Tool,\n\t\t\t\t},\n\t\t\t},\n\t\t\tAction: NewBreakLoopAction(\"SubAgent\"),\n\t\t},\n\t\t{\n\t\t\tAgentName: \"SubAgent\",\n\t\t\tOutput: &AgentOutput{\n\t\t\t\tMessageOutput: &MessageVariant{\n\t\t\t\t\tIsStreaming: false,\n\t\t\t\t\tMessage:     schema.AssistantMessage(\"Final response after tool\", nil),\n\t\t\t\t\tRole:        schema.Assistant,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\tloopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{\n\t\tName:          \"LoopTestAgent\",\n\t\tDescription:   \"Test loop agent\",\n\t\tSubAgents:     []Agent{agent},\n\t\tMaxIterations: 3,\n\t})\n\tassert.NoError(t, err)\n\tassert.NotNil(t, loopAgent)\n\n\tinput := &AgentInput{\n\t\tMessages: []Message{\n\t\t\tschema.UserMessage(\"Test input\"),\n\t\t},\n\t}\n\tctx, _ = initRunCtx(ctx, loopAgent.Name(ctx), input)\n\n\titerator := loopAgent.Run(ctx, input)\n\tassert.NotNil(t, iterator)\n\n\tvar events []*AgentEvent\n\tfor {\n\t\tevent, ok := iterator.Next()\n\t\tif !ok {\n\t\t\tbreak\n\t\t}\n\t\tevents = append(events, event)\n\t}\n\n\tassert.Equal(t, 2, len(events), \"should have 2 events (tool event with BreakLoop + final response) and loop should break\")\n\n\tassert.NotNil(t, events[0].Action, \"first event should have an action\")\n\tassert.NotNil(t, events[0].Action.BreakLoop, \"first event should have BreakLoop action\")\n\tassert.True(t, events[0].Action.BreakLoop.Done, \"BreakLoop should be marked as Done\")\n\tassert.Equal(t, \"SubAgent\", events[0].Action.BreakLoop.From)\n\tassert.Equal(t, 0, events[0].Action.BreakLoop.CurrentIterations)\n\tassert.Equal(t, schema.Tool, events[0].Output.MessageOutput.Role, \"first event should be tool message\")\n\n\tassert.Nil(t, events[1].Action, \"second event should not have an action\")\n\tassert.Equal(t, schema.Assistant, events[1].Output.MessageOutput.Role, \"second event should be assistant message\")\n\tassert.Equal(t, \"Final response after tool\", events[1].Output.MessageOutput.Message.Content)\n}\n"
  },
  {
    "path": "adk/wrappers.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype generateEndpoint func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error)\ntype streamEndpoint func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error)\n\ntype modelWrapperConfig struct {\n\thandlers    []ChatModelAgentMiddleware\n\tmiddlewares []AgentMiddleware\n\tretryConfig *ModelRetryConfig\n\ttoolInfos   []*schema.ToolInfo\n}\n\nfunc buildModelWrappers(m model.BaseChatModel, config *modelWrapperConfig) model.BaseChatModel {\n\tvar wrapped model.BaseChatModel = m\n\n\tif !components.IsCallbacksEnabled(m) {\n\t\twrapped = (&callbackInjectionModelWrapper{}).WrapModel(wrapped)\n\t}\n\n\twrapped = &stateModelWrapper{\n\t\tinner:            wrapped,\n\t\toriginal:         m,\n\t\thandlers:         config.handlers,\n\t\tmiddlewares:      config.middlewares,\n\t\ttoolInfos:        config.toolInfos,\n\t\tmodelRetryConfig: config.retryConfig,\n\t}\n\n\treturn wrapped\n}\n\ntype callbackInjectionModelWrapper struct{}\n\nfunc (w *callbackInjectionModelWrapper) WrapModel(m model.BaseChatModel) model.BaseChatModel {\n\treturn &callbackInjectedModel{inner: m}\n}\n\ntype callbackInjectedModel struct {\n\tinner model.BaseChatModel\n}\n\nfunc (m *callbackInjectedModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tctx = callbacks.OnStart(ctx, input)\n\tresult, err := m.inner.Generate(ctx, input, opts...)\n\tif err != nil {\n\t\tcallbacks.OnError(ctx, err)\n\t\treturn nil, err\n\t}\n\tcallbacks.OnEnd(ctx, result)\n\treturn result, nil\n}\n\nfunc (m *callbackInjectedModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tctx = callbacks.OnStart(ctx, input)\n\tresult, err := m.inner.Stream(ctx, input, opts...)\n\tif err != nil {\n\t\tcallbacks.OnError(ctx, err)\n\t\treturn nil, err\n\t}\n\t_, wrappedStream := callbacks.OnEndWithStreamOutput(ctx, result)\n\treturn wrappedStream, nil\n}\n\nfunc handlersToToolMiddlewares(handlers []ChatModelAgentMiddleware) []compose.ToolMiddleware {\n\tvar middlewares []compose.ToolMiddleware\n\tfor i := len(handlers) - 1; i >= 0; i-- {\n\t\thandler := handlers[i]\n\n\t\tm := compose.ToolMiddleware{}\n\n\t\th := handler\n\t\tm.Invokable = func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\t\t\ttCtx := &ToolContext{\n\t\t\t\t\tName:   input.Name,\n\t\t\t\t\tCallID: input.CallID,\n\t\t\t\t}\n\t\t\t\twrappedEndpoint, err := h.WrapInvokableToolCall(\n\t\t\t\t\tctx,\n\t\t\t\t\tfunc(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\t\t\t\t\t\toutput, err := next(ctx, &compose.ToolInput{\n\t\t\t\t\t\t\tName:        input.Name,\n\t\t\t\t\t\t\tCallID:      input.CallID,\n\t\t\t\t\t\t\tArguments:   argumentsInJSON,\n\t\t\t\t\t\t\tCallOptions: opts,\n\t\t\t\t\t\t})\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\treturn \"\", err\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn output.Result, nil\n\t\t\t\t\t},\n\t\t\t\t\ttCtx,\n\t\t\t\t)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tresult, err := wrappedEndpoint(ctx, input.Arguments, input.CallOptions...)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\treturn &compose.ToolOutput{Result: result}, nil\n\t\t\t}\n\t\t}\n\n\t\tm.Streamable = func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\t\t\ttCtx := &ToolContext{\n\t\t\t\t\tName:   input.Name,\n\t\t\t\t\tCallID: input.CallID,\n\t\t\t\t}\n\t\t\t\twrappedEndpoint, err := h.WrapStreamableToolCall(\n\t\t\t\t\tctx,\n\t\t\t\t\tfunc(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\t\t\t\t\t\toutput, err := next(ctx, &compose.ToolInput{\n\t\t\t\t\t\t\tName:        input.Name,\n\t\t\t\t\t\t\tCallID:      input.CallID,\n\t\t\t\t\t\t\tArguments:   argumentsInJSON,\n\t\t\t\t\t\t\tCallOptions: opts,\n\t\t\t\t\t\t})\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn output.Result, nil\n\t\t\t\t\t},\n\t\t\t\t\ttCtx,\n\t\t\t\t)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tresult, err := wrappedEndpoint(ctx, input.Arguments, input.CallOptions...)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\treturn &compose.StreamToolOutput{Result: result}, nil\n\t\t\t}\n\t\t}\n\n\t\tm.EnhancedInvokable = func(next compose.EnhancedInvokableToolEndpoint) compose.EnhancedInvokableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\t\ttCtx := &ToolContext{\n\t\t\t\t\tName:   input.Name,\n\t\t\t\t\tCallID: input.CallID,\n\t\t\t\t}\n\t\t\t\twrappedEndpoint, err := h.WrapEnhancedInvokableToolCall(\n\t\t\t\t\tctx,\n\t\t\t\t\tfunc(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\t\t\t\toutput, err := next(ctx, &compose.ToolInput{\n\t\t\t\t\t\t\tName:        input.Name,\n\t\t\t\t\t\t\tCallID:      input.CallID,\n\t\t\t\t\t\t\tArguments:   toolArgument.Text,\n\t\t\t\t\t\t\tCallOptions: opts,\n\t\t\t\t\t\t})\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn output.Result, nil\n\t\t\t\t\t},\n\t\t\t\t\ttCtx,\n\t\t\t\t)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tresult, err := wrappedEndpoint(ctx, &schema.ToolArgument{Text: input.Arguments}, input.CallOptions...)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\treturn &compose.EnhancedInvokableToolOutput{Result: result}, nil\n\t\t\t}\n\t\t}\n\n\t\tm.EnhancedStreamable = func(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\t\t\ttCtx := &ToolContext{\n\t\t\t\t\tName:   input.Name,\n\t\t\t\t\tCallID: input.CallID,\n\t\t\t\t}\n\t\t\t\twrappedEndpoint, err := h.WrapEnhancedStreamableToolCall(\n\t\t\t\t\tctx,\n\t\t\t\t\tfunc(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\t\t\t\toutput, err := next(ctx, &compose.ToolInput{\n\t\t\t\t\t\t\tName:        input.Name,\n\t\t\t\t\t\t\tCallID:      input.CallID,\n\t\t\t\t\t\t\tArguments:   toolArgument.Text,\n\t\t\t\t\t\t\tCallOptions: opts,\n\t\t\t\t\t\t})\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn output.Result, nil\n\t\t\t\t\t},\n\t\t\t\t\ttCtx,\n\t\t\t\t)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tresult, err := wrappedEndpoint(ctx, &schema.ToolArgument{Text: input.Arguments}, input.CallOptions...)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\treturn &compose.EnhancedStreamableToolOutput{Result: result}, nil\n\t\t\t}\n\t\t}\n\n\t\tmiddlewares = append(middlewares, m)\n\t}\n\treturn middlewares\n}\n\ntype eventSenderModelWrapper struct {\n\t*BaseChatModelAgentMiddleware\n}\n\n// NewEventSenderModelWrapper returns a ChatModelAgentMiddleware that sends model response events.\n// By default, the framework applies this wrapper after all user middlewares, so events contain\n// modified messages. To send events with original (unmodified) output, pass this as a Handler\n// after the modifying middleware (placing it innermost in the wrapper chain).\n// When detected in Handlers, the framework skips the default event sender to avoid duplicates.\nfunc NewEventSenderModelWrapper() ChatModelAgentMiddleware {\n\treturn &eventSenderModelWrapper{\n\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t}\n}\n\nfunc (w *eventSenderModelWrapper) WrapModel(_ context.Context, m model.BaseChatModel, mc *ModelContext) (model.BaseChatModel, error) {\n\tvar retryConfig *ModelRetryConfig\n\tif mc != nil {\n\t\tretryConfig = mc.ModelRetryConfig\n\t}\n\treturn &eventSenderModel{inner: m, modelRetryConfig: retryConfig}, nil\n}\n\ntype eventSenderModel struct {\n\tinner            model.BaseChatModel\n\tmodelRetryConfig *ModelRetryConfig\n}\n\nfunc (m *eventSenderModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tresult, err := m.inner.Generate(ctx, input, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\texecCtx := getChatModelAgentExecCtx(ctx)\n\tif execCtx == nil || execCtx.generator == nil {\n\t\treturn nil, errors.New(\"generator is nil when sending event in Generate: ensure agent state is properly initialized\")\n\t}\n\n\tmsgCopy := *result\n\tevent := EventFromMessage(&msgCopy, nil, schema.Assistant, \"\")\n\texecCtx.send(event)\n\n\treturn result, nil\n}\n\nfunc (m *eventSenderModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tresult, err := m.inner.Stream(ctx, input, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\texecCtx := getChatModelAgentExecCtx(ctx)\n\tif execCtx == nil || execCtx.generator == nil {\n\t\tresult.Close()\n\t\treturn nil, errors.New(\"generator is nil when sending event in Stream: ensure agent state is properly initialized\")\n\t}\n\n\tvar retryAttempt int\n\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tretryAttempt = st.getRetryAttempt()\n\t\treturn nil\n\t})\n\n\tstreams := result.Copy(2)\n\n\teventStream := streams[0]\n\tif m.modelRetryConfig != nil {\n\t\tconvertOpts := []schema.ConvertOption{\n\t\t\tschema.WithErrWrapper(genErrWrapper(ctx, m.modelRetryConfig.MaxRetries,\n\t\t\t\tretryAttempt, m.modelRetryConfig.IsRetryAble)),\n\t\t}\n\t\teventStream = schema.StreamReaderWithConvert(streams[0],\n\t\t\tfunc(msg *schema.Message) (*schema.Message, error) { return msg, nil },\n\t\t\tconvertOpts...)\n\t}\n\n\tevent := EventFromMessage(nil, eventStream, schema.Assistant, \"\")\n\texecCtx.send(event)\n\n\treturn streams[1], nil\n}\n\nfunc popToolGenAction(ctx context.Context, toolName string) *AgentAction {\n\ttoolCallID := compose.GetToolCallID(ctx)\n\n\tvar action *AgentAction\n\t_ = compose.ProcessState(ctx, func(ctx context.Context, st *State) error {\n\t\tif len(toolCallID) > 0 {\n\t\t\tif a := st.popToolGenAction(toolCallID); a != nil {\n\t\t\t\taction = a\n\t\t\t\treturn nil\n\t\t\t}\n\t\t}\n\n\t\tif a := st.popToolGenAction(toolName); a != nil {\n\t\t\taction = a\n\t\t}\n\n\t\treturn nil\n\t})\n\n\treturn action\n}\n\ntype eventSenderToolHandler struct{}\n\nfunc (h *eventSenderToolHandler) WrapInvokableToolCall(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {\n\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\toutput, err := next(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\ttoolName := input.Name\n\t\tcallID := input.CallID\n\n\t\tprePopAction := popToolGenAction(ctx, toolName)\n\t\tmsg := schema.ToolMessage(output.Result, callID, schema.WithToolName(toolName))\n\t\tevent := EventFromMessage(msg, nil, schema.Tool, toolName)\n\t\tif prePopAction != nil {\n\t\t\tevent.Action = prePopAction\n\t\t}\n\n\t\texecCtx := getChatModelAgentExecCtx(ctx)\n\t\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\t\tif st.getReturnDirectlyToolCallID() == callID {\n\t\t\t\tst.setReturnDirectlyEvent(event)\n\t\t\t} else {\n\t\t\t\texecCtx.send(event)\n\t\t\t}\n\t\t\treturn nil\n\t\t})\n\n\t\treturn output, nil\n\t}\n}\n\nfunc (h *eventSenderToolHandler) WrapStreamableToolCall(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {\n\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\toutput, err := next(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\ttoolName := input.Name\n\t\tcallID := input.CallID\n\n\t\tprePopAction := popToolGenAction(ctx, toolName)\n\t\tstreams := output.Result.Copy(2)\n\n\t\tcvt := func(in string) (Message, error) {\n\t\t\treturn schema.ToolMessage(in, callID, schema.WithToolName(toolName)), nil\n\t\t}\n\t\tmsgStream := schema.StreamReaderWithConvert(streams[0], cvt)\n\t\tevent := EventFromMessage(nil, msgStream, schema.Tool, toolName)\n\t\tevent.Action = prePopAction\n\n\t\texecCtx := getChatModelAgentExecCtx(ctx)\n\t\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\t\tif st.getReturnDirectlyToolCallID() == callID {\n\t\t\t\tst.setReturnDirectlyEvent(event)\n\t\t\t} else {\n\t\t\t\texecCtx.send(event)\n\t\t\t}\n\t\t\treturn nil\n\t\t})\n\n\t\treturn &compose.StreamToolOutput{Result: streams[1]}, nil\n\t}\n}\n\nfunc (h *eventSenderToolHandler) WrapEnhancedInvokableToolCall(next compose.EnhancedInvokableToolEndpoint) compose.EnhancedInvokableToolEndpoint {\n\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\toutput, err := next(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\ttoolName := input.Name\n\t\tcallID := input.CallID\n\n\t\tprePopAction := popToolGenAction(ctx, toolName)\n\t\tmsg := schema.ToolMessage(\"\", callID, schema.WithToolName(toolName))\n\t\tmsg.UserInputMultiContent, err = output.Result.ToMessageInputParts()\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tevent := EventFromMessage(msg, nil, schema.Tool, toolName)\n\t\tif prePopAction != nil {\n\t\t\tevent.Action = prePopAction\n\t\t}\n\n\t\texecCtx := getChatModelAgentExecCtx(ctx)\n\t\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\t\tif st.getReturnDirectlyToolCallID() == callID {\n\t\t\t\tst.setReturnDirectlyEvent(event)\n\t\t\t} else {\n\t\t\t\texecCtx.send(event)\n\t\t\t}\n\t\t\treturn nil\n\t\t})\n\n\t\treturn output, nil\n\t}\n}\n\nfunc (h *eventSenderToolHandler) WrapEnhancedStreamableToolCall(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint {\n\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\toutput, err := next(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\ttoolName := input.Name\n\t\tcallID := input.CallID\n\n\t\tprePopAction := popToolGenAction(ctx, toolName)\n\t\tstreams := output.Result.Copy(2)\n\n\t\tcvt := func(in *schema.ToolResult) (Message, error) {\n\t\t\tmsg := schema.ToolMessage(\"\", callID, schema.WithToolName(toolName))\n\t\t\tvar cvtErr error\n\t\t\tmsg.UserInputMultiContent, cvtErr = in.ToMessageInputParts()\n\t\t\tif cvtErr != nil {\n\t\t\t\treturn nil, cvtErr\n\t\t\t}\n\t\t\treturn msg, nil\n\t\t}\n\t\tmsgStream := schema.StreamReaderWithConvert(streams[0], cvt)\n\t\tevent := EventFromMessage(nil, msgStream, schema.Tool, toolName)\n\t\tevent.Action = prePopAction\n\n\t\texecCtx := getChatModelAgentExecCtx(ctx)\n\t\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\t\tif st.getReturnDirectlyToolCallID() == callID {\n\t\t\t\tst.setReturnDirectlyEvent(event)\n\t\t\t} else {\n\t\t\t\texecCtx.send(event)\n\t\t\t}\n\t\t\treturn nil\n\t\t})\n\n\t\treturn &compose.EnhancedStreamableToolOutput{Result: streams[1]}, nil\n\t}\n}\n\ntype stateModelWrapper struct {\n\tinner            model.BaseChatModel\n\toriginal         model.BaseChatModel\n\thandlers         []ChatModelAgentMiddleware\n\tmiddlewares      []AgentMiddleware\n\ttoolInfos        []*schema.ToolInfo\n\tmodelRetryConfig *ModelRetryConfig\n}\n\nfunc (w *stateModelWrapper) IsCallbacksEnabled() bool {\n\treturn true\n}\n\nfunc (w *stateModelWrapper) GetType() string {\n\tif typer, ok := w.original.(components.Typer); ok {\n\t\treturn typer.GetType()\n\t}\n\treturn generic.ParseTypeName(reflect.ValueOf(w.original))\n}\n\nfunc (w *stateModelWrapper) hasUserEventSender() bool {\n\tfor _, handler := range w.handlers {\n\t\tif _, ok := handler.(*eventSenderModelWrapper); ok {\n\t\t\treturn true\n\t\t}\n\t}\n\treturn false\n}\n\nfunc (w *stateModelWrapper) wrapGenerateEndpoint(endpoint generateEndpoint) generateEndpoint {\n\thasUserEventSender := w.hasUserEventSender()\n\tretryConfig := w.modelRetryConfig\n\n\tfor i := len(w.handlers) - 1; i >= 0; i-- {\n\t\thandler := w.handlers[i]\n\t\tinnerEndpoint := endpoint\n\t\tbaseToolInfos := w.toolInfos\n\t\tendpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\tbaseOpts := &model.Options{Tools: baseToolInfos}\n\t\t\tcommonOpts := model.GetCommonOptions(baseOpts, opts...)\n\t\t\tmc := &ModelContext{Tools: commonOpts.Tools, ModelRetryConfig: retryConfig}\n\t\t\twrappedModel, err := handler.WrapModel(ctx, &endpointModel{generate: innerEndpoint}, mc)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\treturn wrappedModel.Generate(ctx, input, opts...)\n\t\t}\n\t}\n\n\tif !hasUserEventSender {\n\t\tinnerEndpoint := endpoint\n\t\teventSender := NewEventSenderModelWrapper()\n\t\tendpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\texecCtx := getChatModelAgentExecCtx(ctx)\n\t\t\tif execCtx == nil || execCtx.generator == nil {\n\t\t\t\treturn innerEndpoint(ctx, input, opts...)\n\t\t\t}\n\t\t\tmc := &ModelContext{ModelRetryConfig: retryConfig}\n\t\t\twrappedModel, err := eventSender.WrapModel(ctx, &endpointModel{generate: innerEndpoint}, mc)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\treturn wrappedModel.Generate(ctx, input, opts...)\n\t\t}\n\t}\n\n\tif w.modelRetryConfig != nil {\n\t\tinnerEndpoint := endpoint\n\t\tendpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\tretryWrapper := newRetryModelWrapper(&endpointModel{generate: innerEndpoint}, w.modelRetryConfig)\n\t\t\treturn retryWrapper.Generate(ctx, input, opts...)\n\t\t}\n\t}\n\n\treturn endpoint\n}\n\nfunc (w *stateModelWrapper) wrapStreamEndpoint(endpoint streamEndpoint) streamEndpoint {\n\thasUserEventSender := w.hasUserEventSender()\n\tretryConfig := w.modelRetryConfig\n\n\tfor i := len(w.handlers) - 1; i >= 0; i-- {\n\t\thandler := w.handlers[i]\n\t\tinnerEndpoint := endpoint\n\t\tbaseToolInfos := w.toolInfos\n\t\tendpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\t\t\tbaseOpts := &model.Options{Tools: baseToolInfos}\n\t\t\tcommonOpts := model.GetCommonOptions(baseOpts, opts...)\n\t\t\tmc := &ModelContext{Tools: commonOpts.Tools, ModelRetryConfig: retryConfig}\n\t\t\twrappedModel, err := handler.WrapModel(ctx, &endpointModel{stream: innerEndpoint}, mc)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\treturn wrappedModel.Stream(ctx, input, opts...)\n\t\t}\n\t}\n\n\tif !hasUserEventSender {\n\t\tinnerEndpoint := endpoint\n\t\teventSender := NewEventSenderModelWrapper()\n\t\tendpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\t\t\texecCtx := getChatModelAgentExecCtx(ctx)\n\t\t\tif execCtx == nil || execCtx.generator == nil {\n\t\t\t\treturn innerEndpoint(ctx, input, opts...)\n\t\t\t}\n\t\t\tmc := &ModelContext{ModelRetryConfig: retryConfig}\n\t\t\twrappedModel, err := eventSender.WrapModel(ctx, &endpointModel{stream: innerEndpoint}, mc)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\treturn wrappedModel.Stream(ctx, input, opts...)\n\t\t}\n\t}\n\n\tif w.modelRetryConfig != nil {\n\t\tinnerEndpoint := endpoint\n\t\tendpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\t\t\tretryWrapper := newRetryModelWrapper(&endpointModel{stream: innerEndpoint}, w.modelRetryConfig)\n\t\t\treturn retryWrapper.Stream(ctx, input, opts...)\n\t\t}\n\t}\n\n\treturn endpoint\n}\n\nfunc (w *stateModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tvar stateMessages []Message\n\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tstateMessages = st.Messages\n\t\treturn nil\n\t})\n\n\tstate := &ChatModelAgentState{Messages: append(stateMessages, input...)}\n\n\tfor _, m := range w.middlewares {\n\t\tif m.BeforeChatModel != nil {\n\t\t\tif err := m.BeforeChatModel(ctx, state); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\n\tbaseOpts := &model.Options{Tools: w.toolInfos}\n\tcommonOpts := model.GetCommonOptions(baseOpts, opts...)\n\tmc := &ModelContext{Tools: commonOpts.Tools, ModelRetryConfig: w.modelRetryConfig}\n\tfor _, handler := range w.handlers {\n\t\tvar err error\n\t\tctx, state, err = handler.BeforeModelRewriteState(ctx, state, mc)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tst.Messages = state.Messages\n\t\treturn nil\n\t})\n\n\twrappedEndpoint := w.wrapGenerateEndpoint(w.inner.Generate)\n\tresult, err := wrappedEndpoint(ctx, state.Messages, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tstate.Messages = append(state.Messages, result)\n\n\tfor _, handler := range w.handlers {\n\t\tctx, state, err = handler.AfterModelRewriteState(ctx, state, mc)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tfor _, m := range w.middlewares {\n\t\tif m.AfterChatModel != nil {\n\t\t\tif err := m.AfterChatModel(ctx, state); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\n\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tst.Messages = state.Messages\n\t\treturn nil\n\t})\n\n\tif len(state.Messages) == 0 {\n\t\treturn nil, errors.New(\"no messages left in state after model call\")\n\t}\n\treturn state.Messages[len(state.Messages)-1], nil\n}\n\nfunc (w *stateModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tvar stateMessages []Message\n\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tstateMessages = st.Messages\n\t\treturn nil\n\t})\n\n\tstate := &ChatModelAgentState{Messages: append(stateMessages, input...)}\n\n\tfor _, m := range w.middlewares {\n\t\tif m.BeforeChatModel != nil {\n\t\t\tif err := m.BeforeChatModel(ctx, state); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\n\tbaseOpts := &model.Options{Tools: w.toolInfos}\n\tcommonOpts := model.GetCommonOptions(baseOpts, opts...)\n\tmc := &ModelContext{Tools: commonOpts.Tools, ModelRetryConfig: w.modelRetryConfig}\n\tfor _, handler := range w.handlers {\n\t\tvar err error\n\t\tctx, state, err = handler.BeforeModelRewriteState(ctx, state, mc)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tst.Messages = state.Messages\n\t\treturn nil\n\t})\n\n\twrappedEndpoint := w.wrapStreamEndpoint(w.inner.Stream)\n\tstream, err := wrappedEndpoint(ctx, state.Messages, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tresult, err := schema.ConcatMessageStream(stream)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tstate.Messages = append(state.Messages, result)\n\n\tfor _, handler := range w.handlers {\n\t\tctx, state, err = handler.AfterModelRewriteState(ctx, state, mc)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tfor _, m := range w.middlewares {\n\t\tif m.AfterChatModel != nil {\n\t\t\tif err := m.AfterChatModel(ctx, state); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\n\t_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {\n\t\tst.Messages = state.Messages\n\t\treturn nil\n\t})\n\n\tif len(state.Messages) == 0 {\n\t\treturn nil, errors.New(\"no messages left in state after model call\")\n\t}\n\treturn schema.StreamReaderFromArray([]*schema.Message{state.Messages[len(state.Messages)-1]}), nil\n}\n\ntype endpointModel struct {\n\tgenerate generateEndpoint\n\tstream   streamEndpoint\n}\n\nfunc (m *endpointModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tif m.generate != nil {\n\t\treturn m.generate(ctx, input, opts...)\n\t}\n\treturn nil, errors.New(\"generate endpoint not set\")\n}\n\nfunc (m *endpointModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tif m.stream != nil {\n\t\treturn m.stream(ctx, input, opts...)\n\t}\n\treturn nil, errors.New(\"stream endpoint not set\")\n}\n"
  },
  {
    "path": "adk/wrappers_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage adk\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype testEnhancedToolWrapperHandler struct {\n\t*BaseChatModelAgentMiddleware\n\twrapEnhancedInvokableFn  func(context.Context, EnhancedInvokableToolCallEndpoint, *ToolContext) EnhancedInvokableToolCallEndpoint\n\twrapEnhancedStreamableFn func(context.Context, EnhancedStreamableToolCallEndpoint, *ToolContext) EnhancedStreamableToolCallEndpoint\n}\n\nfunc (h *testEnhancedToolWrapperHandler) WrapEnhancedInvokableToolCall(ctx context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error) {\n\tif h.wrapEnhancedInvokableFn != nil {\n\t\treturn h.wrapEnhancedInvokableFn(ctx, endpoint, tCtx), nil\n\t}\n\treturn endpoint, nil\n}\n\nfunc (h *testEnhancedToolWrapperHandler) WrapEnhancedStreamableToolCall(ctx context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error) {\n\tif h.wrapEnhancedStreamableFn != nil {\n\t\treturn h.wrapEnhancedStreamableFn(ctx, endpoint, tCtx), nil\n\t}\n\treturn endpoint, nil\n}\n\nfunc newTestEnhancedInvokableToolCallWrapper(beforeFn, afterFn func()) func(context.Context, EnhancedInvokableToolCallEndpoint, *ToolContext) EnhancedInvokableToolCallEndpoint {\n\treturn func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\tif beforeFn != nil {\n\t\t\t\tbeforeFn()\n\t\t\t}\n\t\t\tresult, err := endpoint(ctx, toolArgument, opts...)\n\t\t\tif afterFn != nil {\n\t\t\t\tafterFn()\n\t\t\t}\n\t\t\treturn result, err\n\t\t}\n\t}\n}\n\nfunc newTestEnhancedStreamableToolCallWrapper(beforeFn, afterFn func()) func(context.Context, EnhancedStreamableToolCallEndpoint, *ToolContext) EnhancedStreamableToolCallEndpoint {\n\treturn func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint {\n\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\tif beforeFn != nil {\n\t\t\t\tbeforeFn()\n\t\t\t}\n\t\t\tresult, err := endpoint(ctx, toolArgument, opts...)\n\t\t\tif afterFn != nil {\n\t\t\t\tafterFn()\n\t\t\t}\n\t\t\treturn result, err\n\t\t}\n\t}\n}\n\nfunc TestHandlersToToolMiddlewaresEnhanced(t *testing.T) {\n\tt.Run(\"OnlyEnhancedInvokableHandler\", func(t *testing.T) {\n\t\tvar called bool\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\t\t\t\tcalled = true\n\t\t\t\t\t\treturn endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tassert.Len(t, middlewares, 1)\n\t\tassert.NotNil(t, middlewares[0].EnhancedInvokable)\n\t\tassert.NotNil(t, middlewares[0].Invokable)\n\t\tassert.NotNil(t, middlewares[0].Streamable)\n\t\tassert.NotNil(t, middlewares[0].EnhancedStreamable)\n\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn &compose.EnhancedInvokableToolOutput{\n\t\t\t\tResult: &schema.ToolResult{\n\t\t\t\t\tParts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: \"test\"}},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedInvokable(mockEndpoint)\n\t\t_, err := wrapped(context.Background(), &compose.ToolInput{\n\t\t\tName:      \"test_tool\",\n\t\t\tCallID:    \"call-1\",\n\t\t\tArguments: `{\"input\": \"test\"}`,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, called)\n\t})\n\n\tt.Run(\"OnlyEnhancedStreamableHandler\", func(t *testing.T) {\n\t\tvar called bool\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\t\t\t\tcalled = true\n\t\t\t\t\t\treturn endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tassert.Len(t, middlewares, 1)\n\t\tassert.NotNil(t, middlewares[0].EnhancedStreamable)\n\t\tassert.NotNil(t, middlewares[0].Invokable)\n\t\tassert.NotNil(t, middlewares[0].Streamable)\n\t\tassert.NotNil(t, middlewares[0].EnhancedInvokable)\n\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\t\treturn &compose.EnhancedStreamableToolOutput{\n\t\t\t\tResult: schema.StreamReaderFromArray([]*schema.ToolResult{\n\t\t\t\t\t{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: \"test\"}}},\n\t\t\t\t}),\n\t\t\t}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedStreamable(mockEndpoint)\n\t\t_, err := wrapped(context.Background(), &compose.ToolInput{\n\t\t\tName:      \"test_tool\",\n\t\t\tCallID:    \"call-1\",\n\t\t\tArguments: `{\"input\": \"test\"}`,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, called)\n\t})\n\n\tt.Run(\"MixedHandlers\", func(t *testing.T) {\n\t\tvar invokableCalled, streamableCalled, enhancedInvokableCalled, enhancedStreamableCalled bool\n\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapInvokableFn: func(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) InvokableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\t\t\t\t\t\tinvokableCalled = true\n\t\t\t\t\t\treturn endpoint(ctx, argumentsInJSON, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t\t&testToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapStreamableFn: func(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) StreamableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\t\t\t\t\t\tstreamableCalled = true\n\t\t\t\t\t\treturn endpoint(ctx, argumentsInJSON, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\t\t\t\tenhancedInvokableCalled = true\n\t\t\t\t\t\treturn endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\t\t\t\tenhancedStreamableCalled = true\n\t\t\t\t\t\treturn endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tassert.Len(t, middlewares, 4)\n\n\t\tinvokableEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\t\treturn &compose.ToolOutput{Result: \"test\"}, nil\n\t\t}\n\t\t_, _ = middlewares[3].Invokable(invokableEndpoint)(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\n\t\tstreamableEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\t\treturn &compose.StreamToolOutput{Result: schema.StreamReaderFromArray([]string{\"test\"})}, nil\n\t\t}\n\t\t_, _ = middlewares[2].Streamable(streamableEndpoint)(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\n\t\tenhancedInvokableEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil\n\t\t}\n\t\t_, _ = middlewares[1].EnhancedInvokable(enhancedInvokableEndpoint)(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\n\t\tenhancedStreamableEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\t\treturn &compose.EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{{}})}, nil\n\t\t}\n\t\t_, _ = middlewares[0].EnhancedStreamable(enhancedStreamableEndpoint)(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\n\t\tassert.True(t, invokableCalled)\n\t\tassert.True(t, streamableCalled)\n\t\tassert.True(t, enhancedInvokableCalled)\n\t\tassert.True(t, enhancedStreamableCalled)\n\t})\n\n\tt.Run(\"NoHandlers\", func(t *testing.T) {\n\t\thandlers := []ChatModelAgentMiddleware{}\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tassert.Len(t, middlewares, 0)\n\t})\n\n\tt.Run(\"HandlerWithNoToolWrappers\", func(t *testing.T) {\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&BaseChatModelAgentMiddleware{},\n\t\t}\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tassert.Len(t, middlewares, 1)\n\t})\n\n\tt.Run(\"EnhancedInvokableToolCallErrorPropagation\", func(t *testing.T) {\n\t\texpectedErr := errors.New(\"test error\")\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\t\t\t\treturn nil, expectedErr\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedInvokable(mockEndpoint)\n\t\t_, err := wrapped(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\t\tassert.Error(t, err)\n\t\tassert.Equal(t, expectedErr, err)\n\t})\n\n\tt.Run(\"EnhancedStreamableToolCallErrorPropagation\", func(t *testing.T) {\n\t\texpectedErr := errors.New(\"test error\")\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\t\t\t\treturn nil, expectedErr\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\t\treturn &compose.EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{})}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedStreamable(mockEndpoint)\n\t\t_, err := wrapped(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\t\tassert.Error(t, err)\n\t\tassert.Equal(t, expectedErr, err)\n\t})\n\n\tt.Run(\"MultipleEnhancedInvokableWrappers\", func(t *testing.T) {\n\t\tvar executionOrder []string\n\t\tvar mu sync.Mutex\n\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: newTestEnhancedInvokableToolCallWrapper(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\texecutionOrder = append(executionOrder, \"handler1-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\texecutionOrder = append(executionOrder, \"handler1-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t),\n\t\t\t},\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: newTestEnhancedInvokableToolCallWrapper(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\texecutionOrder = append(executionOrder, \"handler2-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\texecutionOrder = append(executionOrder, \"handler2-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t),\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tassert.Len(t, middlewares, 2)\n\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedInvokable(middlewares[1].EnhancedInvokable(mockEndpoint))\n\t\t_, err := wrapped(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []string{\"handler2-before\", \"handler1-before\", \"handler1-after\", \"handler2-after\"}, executionOrder)\n\t})\n\n\tt.Run(\"MultipleEnhancedStreamableWrappers\", func(t *testing.T) {\n\t\tvar executionOrder []string\n\t\tvar mu sync.Mutex\n\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedStreamableFn: newTestEnhancedStreamableToolCallWrapper(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\texecutionOrder = append(executionOrder, \"handler1-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\texecutionOrder = append(executionOrder, \"handler1-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t),\n\t\t\t},\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedStreamableFn: newTestEnhancedStreamableToolCallWrapper(\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\texecutionOrder = append(executionOrder, \"handler2-before\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t\tfunc() {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\texecutionOrder = append(executionOrder, \"handler2-after\")\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t},\n\t\t\t\t),\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tassert.Len(t, middlewares, 2)\n\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\t\treturn &compose.EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{{}})}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedStreamable(middlewares[1].EnhancedStreamable(mockEndpoint))\n\t\t_, err := wrapped(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []string{\"handler2-before\", \"handler1-before\", \"handler1-after\", \"handler2-after\"}, executionOrder)\n\t})\n}\n\nfunc TestEnhancedToolContextPropagation(t *testing.T) {\n\tt.Run(\"ToolContextContainsCorrectInfo\", func(t *testing.T) {\n\t\tvar capturedCtx *ToolContext\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\t\t\t\tcapturedCtx = tCtx\n\t\t\t\t\treturn endpoint\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedInvokable(mockEndpoint)\n\t\t_, _ = wrapped(context.Background(), &compose.ToolInput{\n\t\t\tName:      \"my_tool\",\n\t\t\tCallID:    \"call-123\",\n\t\t\tArguments: `{\"key\": \"value\"}`,\n\t\t})\n\n\t\tassert.NotNil(t, capturedCtx)\n\t\tassert.Equal(t, \"my_tool\", capturedCtx.Name)\n\t\tassert.Equal(t, \"call-123\", capturedCtx.CallID)\n\t})\n\n\tt.Run(\"StreamableToolContextContainsCorrectInfo\", func(t *testing.T) {\n\t\tvar capturedCtx *ToolContext\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) EnhancedStreamableToolCallEndpoint {\n\t\t\t\t\tcapturedCtx = tCtx\n\t\t\t\t\treturn endpoint\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\t\treturn &compose.EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{{}})}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedStreamable(mockEndpoint)\n\t\t_, _ = wrapped(context.Background(), &compose.ToolInput{\n\t\t\tName:      \"stream_tool\",\n\t\t\tCallID:    \"call-456\",\n\t\t\tArguments: `{\"data\": \"test\"}`,\n\t\t})\n\n\t\tassert.NotNil(t, capturedCtx)\n\t\tassert.Equal(t, \"stream_tool\", capturedCtx.Name)\n\t\tassert.Equal(t, \"call-456\", capturedCtx.CallID)\n\t})\n}\n\nfunc TestBaseChatModelAgentMiddlewareEnhancedDefaults(t *testing.T) {\n\tt.Run(\"DefaultEnhancedInvokableReturnsEndpoint\", func(t *testing.T) {\n\t\tbase := &BaseChatModelAgentMiddleware{}\n\n\t\tvar called bool\n\t\tendpoint := func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\tcalled = true\n\t\t\treturn &schema.ToolResult{}, nil\n\t\t}\n\n\t\twrapped, wrapErr := base.WrapEnhancedInvokableToolCall(context.Background(), endpoint, &ToolContext{Name: \"test\", CallID: \"1\"})\n\t\tassert.NoError(t, wrapErr)\n\t\t_, err := wrapped(context.Background(), &schema.ToolArgument{Text: \"{}\"})\n\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, called)\n\t})\n\n\tt.Run(\"DefaultEnhancedStreamableReturnsEndpoint\", func(t *testing.T) {\n\t\tbase := &BaseChatModelAgentMiddleware{}\n\n\t\tvar called bool\n\t\tendpoint := func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\tcalled = true\n\t\t\treturn schema.StreamReaderFromArray([]*schema.ToolResult{}), nil\n\t\t}\n\n\t\twrapped, wrapErr := base.WrapEnhancedStreamableToolCall(context.Background(), endpoint, &ToolContext{Name: \"test\", CallID: \"1\"})\n\t\tassert.NoError(t, wrapErr)\n\t\t_, err := wrapped(context.Background(), &schema.ToolArgument{Text: \"{}\"})\n\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, called)\n\t})\n}\n\nfunc TestEnhancedToolArgumentsPropagation(t *testing.T) {\n\tt.Run(\"ArgumentsPassedCorrectly\", func(t *testing.T) {\n\t\tvar capturedArgs string\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\t\t\t\tcapturedArgs = toolArgument.Text\n\t\t\t\t\t\treturn endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedInvokable(mockEndpoint)\n\t\t_, _ = wrapped(context.Background(), &compose.ToolInput{\n\t\t\tName:      \"test_tool\",\n\t\t\tCallID:    \"call-1\",\n\t\t\tArguments: `{\"name\": \"test\", \"value\": 123}`,\n\t\t})\n\n\t\tassert.Equal(t, `{\"name\": \"test\", \"value\": 123}`, capturedArgs)\n\t})\n}\n\nfunc TestEnhancedToolResultPropagation(t *testing.T) {\n\tt.Run(\"ResultPassedThroughMiddleware\", func(t *testing.T) {\n\t\texpectedResult := &schema.ToolResult{\n\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t{Type: schema.ToolPartTypeText, Text: \"original result\"},\n\t\t\t},\n\t\t}\n\n\t\tvar capturedResult *schema.ToolResult\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\t\t\t\tresult, err := endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t\tcapturedResult = result\n\t\t\t\t\t\treturn result, err\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn &compose.EnhancedInvokableToolOutput{Result: expectedResult}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedInvokable(mockEndpoint)\n\t\toutput, err := wrapped(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, expectedResult, capturedResult)\n\t\tassert.Equal(t, expectedResult, output.Result)\n\t})\n\n\tt.Run(\"ModifiedResultPropagated\", func(t *testing.T) {\n\t\tmodifiedResult := &schema.ToolResult{\n\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t{Type: schema.ToolPartTypeText, Text: \"modified result\"},\n\t\t\t},\n\t\t}\n\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\t\t\t\t_, err := endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn modifiedResult, nil\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: \"original\"}},\n\t\t\t}}, nil\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedInvokable(mockEndpoint)\n\t\toutput, err := wrapped(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, modifiedResult, output.Result)\n\t\tassert.Equal(t, \"modified result\", output.Result.Parts[0].Text)\n\t})\n}\n\nfunc TestEnhancedToolEndpointErrorFromNext(t *testing.T) {\n\tt.Run(\"EnhancedInvokableNextError\", func(t *testing.T) {\n\t\texpectedErr := errors.New(\"next endpoint error\")\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\t\t\t\t\t\treturn endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\treturn nil, expectedErr\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedInvokable(mockEndpoint)\n\t\t_, err := wrapped(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\n\t\tassert.Error(t, err)\n\t\tassert.Equal(t, expectedErr, err)\n\t})\n\n\tt.Run(\"EnhancedStreamableNextError\", func(t *testing.T) {\n\t\texpectedErr := errors.New(\"next endpoint error\")\n\t\thandlers := []ChatModelAgentMiddleware{\n\t\t\t&testEnhancedToolWrapperHandler{\n\t\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\t\twrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\t\t\t\treturn endpoint(ctx, toolArgument, opts...)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmiddlewares := handlersToToolMiddlewares(handlers)\n\t\tmockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\t\treturn nil, expectedErr\n\t\t}\n\n\t\twrapped := middlewares[0].EnhancedStreamable(mockEndpoint)\n\t\t_, err := wrapped(context.Background(), &compose.ToolInput{Name: \"test\", CallID: \"1\", Arguments: \"{}\"})\n\n\t\tassert.Error(t, err)\n\t\tassert.Equal(t, expectedErr, err)\n\t})\n}\n\nfunc TestWrapModelStreamChunksPreserved(t *testing.T) {\n\tt.Run(\"AgentEventMessageStreamShouldPreserveChunksWithNoopWrapModel\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tchunk1 := schema.AssistantMessage(\"Hello \", nil)\n\t\tchunk2 := schema.AssistantMessage(\"World\", nil)\n\n\t\tmockModel := &mockStreamingModel{\n\t\t\tchunks: []*schema.Message{chunk1, chunk2},\n\t\t}\n\n\t\tnoopWrapModelHandler := &testModelWrapperHandler{\n\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\tfn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel {\n\t\t\t\treturn m\n\t\t\t},\n\t\t}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       mockModel,\n\t\t\tHandlers:    []ChatModelAgentMiddleware{noopWrapModelHandler},\n\t\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\t\tMaxRetries: 3,\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           agent,\n\t\t\tEnableStreaming: true,\n\t\t})\n\t\titer := r.Run(ctx, []Message{schema.UserMessage(\"test\")})\n\n\t\tvar streamingEvents []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.IsStreaming &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Assistant {\n\t\t\t\tstreamingEvents = append(streamingEvents, event)\n\t\t\t}\n\t\t}\n\n\t\tassert.GreaterOrEqual(t, len(streamingEvents), 1, \"Should have at least one streaming event\")\n\n\t\tif len(streamingEvents) > 0 {\n\t\t\tevent := streamingEvents[0]\n\t\t\tassert.NotNil(t, event.Output.MessageOutput.MessageStream, \"Event should have message stream\")\n\n\t\t\tvar receivedChunks []*schema.Message\n\t\t\tfor {\n\t\t\t\tchunk, recvErr := event.Output.MessageOutput.MessageStream.Recv()\n\t\t\t\tif recvErr != nil {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\treceivedChunks = append(receivedChunks, chunk)\n\t\t\t}\n\n\t\t\tassert.Equal(t, 2, len(receivedChunks),\n\t\t\t\t\"AgentEvent's MessageStream should contain 2 separate chunks, not 1 concatenated chunk. \"+\n\t\t\t\t\t\"Got %d chunks instead. This indicates the stream is being concatenated before being sent to AgentEvent.\",\n\t\t\t\tlen(receivedChunks))\n\n\t\t\tif len(receivedChunks) >= 2 {\n\t\t\t\tassert.Equal(t, \"Hello \", receivedChunks[0].Content, \"First chunk content should be preserved\")\n\t\t\t\tassert.Equal(t, \"World\", receivedChunks[1].Content, \"Second chunk content should be preserved\")\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"AgentEventMessageStreamShouldReflectUserMiddlewareModifications\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tchunk1 := schema.AssistantMessage(\"Hello \", nil)\n\t\tchunk2 := schema.AssistantMessage(\"World\", nil)\n\n\t\tmockModel := &mockStreamingModel{\n\t\t\tchunks: []*schema.Message{chunk1, chunk2},\n\t\t}\n\n\t\tstreamConsumingWrapModelHandler := &testModelWrapperHandler{\n\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\tfn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel {\n\t\t\t\treturn &streamConsumingModelWrapper{inner: m}\n\t\t\t},\n\t\t}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       mockModel,\n\t\t\tHandlers:    []ChatModelAgentMiddleware{streamConsumingWrapModelHandler},\n\t\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\t\tMaxRetries: 3,\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           agent,\n\t\t\tEnableStreaming: true,\n\t\t})\n\t\titer := r.Run(ctx, []Message{schema.UserMessage(\"test\")})\n\n\t\tvar streamingEvents []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.IsStreaming &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Assistant {\n\t\t\t\tstreamingEvents = append(streamingEvents, event)\n\t\t\t}\n\t\t}\n\n\t\tassert.GreaterOrEqual(t, len(streamingEvents), 1, \"Should have at least one streaming event\")\n\n\t\tif len(streamingEvents) > 0 {\n\t\t\tevent := streamingEvents[0]\n\t\t\tassert.NotNil(t, event.Output.MessageOutput.MessageStream, \"Event should have message stream\")\n\n\t\t\tvar receivedChunks []*schema.Message\n\t\t\tfor {\n\t\t\t\tchunk, recvErr := event.Output.MessageOutput.MessageStream.Recv()\n\t\t\t\tif recvErr != nil {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\treceivedChunks = append(receivedChunks, chunk)\n\t\t\t}\n\n\t\t\tassert.Equal(t, 1, len(receivedChunks),\n\t\t\t\t\"AgentEvent's MessageStream should contain 1 concatenated chunk (modified by user middleware). \"+\n\t\t\t\t\t\"Got %d chunks instead.\",\n\t\t\t\tlen(receivedChunks))\n\n\t\t\tif len(receivedChunks) >= 1 {\n\t\t\t\tassert.Equal(t, \"Hello World\", receivedChunks[0].Content, \"Chunk content should be concatenated by user middleware\")\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"AgentEventMessageStreamShouldReflectMultipleUserMiddlewareModifications\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tchunk1 := schema.AssistantMessage(\"Hello \", nil)\n\t\tchunk2 := schema.AssistantMessage(\"World\", nil)\n\n\t\tmockModel := &mockStreamingModel{\n\t\t\tchunks: []*schema.Message{chunk1, chunk2},\n\t\t}\n\n\t\thandler1 := &testModelWrapperHandler{\n\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\tfn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel {\n\t\t\t\treturn &streamConsumingModelWrapper{inner: m}\n\t\t\t},\n\t\t}\n\n\t\thandler2 := &testModelWrapperHandler{\n\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\tfn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel {\n\t\t\t\treturn m\n\t\t\t},\n\t\t}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       mockModel,\n\t\t\tHandlers:    []ChatModelAgentMiddleware{handler1, handler2},\n\t\t\tModelRetryConfig: &ModelRetryConfig{\n\t\t\t\tMaxRetries: 3,\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           agent,\n\t\t\tEnableStreaming: true,\n\t\t})\n\t\titer := r.Run(ctx, []Message{schema.UserMessage(\"test\")})\n\n\t\tvar streamingEvents []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.IsStreaming &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Assistant {\n\t\t\t\tstreamingEvents = append(streamingEvents, event)\n\t\t\t}\n\t\t}\n\n\t\tassert.GreaterOrEqual(t, len(streamingEvents), 1, \"Should have at least one streaming event\")\n\n\t\tif len(streamingEvents) > 0 {\n\t\t\tevent := streamingEvents[0]\n\t\t\tassert.NotNil(t, event.Output.MessageOutput.MessageStream, \"Event should have message stream\")\n\n\t\t\tvar receivedChunks []*schema.Message\n\t\t\tfor {\n\t\t\t\tchunk, recvErr := event.Output.MessageOutput.MessageStream.Recv()\n\t\t\t\tif recvErr != nil {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\treceivedChunks = append(receivedChunks, chunk)\n\t\t\t}\n\n\t\t\tassert.Equal(t, 1, len(receivedChunks),\n\t\t\t\t\"AgentEvent's MessageStream should contain 1 concatenated chunk (modified by user middleware). \"+\n\t\t\t\t\t\"Got %d chunks instead.\",\n\t\t\t\tlen(receivedChunks))\n\n\t\t\tif len(receivedChunks) >= 1 {\n\t\t\t\tassert.Equal(t, \"Hello World\", receivedChunks[0].Content, \"Chunk content should be concatenated by user middleware\")\n\t\t\t}\n\t\t}\n\t})\n}\n\ntype mockStreamingModel struct {\n\tchunks []*schema.Message\n}\n\nfunc (m *mockStreamingModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\treturn schema.ConcatMessages(m.chunks)\n}\n\nfunc (m *mockStreamingModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tsr, sw := schema.Pipe[*schema.Message](len(m.chunks))\n\tgo func() {\n\t\tdefer sw.Close()\n\t\tfor _, chunk := range m.chunks {\n\t\t\tsw.Send(chunk, nil)\n\t\t}\n\t}()\n\treturn sr, nil\n}\n\ntype streamConsumingModelWrapper struct {\n\tinner model.BaseChatModel\n}\n\nfunc (m *streamConsumingModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\treturn m.inner.Generate(ctx, input, opts...)\n}\n\nfunc (m *streamConsumingModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tstream, err := m.inner.Stream(ctx, input, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tresult, err := schema.ConcatMessageStream(stream)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn schema.StreamReaderFromArray([]*schema.Message{result}), nil\n}\n\nfunc TestEventSenderModelWrapperCustomPosition(t *testing.T) {\n\tt.Run(\"UserConfiguredEventSenderSkipsDefaultEventSender\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tchunk1 := schema.AssistantMessage(\"Hello \", nil)\n\t\tchunk2 := schema.AssistantMessage(\"World\", nil)\n\n\t\tmockModel := &mockStreamingModel{\n\t\t\tchunks: []*schema.Message{chunk1, chunk2},\n\t\t}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       mockModel,\n\t\t\tHandlers:    []ChatModelAgentMiddleware{NewEventSenderModelWrapper()},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           agent,\n\t\t\tEnableStreaming: true,\n\t\t})\n\t\titer := r.Run(ctx, []Message{schema.UserMessage(\"test\")})\n\n\t\tvar streamingEvents []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.IsStreaming &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Assistant {\n\t\t\t\tstreamingEvents = append(streamingEvents, event)\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 1, len(streamingEvents), \"Should have exactly one streaming event (no duplicate from default event sender)\")\n\t})\n\n\tt.Run(\"EventSenderAfterUserMiddlewareByDefault\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tmockModel := &mockStreamingModel{\n\t\t\tchunks: []*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"Original\", nil),\n\t\t\t},\n\t\t}\n\n\t\tmodifiedContent := \"Modified\"\n\t\tcontentModifyingHandler := &testModelWrapperHandler{\n\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\tfn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel {\n\t\t\t\treturn &contentModifyingModelWrapper{inner: m, newContent: modifiedContent}\n\t\t\t},\n\t\t}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       mockModel,\n\t\t\tHandlers:    []ChatModelAgentMiddleware{contentModifyingHandler},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           agent,\n\t\t\tEnableStreaming: false,\n\t\t})\n\t\titer := r.Run(ctx, []Message{schema.UserMessage(\"test\")})\n\n\t\tvar assistantEvents []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Assistant {\n\t\t\t\tassistantEvents = append(assistantEvents, event)\n\t\t\t}\n\t\t}\n\n\t\tassert.GreaterOrEqual(t, len(assistantEvents), 1, \"Should have at least one assistant event\")\n\t\tif len(assistantEvents) > 0 {\n\t\t\tmsg := assistantEvents[0].Output.MessageOutput.Message\n\t\t\tassert.Equal(t, modifiedContent, msg.Content, \"Event should contain modified content from user middleware\")\n\t\t}\n\t})\n\n\tt.Run(\"EventSenderInnermostGetsOriginalOutput\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\toriginalContent := \"Original\"\n\t\tmockModel := &mockStreamingModel{\n\t\t\tchunks: []*schema.Message{\n\t\t\t\tschema.AssistantMessage(originalContent, nil),\n\t\t\t},\n\t\t}\n\n\t\tmodifiedContent := \"Modified\"\n\t\tcontentModifyingHandler := &testModelWrapperHandler{\n\t\t\tBaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},\n\t\t\tfn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel {\n\t\t\t\treturn &contentModifyingModelWrapper{inner: m, newContent: modifiedContent}\n\t\t\t},\n\t\t}\n\n\t\tagent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{\n\t\t\tName:        \"TestAgent\",\n\t\t\tDescription: \"Test agent\",\n\t\t\tModel:       mockModel,\n\t\t\tHandlers: []ChatModelAgentMiddleware{\n\t\t\t\tcontentModifyingHandler,\n\t\t\t\tNewEventSenderModelWrapper(),\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tr := NewRunner(ctx, RunnerConfig{\n\t\t\tAgent:           agent,\n\t\t\tEnableStreaming: false,\n\t\t})\n\t\titer := r.Run(ctx, []Message{schema.UserMessage(\"test\")})\n\n\t\tvar assistantEvents []*AgentEvent\n\t\tfor {\n\t\t\tevent, ok := iter.Next()\n\t\t\tif !ok {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif event.Output != nil && event.Output.MessageOutput != nil &&\n\t\t\t\tevent.Output.MessageOutput.Role == schema.Assistant {\n\t\t\t\tassistantEvents = append(assistantEvents, event)\n\t\t\t}\n\t\t}\n\n\t\tassert.GreaterOrEqual(t, len(assistantEvents), 1, \"Should have at least one assistant event\")\n\t\tif len(assistantEvents) > 0 {\n\t\t\tmsg := assistantEvents[0].Output.MessageOutput.Message\n\t\t\tassert.Equal(t, originalContent, msg.Content, \"Event should contain original content (EventSenderModelWrapper is innermost)\")\n\t\t}\n\t})\n}\n\ntype contentModifyingModelWrapper struct {\n\tinner      model.BaseChatModel\n\tnewContent string\n}\n\nfunc (m *contentModifyingModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tresult, err := m.inner.Generate(ctx, input, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tresult.Content = m.newContent\n\treturn result, nil\n}\n\nfunc (m *contentModifyingModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tstream, err := m.inner.Stream(ctx, input, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tresult, err := schema.ConcatMessageStream(stream)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tresult.Content = m.newContent\n\treturn schema.StreamReaderFromArray([]*schema.Message{result}), nil\n}\n"
  },
  {
    "path": "callbacks/aspect_inject.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/internal/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// OnStart Fast inject callback input / output aspect for component developer\n// e.g.\n//\n//\tfunc (t *testChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (resp *schema.Message, err error) {\n//\t\tdefer func() {\n//\t\t\tif err != nil {\n//\t\t\t\tcallbacks.OnError(ctx, err)\n//\t\t\t}\n//\t\t}()\n//\n//\t\tctx = callbacks.OnStart(ctx, &model.CallbackInput{\n//\t\t\tMessages: input,\n//\t\t\tTools:    nil,\n//\t\t\tExtra:    nil,\n//\t\t})\n//\n//\t\t// do smt\n//\n//\t\tctx = callbacks.OnEnd(ctx, &model.CallbackOutput{\n//\t\t\tMessage: resp,\n//\t\t\tExtra:   nil,\n//\t\t})\n//\n//\t\treturn resp, nil\n//\t}\n\n// OnStart invokes the OnStart timing for all registered handlers in the\n// context. This is called by component implementations that manage their own\n// callbacks (i.e. implement [components.Checker] and return true from\n// IsCallbacksEnabled). The returned context must be propagated to subsequent\n// OnEnd/OnError calls so handlers can correlate start and end events.\n//\n// Handlers are invoked in reverse registration order (last registered = first\n// called) to match the middleware wrapping convention.\n//\n// Example — typical usage inside a component's Generate method:\n//\n//\tfunc (m *myChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n//\t    ctx = callbacks.OnStart(ctx, &model.CallbackInput{Messages: input})\n//\t    resp, err := m.doGenerate(ctx, input, opts...)\n//\t    if err != nil {\n//\t        callbacks.OnError(ctx, err)\n//\t        return nil, err\n//\t    }\n//\t    callbacks.OnEnd(ctx, &model.CallbackOutput{Message: resp})\n//\t    return resp, nil\n//\t}\nfunc OnStart[T any](ctx context.Context, input T) context.Context {\n\tctx, _ = callbacks.On(ctx, input, callbacks.OnStartHandle[T], TimingOnStart, true)\n\n\treturn ctx\n}\n\n// OnEnd invokes the OnEnd timing for all registered handlers. Call this after\n// the component produces a successful result. Handlers run in registration\n// order (first registered = first called).\n//\n// Do not call both OnEnd and OnError for the same invocation — OnEnd signals\n// success; OnError signals failure.\nfunc OnEnd[T any](ctx context.Context, output T) context.Context {\n\tctx, _ = callbacks.On(ctx, output, callbacks.OnEndHandle[T], TimingOnEnd, false)\n\n\treturn ctx\n}\n\n// OnStartWithStreamInput invokes the OnStartWithStreamInput timing. Use this\n// when the component's input is itself a stream (Collect / Transform\n// paradigms). The framework automatically copies the stream so each handler\n// receives an independent reader; handlers MUST close their copy or the\n// underlying goroutine will leak.\n//\n// Returns the updated context and a new StreamReader that the component should\n// use going forward (the original is consumed by the framework).\nfunc OnStartWithStreamInput[T any](ctx context.Context, input *schema.StreamReader[T]) (\n\tnextCtx context.Context, newStreamReader *schema.StreamReader[T]) {\n\n\treturn callbacks.On(ctx, input, callbacks.OnStartWithStreamInputHandle[T], TimingOnStartWithStreamInput, true)\n}\n\n// OnEndWithStreamOutput invokes the OnEndWithStreamOutput timing. Use this\n// when the component produces a streaming output (Stream / Transform\n// paradigms). Like OnStartWithStreamInput, stream copies are made per\n// handler; each handler must close its copy.\n//\n// Returns the updated context and the StreamReader the component should return\n// to its caller.\nfunc OnEndWithStreamOutput[T any](ctx context.Context, output *schema.StreamReader[T]) (\n\tnextCtx context.Context, newStreamReader *schema.StreamReader[T]) {\n\n\treturn callbacks.On(ctx, output, callbacks.OnEndWithStreamOutputHandle[T], TimingOnEndWithStreamOutput, false)\n}\n\n// OnError invokes the OnError timing for all registered handlers. Call this\n// when the component returns an error. Errors that occur mid-stream (after the\n// StreamReader has been returned) are NOT routed through OnError; they surface\n// as errors inside Recv.\n//\n// Handlers run in registration order (same as OnEnd).\nfunc OnError(ctx context.Context, err error) context.Context {\n\tctx, _ = callbacks.On(ctx, err, callbacks.OnErrorHandle, TimingOnError, false)\n\n\treturn ctx\n}\n\n// EnsureRunInfo ensures the context carries a [RunInfo] for the given type and\n// component kind. If the context already has a matching RunInfo, it is\n// returned unchanged. Otherwise, a new callback manager is created that\n// inherits the global handlers plus any handlers already in ctx.\n//\n// Component implementations that set IsCallbacksEnabled() = true should call\n// this at the start of every public method (Generate, Stream, etc.) before\n// calling [OnStart], so that the RunInfo is never missing from callbacks.\nfunc EnsureRunInfo(ctx context.Context, typ string, comp components.Component) context.Context {\n\treturn callbacks.EnsureRunInfo(ctx, typ, comp)\n}\n\n// ReuseHandlers creates a new context that inherits all handlers already\n// present in ctx and sets a new RunInfo. Global handlers are added if ctx\n// carries none yet.\n//\n// Use this when a component calls another component internally and wants the\n// inner component's callbacks to share the same set of handlers as the outer\n// component, but with the inner component's own identity in RunInfo:\n//\n//\tinnerCtx := callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{\n//\t    Type:      \"InnerChatModel\",\n//\t    Component: components.ComponentOfChatModel,\n//\t    Name:      \"inner-cm\",\n//\t})\nfunc ReuseHandlers(ctx context.Context, info *RunInfo) context.Context {\n\treturn callbacks.ReuseHandlers(ctx, info)\n}\n\n// InitCallbacks creates a new context with the given RunInfo and handlers,\n// completely replacing any RunInfo and handlers already in ctx.\n//\n// Use this when running a component standalone outside a Graph — the Graph\n// normally manages RunInfo injection automatically, but standalone callers must\n// set it up themselves:\n//\n//\tctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{\n//\t    Type:      myModel.GetType(),\n//\t    Component: components.ComponentOfChatModel,\n//\t    Name:      \"my-model\",\n//\t}, myHandler)\nfunc InitCallbacks(ctx context.Context, info *RunInfo, handlers ...Handler) context.Context {\n\treturn callbacks.InitCallbacks(ctx, info, handlers...)\n}\n"
  },
  {
    "path": "callbacks/aspect_inject_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"io\"\n\t\"strconv\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/internal/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestAspectInject(t *testing.T) {\n\tt.Run(\"ctx without manager\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tctx = OnStart(ctx, 1)\n\t\tctx = OnEnd(ctx, 2)\n\t\tctx = OnError(ctx, fmt.Errorf(\"3\"))\n\t\tisr, isw := schema.Pipe[int](2)\n\t\tgo func() {\n\t\t\tfor i := 0; i < 10; i++ {\n\t\t\t\tisw.Send(i, nil)\n\t\t\t}\n\t\t\tisw.Close()\n\t\t}()\n\n\t\tvar nisr *schema.StreamReader[int]\n\t\tctx, nisr = OnStartWithStreamInput(ctx, isr)\n\t\tj := 0\n\t\tfor {\n\t\t\ti, err := nisr.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tassert.NoError(t, err)\n\t\t\tassert.Equal(t, j, i)\n\t\t\tj++\n\t\t}\n\t\tnisr.Close()\n\n\t\tosr, osw := schema.Pipe[int](2)\n\t\tgo func() {\n\t\t\tfor i := 0; i < 10; i++ {\n\t\t\t\tosw.Send(i, nil)\n\t\t\t}\n\t\t\tosw.Close()\n\t\t}()\n\n\t\tvar nosr *schema.StreamReader[int]\n\t\tctx, nosr = OnEndWithStreamOutput(ctx, osr)\n\t\tj = 0\n\t\tfor {\n\t\t\ti, err := nosr.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tassert.NoError(t, err)\n\t\t\tassert.Equal(t, j, i)\n\t\t\tj++\n\t\t}\n\t\tnosr.Close()\n\t})\n\n\tt.Run(\"ctx with manager\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tcnt := 0\n\n\t\thb := NewHandlerBuilder().\n\t\t\tOnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context {\n\t\t\t\tcnt += input.(int)\n\t\t\t\treturn ctx\n\t\t\t}).\n\t\t\tOnEndFn(func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context {\n\t\t\t\tcnt += output.(int)\n\t\t\t\treturn ctx\n\t\t\t}).\n\t\t\tOnErrorFn(func(ctx context.Context, info *RunInfo, err error) context.Context {\n\t\t\t\tv, _ := strconv.ParseInt(err.Error(), 10, 64)\n\t\t\t\tcnt += int(v)\n\t\t\t\treturn ctx\n\t\t\t}).\n\t\t\tOnStartWithStreamInputFn(func(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context {\n\t\t\t\tfor {\n\t\t\t\t\ti, err := input.Recv()\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\n\t\t\t\t\tcnt += i.(int)\n\t\t\t\t}\n\n\t\t\t\tinput.Close()\n\t\t\t\treturn ctx\n\t\t\t}).\n\t\t\tOnEndWithStreamOutputFn(func(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context {\n\t\t\t\tfor {\n\t\t\t\t\to, err := output.Recv()\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\n\t\t\t\t\tcnt += o.(int)\n\t\t\t\t}\n\n\t\t\t\toutput.Close()\n\t\t\t\treturn ctx\n\t\t\t}).Build()\n\n\t\tctx = InitCallbacks(ctx, nil, hb)\n\n\t\tctx = OnStart(ctx, 1)\n\t\tctx = OnEnd(ctx, 2)\n\t\tctx = OnError(ctx, fmt.Errorf(\"3\"))\n\t\tisr, isw := schema.Pipe[int](2)\n\t\tgo func() {\n\t\t\tfor i := 0; i < 10; i++ {\n\t\t\t\tisw.Send(i, nil)\n\t\t\t}\n\t\t\tisw.Close()\n\t\t}()\n\n\t\tctx = ReuseHandlers(ctx, &RunInfo{})\n\t\tvar nisr *schema.StreamReader[int]\n\t\tctx, nisr = OnStartWithStreamInput(ctx, isr)\n\t\tj := 0\n\t\tfor {\n\t\t\ti, err := nisr.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tassert.NoError(t, err)\n\t\t\tassert.Equal(t, j, i)\n\t\t\tj++\n\t\t\tcnt += i\n\t\t}\n\t\tnisr.Close()\n\n\t\tosr, osw := schema.Pipe[int](2)\n\t\tgo func() {\n\t\t\tfor i := 0; i < 10; i++ {\n\t\t\t\tosw.Send(i, nil)\n\t\t\t}\n\t\t\tosw.Close()\n\t\t}()\n\n\t\tvar nosr *schema.StreamReader[int]\n\t\tctx, nosr = OnEndWithStreamOutput(ctx, osr)\n\t\tj = 0\n\t\tfor {\n\t\t\ti, err := nosr.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tassert.NoError(t, err)\n\t\t\tassert.Equal(t, j, i)\n\t\t\tj++\n\t\t\tcnt += i\n\t\t}\n\t\tnosr.Close()\n\t\tassert.Equal(t, 186, cnt)\n\t})\n}\n\nfunc TestGlobalCallbacksRepeated(t *testing.T) {\n\ttimes := 0\n\ttestHandler := NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\ttimes++\n\t\treturn ctx\n\t}).Build()\n\tcallbacks.GlobalHandlers = append(callbacks.GlobalHandlers, testHandler)\n\n\tctx := context.Background()\n\tctx = callbacks.AppendHandlers(ctx, &RunInfo{})\n\tctx = callbacks.AppendHandlers(ctx, &RunInfo{})\n\n\tcallbacks.On(ctx, \"test\", callbacks.OnStartHandle[string], TimingOnStart, true)\n\tassert.Equal(t, times, 1)\n}\n\nfunc TestEnsureRunInfo(t *testing.T) {\n\tctx := context.Background()\n\n\tvar name, typ, comp string\n\tctx = InitCallbacks(ctx, &RunInfo{Name: \"name\", Type: \"type\", Component: \"component\"}, NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context {\n\t\tname = info.Name\n\t\ttyp = info.Type\n\t\tcomp = string(info.Component)\n\t\treturn ctx\n\t}).Build())\n\n\tctx = OnStart(ctx, \"\")\n\tassert.Equal(t, \"name\", name)\n\tassert.Equal(t, \"type\", typ)\n\tassert.Equal(t, \"component\", comp)\n\tctx2 := EnsureRunInfo(ctx, \"type2\", \"component2\")\n\tOnStart(ctx2, \"\")\n\tassert.Equal(t, \"\", name)\n\tassert.Equal(t, \"type2\", typ)\n\tassert.Equal(t, \"component2\", comp)\n\n\t// EnsureRunInfo on an empty Context\n\tAppendGlobalHandlers(NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context {\n\t\ttyp = info.Type\n\t\tcomp = string(info.Component)\n\t\treturn ctx\n\t}).Build())\n\tctx3 := EnsureRunInfo(context.Background(), \"type3\", \"component3\")\n\tOnStart(ctx3, 0)\n\tassert.Equal(t, \"type3\", typ)\n\tassert.Equal(t, \"component3\", comp)\n\tcallbacks.GlobalHandlers = []Handler{}\n}\n\nfunc TestNesting(t *testing.T) {\n\tctx := context.Background()\n\tcb := &myCallback{t: t}\n\tctx = InitCallbacks(ctx, &RunInfo{\n\t\tName: \"test\",\n\t}, cb)\n\n\t// jumped\n\tctx1 := OnStart(ctx, 0)\n\tctx2 := OnStart(ctx1, 1)\n\tOnEnd(ctx2, 1)\n\tOnEnd(ctx1, 0)\n\tassert.Equal(t, 4, cb.times)\n\n\t// reused\n\tcb.times = 0\n\tctx1 = OnStart(ctx, 0)\n\tctx2 = ReuseHandlers(ctx1, &RunInfo{Name: \"test2\"})\n\tctx3 := OnStart(ctx2, 1)\n\tOnEnd(ctx3, 1)\n\tOnEnd(ctx1, 0)\n\tassert.Equal(t, 4, cb.times)\n\n}\n\nfunc TestReuseHandlersOnEmptyCtx(t *testing.T) {\n\tcallbacks.GlobalHandlers = []Handler{}\n\tcb := &myCallback{t: t}\n\tAppendGlobalHandlers(cb)\n\tctx := ReuseHandlers(context.Background(), &RunInfo{Name: \"test\"})\n\tOnStart(ctx, 0)\n\tassert.Equal(t, 1, cb.times)\n}\n\nfunc TestAppendHandlersTwiceOnSameCtx(t *testing.T) {\n\tcallbacks.GlobalHandlers = []Handler{}\n\tcb := &myCallback{t: t}\n\tcb1 := &myCallback{t: t}\n\tcb2 := &myCallback{t: t}\n\tctx := InitCallbacks(context.Background(), &RunInfo{Name: \"test\"}, cb)\n\tctx1 := callbacks.AppendHandlers(ctx, &RunInfo{Name: \"test\"}, cb1)\n\tctx2 := callbacks.AppendHandlers(ctx, &RunInfo{Name: \"test\"}, cb2)\n\tOnStart(ctx1, 0)\n\tOnStart(ctx2, 0)\n\tassert.Equal(t, 2, cb.times)\n\tassert.Equal(t, 1, cb1.times)\n\tassert.Equal(t, 1, cb2.times)\n}\n\ntype myCallback struct {\n\tt     *testing.T\n\ttimes int\n}\n\nfunc (m *myCallback) OnStart(ctx context.Context, info *RunInfo, input CallbackInput) context.Context {\n\tm.times++\n\tif info == nil {\n\t\tassert.Equal(m.t, 2, m.times)\n\t\treturn ctx\n\t}\n\tif info.Name == \"test\" {\n\t\tassert.Equal(m.t, 0, input)\n\t} else {\n\t\tassert.Equal(m.t, 1, input)\n\t}\n\n\treturn ctx\n}\n\nfunc (m *myCallback) OnEnd(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context {\n\tm.times++\n\tif info == nil {\n\t\tassert.Equal(m.t, 3, m.times)\n\t\treturn ctx\n\t}\n\tif info.Name == \"test\" {\n\t\tassert.Equal(m.t, 0, output)\n\t} else {\n\t\tassert.Equal(m.t, 1, output)\n\t}\n\treturn ctx\n}\n\nfunc (m *myCallback) OnError(ctx context.Context, info *RunInfo, err error) context.Context {\n\tpanic(\"implement me\")\n}\n\nfunc (m *myCallback) OnStartWithStreamInput(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context {\n\tpanic(\"implement me\")\n}\n\nfunc (m *myCallback) OnEndWithStreamOutput(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context {\n\tpanic(\"implement me\")\n}\n"
  },
  {
    "path": "callbacks/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package callbacks provides observability hooks for component execution in Eino.\n//\n// Callbacks fire at five lifecycle timings around every component invocation:\n//   - [TimingOnStart] / [TimingOnEnd]: non-streaming input and output.\n//   - [TimingOnStartWithStreamInput] / [TimingOnEndWithStreamOutput]: streaming\n//     variants — handlers receive a copy of the stream and MUST close it.\n//   - [TimingOnError]: component returned a non-nil error (stream-internal\n//     errors are NOT reported here).\n//\n// # Attaching Handlers\n//\n// Global handlers (observe every node in every graph):\n//\n//\tcallbacks.AppendGlobalHandlers(myHandler) // call once, at startup — NOT thread-safe\n//\n// Per-invocation handlers (observe one graph run):\n//\n//\trunnable.Invoke(ctx, input, compose.WithCallbacks(myHandler))\n//\n// Target a specific node:\n//\n//\tcompose.WithCallbacks(myHandler).DesignateNode(\"nodeName\")\n//\n// Handler inheritance: if the context passed to a graph run already carries\n// handlers (e.g. from a parent graph), those handlers are inherited by the\n// entire child run automatically.\n//\n// # Building Handlers\n//\n// Option 1 — [NewHandlerBuilder]: register raw functions for the timings you\n// need. Input/output are untyped; use the component package's ConvCallbackInput\n// helper to cast to a concrete type:\n//\n//\thandler := callbacks.NewHandlerBuilder().\n//\t\tOnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context {\n//\t\t\t// Handle component start\n//\t\t\treturn ctx\n//\t\t}).\n//\t\tOnEndFn(func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context {\n//\t\t\t// Handle component end\n//\t\t\treturn ctx\n//\t\t}).\n//\t\tOnErrorFn(func(ctx context.Context, info *RunInfo, err error) context.Context {\n//\t\t\t// Handle component error\n//\t\t\treturn ctx\n//\t\t}).\n//\t\tOnStartWithStreamInputFn(func(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context {\n//\t\t\tdefer input.Close() // MUST close — failure causes pipeline goroutine leak\n//\t\t\treturn ctx\n//\t\t}).\n//\t\tOnEndWithStreamOutputFn(func(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context {\n//\t\t\tdefer output.Close() // MUST close\n//\t\t\treturn ctx\n//\t\t}).\n//\t\tBuild()\n//\n// Option 2 — utils/callbacks.NewHandlerHelper: dispatches by component type, so\n// each handler function receives the concrete typed input/output directly:\n//\n//\thandler := callbacks.NewHandlerHelper().\n//\t\tChatModel(&model.CallbackHandler{\n//\t\t\tOnStart: func(ctx context.Context, info *RunInfo, input *model.CallbackInput) context.Context {\n//\t\t\t\tlog.Printf(\"Model started: %s, messages: %d\", info.Name, len(input.Messages))\n//\t\t\t\treturn ctx\n//\t\t\t},\n//\t\t}).\n//\t\tPrompt(&prompt.CallbackHandler{\n//\t\t\tOnEnd: func(ctx context.Context, info *RunInfo, output *prompt.CallbackOutput) context.Context {\n//\t\t\t\tlog.Printf(\"Prompt completed\")\n//\t\t\t\treturn ctx\n//\t\t\t},\n//\t\t}).\n//\t\tHandler()\n//\n// # Passing State Within a Handler\n//\n// The ctx returned by one timing is passed to the next timing of the SAME\n// handler, enabling OnStart→OnEnd state transfer via context.WithValue:\n//\n//\tNewHandlerBuilder().\n//\t\tOnStartFn(func(ctx context.Context, info *RunInfo, _ CallbackInput) context.Context {\n//\t\t\treturn context.WithValue(ctx, startTimeKey{}, time.Now())\n//\t\t}).\n//\t\tOnEndFn(func(ctx context.Context, info *RunInfo, _ CallbackOutput) context.Context {\n//\t\t\tstart := ctx.Value(startTimeKey{}).(time.Time)\n//\t\t\tlog.Printf(\"duration: %v\", time.Since(start))\n//\t\t\treturn ctx\n//\t\t}).Build()\n//\n// Between DIFFERENT handlers there is no guaranteed execution order and no\n// context chain. To share state between handlers, store it in a\n// concurrency-safe variable in the outermost context instead.\n//\n// # Common Pitfalls\n//\n//   - Stream copies must be closed: when N handlers register for a streaming\n//     timing, the stream is copied N+1 times (one per handler + one for\n//     downstream). If any handler's copy is not closed, the original stream\n//     cannot be freed and the entire pipeline leaks.\n//\n//   - Do NOT mutate Input/Output: all downstream nodes and handlers share the\n//     same pointer. Mutations cause data races in concurrent graph execution.\n//\n//   - AppendGlobalHandlers is NOT thread-safe: call only during initialization,\n//     never concurrently with graph execution.\n//\n//   - Stream errors are invisible to OnError: errors that occur while a\n//     consumer reads from a StreamReader are not routed through OnError.\n//\n//   - RunInfo may be nil: always nil-check before dereferencing in handlers,\n//     especially when a component is used standalone outside a graph without\n//     InitCallbacks being called.\npackage callbacks\n"
  },
  {
    "path": "callbacks/handler_builder.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// HandlerBuilder constructs a [Handler] by registering callback functions for\n// individual timings. Only set the timings you care about; the built handler\n// implements [TimingChecker] and returns false for unregistered timings, so\n// the framework skips those timings with no overhead.\n//\n// The input/output values are untyped (CallbackInput / CallbackOutput). To\n// work with a specific component's payload, use the component package's\n// ConvCallbackInput / ConvCallbackOutput helpers inside your function. For a\n// higher-level API that dispatches by component type automatically, see\n// utils/callbacks.NewHandlerHelper.\n//\n// Example:\n//\n//\thandler := callbacks.NewHandlerBuilder().\n//\t    OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n//\t        mi := model.ConvCallbackInput(input)\n//\t        if mi != nil {\n//\t            log.Printf(\"[%s] model start: %d messages\", info.Name, len(mi.Messages))\n//\t        }\n//\t        return ctx\n//\t    }).\n//\t    OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n//\t        mo := model.ConvCallbackOutput(output)\n//\t        if mo != nil && mo.Message.ResponseMeta != nil {\n//\t            log.Printf(\"[%s] tokens: %d\", info.Name, mo.Message.ResponseMeta.Usage.TotalTokens)\n//\t        }\n//\t        return ctx\n//\t    }).\n//\t    Build()\ntype HandlerBuilder struct {\n\tonStartFn                func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context\n\tonEndFn                  func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context\n\tonErrorFn                func(ctx context.Context, info *RunInfo, err error) context.Context\n\tonStartWithStreamInputFn func(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context\n\tonEndWithStreamOutputFn  func(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context\n}\n\ntype handlerImpl struct {\n\tHandlerBuilder\n}\n\nfunc (hb *handlerImpl) OnStart(ctx context.Context, info *RunInfo, input CallbackInput) context.Context {\n\treturn hb.onStartFn(ctx, info, input)\n}\n\nfunc (hb *handlerImpl) OnEnd(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context {\n\treturn hb.onEndFn(ctx, info, output)\n}\n\nfunc (hb *handlerImpl) OnError(ctx context.Context, info *RunInfo, err error) context.Context {\n\treturn hb.onErrorFn(ctx, info, err)\n}\n\nfunc (hb *handlerImpl) OnStartWithStreamInput(ctx context.Context, info *RunInfo,\n\tinput *schema.StreamReader[CallbackInput]) context.Context {\n\n\treturn hb.onStartWithStreamInputFn(ctx, info, input)\n}\n\nfunc (hb *handlerImpl) OnEndWithStreamOutput(ctx context.Context, info *RunInfo,\n\toutput *schema.StreamReader[CallbackOutput]) context.Context {\n\n\treturn hb.onEndWithStreamOutputFn(ctx, info, output)\n}\n\nfunc (hb *handlerImpl) Needed(_ context.Context, _ *RunInfo, timing CallbackTiming) bool {\n\tswitch timing {\n\tcase TimingOnStart:\n\t\treturn hb.onStartFn != nil\n\tcase TimingOnEnd:\n\t\treturn hb.onEndFn != nil\n\tcase TimingOnError:\n\t\treturn hb.onErrorFn != nil\n\tcase TimingOnStartWithStreamInput:\n\t\treturn hb.onStartWithStreamInputFn != nil\n\tcase TimingOnEndWithStreamOutput:\n\t\treturn hb.onEndWithStreamOutputFn != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// NewHandlerBuilder creates and returns a new HandlerBuilder instance.\n// HandlerBuilder is used to construct a Handler with custom callback functions\nfunc NewHandlerBuilder() *HandlerBuilder {\n\treturn &HandlerBuilder{}\n}\n\n// OnStartFn sets the handler for the start timing.\nfunc (hb *HandlerBuilder) OnStartFn(\n\tfn func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context) *HandlerBuilder {\n\n\thb.onStartFn = fn\n\treturn hb\n}\n\n// OnEndFn sets the handler for the end timing.\nfunc (hb *HandlerBuilder) OnEndFn(\n\tfn func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context) *HandlerBuilder {\n\n\thb.onEndFn = fn\n\treturn hb\n}\n\n// OnErrorFn sets the handler for the error timing.\nfunc (hb *HandlerBuilder) OnErrorFn(\n\tfn func(ctx context.Context, info *RunInfo, err error) context.Context) *HandlerBuilder {\n\n\thb.onErrorFn = fn\n\treturn hb\n}\n\n// OnStartWithStreamInputFn sets the callback invoked when a component receives\n// streaming input. The handler receives a [*schema.StreamReader] that is a\n// private copy; it MUST close the reader after consuming it to avoid goroutine\n// and memory leaks.\nfunc (hb *HandlerBuilder) OnStartWithStreamInputFn(\n\tfn func(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context) *HandlerBuilder {\n\n\thb.onStartWithStreamInputFn = fn\n\treturn hb\n}\n\n// OnEndWithStreamOutputFn sets the callback invoked when a component produces\n// streaming output. Like OnStartWithStreamInputFn, the handler receives a\n// private copy of the stream and MUST close it after reading to prevent\n// goroutine and memory leaks. This is the right place to implement streaming\n// token-usage accounting or streaming log capture.\nfunc (hb *HandlerBuilder) OnEndWithStreamOutputFn(\n\tfn func(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context) *HandlerBuilder {\n\n\thb.onEndWithStreamOutputFn = fn\n\treturn hb\n}\n\n// Build returns a Handler with the functions set in the builder.\nfunc (hb *HandlerBuilder) Build() Handler {\n\treturn &handlerImpl{*hb}\n}\n"
  },
  {
    "path": "callbacks/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport (\n\t\"github.com/cloudwego/eino/internal/callbacks\"\n)\n\n// RunInfo describes the entity that triggered a callback. Always nil-check\n// before dereferencing — a component that calls OnStart without first calling\n// EnsureRunInfo or InitCallbacks will leave RunInfo absent in the context.\n//\n// Fields:\n//   - Name: business-meaningful name specified by the user. For nodes in a\n//     graph this is the node name (compose.WithNodeName). For standalone\n//     components it must be set explicitly via [InitCallbacks] or\n//     [ReuseHandlers]; it is empty string if not set.\n//   - Type: implementation identity, e.g. \"OpenAI\". Set by the component via\n//     [components.Typer]; falls back to reflection (struct/func name) if the\n//     interface is not implemented. Empty for Graph itself.\n//   - Component: category constant, e.g. components.ComponentOfChatModel.\n//     Fixed value \"Lambda\" for lambdas, \"Graph\"/\"Chain\"/\"Workflow\" for graphs.\n//     Use this to branch on component kind without caring about implementation.\n//\n// Handlers should filter using RunInfo rather than assuming a fixed execution\n// order — there is no guaranteed ordering between different Handlers.\ntype RunInfo = callbacks.RunInfo\n\n// CallbackInput is the value passed to OnStart and OnStartWithStreamInput\n// handlers. The concrete type is defined by the component — for example,\n// ChatModel callbacks carry *model.CallbackInput. Use the component package's\n// ConvCallbackInput helper (e.g. model.ConvCallbackInput) to cast safely; it\n// returns nil if the type does not match, so you can ignore irrelevant\n// component types:\n//\n//\tmodelInput := model.ConvCallbackInput(in)\n//\tif modelInput == nil {\n//\t    return ctx // not a model invocation, skip\n//\t}\n//\tlog.Printf(\"prompt: %v\", modelInput.Messages)\ntype CallbackInput = callbacks.CallbackInput\n\n// CallbackOutput is the value passed to OnEnd and OnEndWithStreamOutput\n// handlers. Like CallbackInput, the concrete type is component-defined.\n// Use the component package's ConvCallbackOutput helper to cast safely.\ntype CallbackOutput = callbacks.CallbackOutput\n\n// Handler is the unified callback handler interface. Implement all five\n// methods (OnStart, OnEnd, OnError, OnStartWithStreamInput,\n// OnEndWithStreamOutput) or use [NewHandlerBuilder] to set only the timings\n// you care about.\n//\n// Each method receives the context returned by the previous timing of the\n// SAME handler, which lets a single handler pass state between its OnStart\n// and OnEnd calls via context.WithValue. There is NO guaranteed execution\n// order between DIFFERENT handlers, and the context chain does not flow\n// from one handler to the next — do not rely on handler ordering.\n//\n// Implement [TimingChecker] (the Needed method) on your handler so the\n// framework can skip timings you have not registered; this avoids unnecessary\n// stream copies and goroutine allocations on every component invocation.\n//\n// Stream handlers (OnStartWithStreamInput, OnEndWithStreamOutput) receive a\n// [*schema.StreamReader] that has already been copied; they MUST close their\n// copy after reading. If any handler's copy is not closed, the original stream\n// cannot be freed, causing a goroutine/memory leak for the entire pipeline.\n//\n// Important: do NOT mutate the Input or Output values. All downstream nodes\n// and handlers share the same pointer (direct assignment, not a deep copy).\n// Mutations cause data races in concurrent graph execution.\ntype Handler = callbacks.Handler\n\n// InitCallbackHandlers sets the global callback handlers.\n// It should be called BEFORE any callback handler by user.\n// It's useful when you want to inject some basic callbacks to all nodes.\n// Deprecated: Use AppendGlobalHandlers instead.\nfunc InitCallbackHandlers(handlers []Handler) {\n\tcallbacks.GlobalHandlers = handlers\n}\n\n// AppendGlobalHandlers appends handlers to the process-wide list of callback\n// handlers. Global handlers run before per-invocation handlers provided via\n// compose.WithCallbacks, giving them higher priority for instrumentation that\n// must observe every component invocation (e.g. distributed tracing, metrics).\n//\n// This function is NOT thread-safe. Call it once during program initialization\n// (e.g. in main or TestMain), before any graph executions begin.\n// Calling it concurrently with ongoing graph executions leads to data races.\nfunc AppendGlobalHandlers(handlers ...Handler) {\n\tcallbacks.GlobalHandlers = append(callbacks.GlobalHandlers, handlers...)\n}\n\n// CallbackTiming enumerates the lifecycle moments at which a callback handler\n// is invoked. Implement [TimingChecker] on your handler and return false for\n// timings you do not handle, so the framework skips the overhead of stream\n// copying and goroutine spawning for those timings.\ntype CallbackTiming = callbacks.CallbackTiming\n\n// Callback timing constants.\nconst (\n\t// TimingOnStart fires just before the component begins processing.\n\t// Receives a fully-formed input value (non-streaming).\n\tTimingOnStart CallbackTiming = iota\n\t// TimingOnEnd fires after the component returns a result successfully.\n\t// Receives the output value. Only fires on success — not on error.\n\tTimingOnEnd\n\t// TimingOnError fires when the component returns a non-nil error.\n\t// Stream errors (mid-stream panics) are NOT reported here; they surface\n\t// as errors inside the stream reader.\n\tTimingOnError\n\t// TimingOnStartWithStreamInput fires when the component receives a\n\t// streaming input (Collect / Transform paradigms). The handler receives a\n\t// copy of the input stream and must close it after reading.\n\tTimingOnStartWithStreamInput\n\t// TimingOnEndWithStreamOutput fires after the component returns a\n\t// streaming output (Stream / Transform paradigms). The handler receives a\n\t// copy of the output stream and must close it after reading. This is\n\t// typically where you implement streaming metrics or logging.\n\tTimingOnEndWithStreamOutput\n)\n\n// TimingChecker is an optional interface for [Handler] implementations.\n// When a handler implements Needed, the framework calls it before each\n// component invocation to decide whether to set up callback infrastructure\n// (stream copying, goroutine allocation) for that timing. Returning false\n// avoids unnecessary overhead.\n//\n// Handlers built with [NewHandlerBuilder] or\n// utils/callbacks.NewHandlerHelper automatically implement TimingChecker\n// based on which callback functions were set.\ntype TimingChecker = callbacks.TimingChecker\n"
  },
  {
    "path": "callbacks/interface_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/internal/callbacks\"\n)\n\nfunc TestAppendGlobalHandlers(t *testing.T) {\n\t// Clear global handlers before test\n\tcallbacks.GlobalHandlers = nil\n\n\t// Create test handlers\n\thandler1 := NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context {\n\t\t\treturn ctx\n\t\t}).Build()\n\thandler2 := NewHandlerBuilder().\n\t\tOnEndFn(func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context {\n\t\t\treturn ctx\n\t\t}).Build()\n\n\t// Test appending first handler\n\tAppendGlobalHandlers(handler1)\n\tassert.Equal(t, 1, len(callbacks.GlobalHandlers))\n\tassert.Contains(t, callbacks.GlobalHandlers, handler1)\n\n\t// Test appending second handler\n\tAppendGlobalHandlers(handler2)\n\tassert.Equal(t, 2, len(callbacks.GlobalHandlers))\n\tassert.Contains(t, callbacks.GlobalHandlers, handler1)\n\tassert.Contains(t, callbacks.GlobalHandlers, handler2)\n\n\t// Test appending nil handler\n\tAppendGlobalHandlers([]Handler{}...)\n\tassert.Equal(t, 2, len(callbacks.GlobalHandlers))\n}\n"
  },
  {
    "path": "components/document/callback_extra_loader.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage document\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// LoaderCallbackInput is the input for the loader callback.\ntype LoaderCallbackInput struct {\n\t// Source is the source of the documents.\n\tSource Source\n\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// LoaderCallbackOutput is the output for the loader callback.\ntype LoaderCallbackOutput struct {\n\t// Source is the source of the documents.\n\tSource Source\n\n\t// Docs is the documents to be loaded.\n\tDocs []*schema.Document\n\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// ConvLoaderCallbackInput converts the callback input to the loader callback input.\nfunc ConvLoaderCallbackInput(src callbacks.CallbackInput) *LoaderCallbackInput {\n\tswitch t := src.(type) {\n\tcase *LoaderCallbackInput:\n\t\treturn t\n\tcase Source:\n\t\treturn &LoaderCallbackInput{\n\t\t\tSource: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// ConvLoaderCallbackOutput converts the callback output to the loader callback output.\nfunc ConvLoaderCallbackOutput(src callbacks.CallbackOutput) *LoaderCallbackOutput {\n\tswitch t := src.(type) {\n\tcase *LoaderCallbackOutput:\n\t\treturn t\n\tcase []*schema.Document:\n\t\treturn &LoaderCallbackOutput{\n\t\t\tDocs: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n"
  },
  {
    "path": "components/document/callback_extra_transformer.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage document\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// TransformerCallbackInput is the input for the transformer callback.\ntype TransformerCallbackInput struct {\n\t// Input is the input documents.\n\tInput []*schema.Document\n\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// TransformerCallbackOutput is the output for the transformer callback.\ntype TransformerCallbackOutput struct {\n\t// Output is the output documents.\n\tOutput []*schema.Document\n\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// ConvTransformerCallbackInput converts the callback input to the transformer callback input.\nfunc ConvTransformerCallbackInput(src callbacks.CallbackInput) *TransformerCallbackInput {\n\tswitch t := src.(type) {\n\tcase *TransformerCallbackInput:\n\t\treturn t\n\tcase []*schema.Document:\n\t\treturn &TransformerCallbackInput{\n\t\t\tInput: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// ConvTransformerCallbackOutput converts the callback output to the transformer callback output.\nfunc ConvTransformerCallbackOutput(src callbacks.CallbackOutput) *TransformerCallbackOutput {\n\tswitch t := src.(type) {\n\tcase *TransformerCallbackOutput:\n\t\treturn t\n\tcase []*schema.Document:\n\t\treturn &TransformerCallbackOutput{\n\t\t\tOutput: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n"
  },
  {
    "path": "components/document/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package document defines the Loader and Transformer component interfaces\n// for ingesting and processing documents in an eino pipeline.\n//\n// # Components\n//\n//   - [Loader]: reads raw content from an external source (file, URL, S3, …)\n//     and returns [schema.Document] values. Parsing is typically delegated to\n//     a [parser.Parser] configured on the loader.\n//   - [Transformer]: takes a slice of [schema.Document] values and transforms\n//     them — splitting, filtering, merging, re-ranking, etc.\n//\n// Concrete implementations live in eino-ext:\n//\n//\tgithub.com/cloudwego/eino-ext/components/document/\n//\n// # Document Metadata\n//\n// [schema.Document].MetaData is the primary mechanism for carrying contextual\n// information (source URI, scores, chunk indices, embeddings) through the\n// pipeline. Transformers should preserve existing metadata and merge rather\n// than replace when adding their own keys.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/document_loader_guide/\n// See https://www.cloudwego.io/docs/eino/core_modules/components/document_transformer_guide/\npackage document\n"
  },
  {
    "path": "components/document/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage document\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Source identifies the external location of a document.\n// URI can be a local file path or a remote URL reachable by the loader.\ntype Source struct {\n\tURI string\n}\n\n//go:generate  mockgen -destination ../../internal/mock/components/document/document_mock.go --package document -source interface.go\n\n// Loader reads raw content from an external source and returns it as a slice\n// of [schema.Document] values.\n//\n// The Source.URI may be a local file path or a remote URL. The loader is\n// responsible for fetching the raw bytes; actual format parsing is typically\n// delegated to a [parser.Parser] configured on the loader via\n// [WithParserOptions].\n//\n// Document metadata ([schema.Document].MetaData) should be populated with at\n// least the source URI so that downstream nodes can trace document provenance.\ntype Loader interface {\n\tLoad(ctx context.Context, src Source, opts ...LoaderOption) ([]*schema.Document, error)\n}\n\n// Transformer converts a slice of [schema.Document] values into another slice,\n// applying operations such as splitting, filtering, merging, or re-ranking.\n//\n// Implementations should preserve existing MetaData keys and merge rather than\n// replace when adding their own metadata. Downstream nodes (e.g. Indexer,\n// Retriever) may depend on metadata set by earlier pipeline stages.\ntype Transformer interface {\n\tTransform(ctx context.Context, src []*schema.Document, opts ...TransformerOption) ([]*schema.Document, error)\n}\n"
  },
  {
    "path": "components/document/option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage document\n\nimport \"github.com/cloudwego/eino/components/document/parser\"\n\n// LoaderOptions configures document loaders, including parser options.\ntype LoaderOptions struct {\n\tParserOptions []parser.Option\n}\n\n// LoaderOption defines call option for Loader component, which is part of the component interface signature.\n// Each Loader implementation could define its own options struct and option funcs within its own package,\n// then wrap the impl specific option funcs into this type, before passing to Load.\ntype LoaderOption struct {\n\tapply func(opts *LoaderOptions)\n\n\timplSpecificOptFn any\n}\n\n// WrapLoaderImplSpecificOptFn wraps the impl specific option functions into LoaderOption type.\n// T: the type of the impl specific options struct.\n// Loader implementations are required to use this function to convert its own option functions into the unified LoaderOption type.\n// For example, if the Loader impl defines its own options struct:\n//\n//\ttype customOptions struct {\n//\t    conf string\n//\t}\n//\n// Then the impl needs to provide an option function as such:\n//\n//\tfunc WithConf(conf string) Option {\n//\t    return WrapLoaderImplSpecificOptFn(func(o *customOptions) {\n//\t\t\to.conf = conf\n//\t\t}\n//\t}\nfunc WrapLoaderImplSpecificOptFn[T any](optFn func(*T)) LoaderOption {\n\treturn LoaderOption{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetLoaderImplSpecificOptions provides Loader author the ability to extract their own custom options from the unified LoaderOption type.\n// T: the type of the impl specific options struct.\n// This function should be used within the Loader implementation's Load function.\n// It is recommended to provide a base T as the first argument, within which the Loader author can provide default values for the impl specific options.\n// eg.\n//\n//\tmyOption := &MyOption{\n//\t\tField1: \"default_value\",\n//\t}\n//\tmyOption := loader.GetLoaderImplSpecificOptions(myOption, opts...)\nfunc GetLoaderImplSpecificOptions[T any](base *T, opts ...LoaderOption) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\ts, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\ts(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n\n// GetLoaderCommonOptions extract loader Options from Option list, optionally providing a base Options with default values.\nfunc GetLoaderCommonOptions(base *LoaderOptions, opts ...LoaderOption) *LoaderOptions {\n\tif base == nil {\n\t\tbase = &LoaderOptions{}\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.apply != nil {\n\t\t\topt.apply(base)\n\t\t}\n\t}\n\n\treturn base\n}\n\n// WithParserOptions attaches parser options to a loader request.\nfunc WithParserOptions(opts ...parser.Option) LoaderOption {\n\treturn LoaderOption{\n\t\tapply: func(o *LoaderOptions) {\n\t\t\to.ParserOptions = opts\n\t\t},\n\t}\n}\n\n// TransformerOption defines call option for Transformer component, which is part of the component interface signature.\n// Each Transformer implementation could define its own options struct and option funcs within its own package,\n// then wrap the impl specific option funcs into this type, before passing to Transform.\ntype TransformerOption struct {\n\timplSpecificOptFn any\n}\n\n// WrapTransformerImplSpecificOptFn wraps the impl specific option functions into TransformerOption type.\n// T: the type of the impl specific options struct.\n// Transformer implementations are required to use this function to convert its own option functions into the unified TransformerOption type.\n// For example, if the Transformer impl defines its own options struct:\n//\n//\ttype customOptions struct {\n//\t    conf string\n//\t}\n//\n// Then the impl needs to provide an option function as such:\n//\n//\tfunc WithConf(conf string) TransformerOption {\n//\t    return WrapTransformerImplSpecificOptFn(func(o *customOptions) {\n//\t\t\to.conf = conf\n//\t\t}\n//\t}\n//\n// .\nfunc WrapTransformerImplSpecificOptFn[T any](optFn func(*T)) TransformerOption {\n\treturn TransformerOption{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetTransformerImplSpecificOptions provides Transformer author the ability to extract their own custom options from the unified TransformerOption type.\n// T: the type of the impl specific options struct.\n// This function should be used within the Transformer implementation's Transform function.\n// It is recommended to provide a base T as the first argument, within which the Transformer author can provide default values for the impl specific options.\n// eg.\n//\n//\tmyOption := &MyOption{\n//\t\tField1: \"default_value\",\n//\t}\n//\tmyOption := transformer.GetTransformerImplSpecificOptions(myOption, opts...)\nfunc GetTransformerImplSpecificOptions[T any](base *T, opts ...TransformerOption) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\ts, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\ts(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "components/document/option_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage document\n\nimport (\n\t\"testing\"\n\n\t\"github.com/smartystreets/goconvey/convey\"\n\n\t\"github.com/cloudwego/eino/components/document/parser\"\n)\n\nfunc TestImplSpecificOpts(t *testing.T) {\n\ttype implSpecificOptions struct {\n\t\tconf  string\n\t\tindex int\n\t}\n\n\twithConf := func(conf string) func(o *implSpecificOptions) {\n\t\treturn func(o *implSpecificOptions) {\n\t\t\to.conf = conf\n\t\t}\n\t}\n\n\twithIndex := func(index int) func(o *implSpecificOptions) {\n\t\treturn func(o *implSpecificOptions) {\n\t\t\to.index = index\n\t\t}\n\t}\n\n\tconvey.Convey(\"TestLoaderImplSpecificOpts\", t, func() {\n\t\tdocumentOption1 := WrapLoaderImplSpecificOptFn(withConf(\"test_conf\"))\n\t\tdocumentOption2 := WrapLoaderImplSpecificOptFn(withIndex(1))\n\n\t\timplSpecificOpts := GetLoaderImplSpecificOptions(&implSpecificOptions{}, documentOption1, documentOption2)\n\n\t\tconvey.So(implSpecificOpts, convey.ShouldResemble, &implSpecificOptions{\n\t\t\tconf:  \"test_conf\",\n\t\t\tindex: 1,\n\t\t})\n\t})\n\tconvey.Convey(\"TestTransformerImplSpecificOpts\", t, func() {\n\t\tdocumentOption1 := WrapTransformerImplSpecificOptFn(withConf(\"test_conf\"))\n\t\tdocumentOption2 := WrapTransformerImplSpecificOptFn(withIndex(1))\n\n\t\timplSpecificOpts := GetTransformerImplSpecificOptions(&implSpecificOptions{}, documentOption1, documentOption2)\n\n\t\tconvey.So(implSpecificOpts, convey.ShouldResemble, &implSpecificOptions{\n\t\t\tconf:  \"test_conf\",\n\t\t\tindex: 1,\n\t\t})\n\t})\n}\n\nfunc TestCommonOptions(t *testing.T) {\n\tconvey.Convey(\"TestCommonOptions\", t, func() {\n\t\to := &LoaderOptions{ParserOptions: []parser.Option{{}}}\n\t\to1 := GetLoaderCommonOptions(o)\n\t\tconvey.So(len(o1.ParserOptions), convey.ShouldEqual, 1)\n\n\t\to2 := GetLoaderCommonOptions(o, WithParserOptions(parser.Option{}, parser.Option{}))\n\t\tconvey.So(len(o2.ParserOptions), convey.ShouldEqual, 2)\n\t})\n}\n"
  },
  {
    "path": "components/document/parser/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package parser defines the Parser interface for converting raw byte streams\n// into [schema.Document] values.\n//\n// # Overview\n//\n// A Parser is not a standalone pipeline component — it is used inside a\n// [document.Loader] to handle format-specific decoding. The loader fetches\n// raw bytes; the parser converts them into documents.\n//\n// # Built-in Implementations\n//\n//   - TextParser: treats the entire reader as plain text, one document per call\n//   - ExtParser: selects a parser by file extension (from [Options.URI]), with\n//     a configurable fallback for unknown extensions\n//\n// Use ExtParser when you want format-agnostic loading: pass the source URI\n// via [WithURI] and ExtParser picks the right sub-parser automatically.\n//\n// # Reader Contract\n//\n// The [io.Reader] passed to [Parser.Parse] is consumed during the call —\n// it cannot be read again. Loaders must not reuse the same reader across\n// multiple Parse calls.\n//\n// # Metadata Propagation\n//\n// Use [WithExtraMeta] to attach key-value pairs that are merged into every\n// document's MetaData. This is the standard way to tag documents with source\n// information (URI, content type, etc.) at parse time.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/document_loader_guide/document_parser_interface_guide/\npackage parser\n"
  },
  {
    "path": "components/document/parser/ext_parser.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parser\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"io\"\n\t\"path/filepath\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// ExtParserConfig defines the configuration for the ExtParser.\ntype ExtParserConfig struct {\n\t// ext -> parser.\n\t// eg: map[string]Parser{\n\t// \t\".pdf\": &PDFParser{},\n\t// \t\".md\": &MarkdownParser{},\n\t// }\n\tParsers map[string]Parser\n\n\t// Fallback parser to use when no other parser is found.\n\t// Default is TextParser if not set.\n\tFallbackParser Parser\n}\n\n// ExtParser is a parser that uses the file extension to determine which parser to use.\n// You can register your own parsers by calling RegisterParser.\n// Default parser is TextParser.\n// Note:\n//\n//\tparse 时，是通过 filepath.Ext(uri) 的方式找到对应的 parser，因此使用时需要：\n//\t \t① 必须使用 parser.WithURI 在请求时传入 URI\n//\t \t② URI 必须能通过 filepath.Ext 来解析出符合预期的 ext\n//\n// eg:\n//\n//\tpdf, _ := os.Open(\"./testdata/test.pdf\")\n//\tdocs, err := ExtParser.Parse(ctx, pdf, parser.WithURI(\"./testdata/test.pdf\"))\ntype ExtParser struct {\n\tparsers map[string]Parser\n\n\tfallbackParser Parser\n}\n\n// NewExtParser creates a new ExtParser.\nfunc NewExtParser(ctx context.Context, conf *ExtParserConfig) (*ExtParser, error) {\n\tif conf == nil {\n\t\tconf = &ExtParserConfig{}\n\t}\n\n\tp := &ExtParser{\n\t\tparsers:        conf.Parsers,\n\t\tfallbackParser: conf.FallbackParser,\n\t}\n\n\tif p.fallbackParser == nil {\n\t\tp.fallbackParser = TextParser{}\n\t}\n\n\tif p.parsers == nil {\n\t\tp.parsers = make(map[string]Parser)\n\t}\n\n\treturn p, nil\n}\n\n// GetParsers returns a copy of the registered parsers.\n// It is safe to modify the returned parsers.\nfunc (p *ExtParser) GetParsers() map[string]Parser {\n\tres := make(map[string]Parser, len(p.parsers))\n\tfor k, v := range p.parsers {\n\t\tres[k] = v\n\t}\n\n\treturn res\n}\n\n// Parse parses the given reader and returns a list of documents.\nfunc (p *ExtParser) Parse(ctx context.Context, reader io.Reader, opts ...Option) ([]*schema.Document, error) {\n\topt := GetCommonOptions(&Options{}, opts...)\n\n\text := filepath.Ext(opt.URI)\n\n\tparser, ok := p.parsers[ext]\n\n\tif !ok {\n\t\tparser = p.fallbackParser\n\t}\n\n\tif parser == nil {\n\t\treturn nil, errors.New(\"no parser found for extension \" + ext)\n\t}\n\n\tdocs, err := parser.Parse(ctx, reader, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tfor _, doc := range docs {\n\t\tif doc == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tif doc.MetaData == nil {\n\t\t\tdoc.MetaData = make(map[string]any)\n\t\t}\n\n\t\tfor k, v := range opt.ExtraMeta {\n\t\t\tdoc.MetaData[k] = v\n\t\t}\n\t}\n\n\treturn docs, nil\n}\n"
  },
  {
    "path": "components/document/parser/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parser\n\nimport (\n\t\"context\"\n\t\"io\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Parser converts raw content from an [io.Reader] into [schema.Document] values.\n//\n// Parse may return multiple documents from a single reader (e.g. a PDF with\n// per-page splitting). The reader is consumed during Parse and must not be\n// reused.\n//\n// Parsers are typically not called directly — they are configured on a\n// [document.Loader] and invoked via [document.WithParserOptions].\ntype Parser interface {\n\tParse(ctx context.Context, reader io.Reader, opts ...Option) ([]*schema.Document, error)\n}\n"
  },
  {
    "path": "components/document/parser/option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parser\n\n// Options configures the document parser with source URI and extra metadata.\ntype Options struct {\n\t// uri of source.\n\tURI string\n\n\t// extra metadata will merge to each document.\n\tExtraMeta map[string]any\n}\n\n// Option defines call option for Parser component, which is part of the component interface signature.\n// Each Parser implementation could define its own options struct and option funcs within its own package,\n// then wrap the impl specific option funcs into this type, before passing to Transform.\ntype Option struct {\n\tapply func(opts *Options)\n\n\timplSpecificOptFn any\n}\n\n// WithURI specifies the source URI of the document.\n// It will be used as to select parser in ExtParser.\nfunc WithURI(uri string) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.URI = uri\n\t\t},\n\t}\n}\n\n// WithExtraMeta attaches extra metadata to the parsed document.\nfunc WithExtraMeta(meta map[string]any) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.ExtraMeta = meta\n\t\t},\n\t}\n}\n\n// GetCommonOptions extract parser Options from Option list, optionally providing a base Options with default values.\nfunc GetCommonOptions(base *Options, opts ...Option) *Options {\n\tif base == nil {\n\t\tbase = &Options{}\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.apply != nil {\n\t\t\topt.apply(base)\n\t\t}\n\t}\n\n\treturn base\n}\n\n// WrapImplSpecificOptFn wraps the impl specific option functions into Option type.\n// T: the type of the impl specific options struct.\n// Parser implementations are required to use this function to convert its own option functions into the unified Option type.\n// For example, if the Parser impl defines its own options struct:\n//\n//\ttype customOptions struct {\n//\t    conf string\n//\t}\n//\n// Then the impl needs to provide an option function as such:\n//\n//\tfunc WithConf(conf string) Option {\n//\t    return WrapImplSpecificOptFn(func(o *customOptions) {\n//\t\t\to.conf = conf\n//\t\t}\n//\t}\n//\n// .\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) Option {\n\treturn Option{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetImplSpecificOptions provides Parser author the ability to extract their own custom options from the unified Option type.\n// T: the type of the impl specific options struct.\n// This function should be used within the Parser implementation's Transform function.\n// It is recommended to provide a base T as the first argument, within which the Parser author can provide default values for the impl specific options.\nfunc GetImplSpecificOptions[T any](base *T, opts ...Option) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\ts, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\ts(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "components/document/parser/option_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parser\n\nimport (\n\t\"testing\"\n\n\t\"github.com/smartystreets/goconvey/convey\"\n)\n\nfunc TestImplSpecificOpts(t *testing.T) {\n\ttype implSpecificOptions struct {\n\t\tconf  string\n\t\tindex int\n\t}\n\n\twithConf := func(conf string) func(o *implSpecificOptions) {\n\t\treturn func(o *implSpecificOptions) {\n\t\t\to.conf = conf\n\t\t}\n\t}\n\n\twithIndex := func(index int) func(o *implSpecificOptions) {\n\t\treturn func(o *implSpecificOptions) {\n\t\t\to.index = index\n\t\t}\n\t}\n\n\tconvey.Convey(\"TestImplSpecificOpts\", t, func() {\n\t\tparserOption1 := WrapImplSpecificOptFn(withConf(\"test_conf\"))\n\t\tparserOption2 := WrapImplSpecificOptFn(withIndex(1))\n\n\t\timplSpecificOpts := GetImplSpecificOptions(&implSpecificOptions{}, parserOption1, parserOption2)\n\n\t\tconvey.So(implSpecificOpts, convey.ShouldResemble, &implSpecificOptions{\n\t\t\tconf:  \"test_conf\",\n\t\t\tindex: 1,\n\t\t})\n\t})\n}\n"
  },
  {
    "path": "components/document/parser/parser_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parser\n\nimport (\n\t\"context\"\n\t\"io\"\n\t\"os\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype ParserForTest struct {\n\tmock func() ([]*schema.Document, error)\n}\n\nfunc (p *ParserForTest) Parse(ctx context.Context, reader io.Reader, opts ...Option) ([]*schema.Document, error) {\n\treturn p.mock()\n}\n\nfunc TestParser(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"Test default parser\", func(t *testing.T) {\n\t\tconf := &ExtParserConfig{}\n\n\t\tp, err := NewExtParser(ctx, conf)\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tf, err := os.Open(\"testdata/test.md\")\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tdefer f.Close()\n\n\t\tdocs, err := p.Parse(ctx, f, WithURI(\"testdata/test.md\"))\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tassert.Equal(t, 1, len(docs))\n\t\tassert.Equal(t, \"# Title\\nhello world\", docs[0].Content)\n\t})\n\n\tt.Run(\"test types\", func(t *testing.T) {\n\t\tmockParser := &ParserForTest{\n\t\t\tmock: func() ([]*schema.Document, error) {\n\t\t\t\treturn []*schema.Document{\n\t\t\t\t\t{\n\t\t\t\t\t\tContent: \"hello world\",\n\t\t\t\t\t\tMetaData: map[string]any{\n\t\t\t\t\t\t\t\"type\": \"text\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}, nil\n\t\t\t},\n\t\t}\n\n\t\tconf := &ExtParserConfig{\n\t\t\tParsers: map[string]Parser{\n\t\t\t\t\".md\": mockParser,\n\t\t\t},\n\t\t}\n\n\t\tp, err := NewExtParser(ctx, conf)\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tf, err := os.Open(\"testdata/test.md\")\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tdefer f.Close()\n\n\t\tdocs, err := p.Parse(ctx, f, WithURI(\"x/test.md\"))\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tassert.Equal(t, 1, len(docs))\n\t\tassert.Equal(t, \"hello world\", docs[0].Content)\n\t\tassert.Equal(t, \"text\", docs[0].MetaData[\"type\"])\n\t})\n\n\tt.Run(\"test get parsers\", func(t *testing.T) {\n\t\tp, err := NewExtParser(ctx, &ExtParserConfig{\n\t\t\tParsers: map[string]Parser{\n\t\t\t\t\".md\": &TextParser{},\n\t\t\t},\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tps := p.GetParsers()\n\t\tassert.Equal(t, 1, len(ps))\n\t})\n}\n"
  },
  {
    "path": "components/document/parser/testdata/test.md",
    "content": "# Title\nhello world"
  },
  {
    "path": "components/document/parser/text_parser.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parser\n\nimport (\n\t\"context\"\n\t\"io\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nconst (\n\t// MetaKeySource is the metadata key storing the document's source URI.\n\tMetaKeySource = \"_source\"\n)\n\n// TextParser is a simple parser that reads the text from a reader and returns a single document.\n// eg:\n//\n//\tdocs, err := TextParser.Parse(ctx, strings.NewReader(\"hello world\"))\n//\tfmt.Println(docs[0].Content) // \"hello world\"\ntype TextParser struct{}\n\n// Parse reads the text from a reader and returns a single document.\nfunc (dp TextParser) Parse(ctx context.Context, reader io.Reader, opts ...Option) ([]*schema.Document, error) {\n\tdata, err := io.ReadAll(reader)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\topt := GetCommonOptions(&Options{}, opts...)\n\n\tmeta := make(map[string]any)\n\tmeta[MetaKeySource] = opt.URI\n\n\tfor k, v := range opt.ExtraMeta {\n\t\tmeta[k] = v\n\t}\n\n\tdoc := &schema.Document{\n\t\tContent:  string(data),\n\t\tMetaData: meta,\n\t}\n\n\treturn []*schema.Document{doc}, nil\n}\n"
  },
  {
    "path": "components/embedding/callback_extra.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage embedding\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n)\n\n// TokenUsage is the token usage for the embedding.\ntype TokenUsage struct {\n\t// PromptTokens is the number of prompt tokens.\n\tPromptTokens int\n\t// CompletionTokens is the number of completion tokens.\n\tCompletionTokens int\n\t// TotalTokens is the total number of tokens.\n\tTotalTokens int\n}\n\n// Config is the config for the embedding.\ntype Config struct {\n\t// Model is the model name.\n\tModel string\n\t// EncodingFormat is the encoding format.\n\tEncodingFormat string\n}\n\n// ComponentExtra is the extra information for the embedding.\ntype ComponentExtra struct {\n\t// Config is the config for the embedding.\n\tConfig *Config\n\t// TokenUsage is the token usage for the embedding.\n\tTokenUsage *TokenUsage\n}\n\n// CallbackInput is the input for the embedding callback.\ntype CallbackInput struct {\n\t// Texts is the texts to be embedded.\n\tTexts []string\n\t// Config is the config for the embedding.\n\tConfig *Config\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// CallbackOutput is the output for the embedding callback.\ntype CallbackOutput struct {\n\t// Embeddings is the embeddings.\n\tEmbeddings [][]float64\n\t// Config is the config for creating the embedding.\n\tConfig *Config\n\t// TokenUsage is the token usage for the embedding.\n\tTokenUsage *TokenUsage\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// ConvCallbackInput converts the callback input to the embedding callback input.\nfunc ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput {\n\tswitch t := src.(type) {\n\tcase *CallbackInput:\n\t\treturn t\n\tcase []string:\n\t\treturn &CallbackInput{\n\t\t\tTexts: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// ConvCallbackOutput converts the callback output to the embedding callback output.\nfunc ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {\n\tswitch t := src.(type) {\n\tcase *CallbackOutput:\n\t\treturn t\n\tcase [][]float64:\n\t\treturn &CallbackOutput{\n\t\t\tEmbeddings: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n"
  },
  {
    "path": "components/embedding/callback_extra_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage embedding\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestConvEmbedding(t *testing.T) {\n\tassert.NotNil(t, ConvCallbackInput(&CallbackInput{}))\n\tassert.NotNil(t, ConvCallbackInput([]string{}))\n\tassert.Nil(t, ConvCallbackInput(\"asd\"))\n\n\tassert.NotNil(t, ConvCallbackOutput(&CallbackOutput{}))\n\tassert.NotNil(t, ConvCallbackOutput([][]float64{}))\n\tassert.Nil(t, ConvCallbackOutput(\"asd\"))\n}\n"
  },
  {
    "path": "components/embedding/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package embedding defines the Embedder component interface for converting\n// text into vector representations.\n//\n// # Overview\n//\n// An Embedder converts a batch of strings into dense float vectors. Semantically\n// similar texts produce vectors that are close in the vector space, making\n// embeddings the backbone of semantic search, RAG pipelines, and clustering.\n//\n// Concrete implementations (OpenAI, Ark, Ollama, …) live in eino-ext:\n//\n//\tgithub.com/cloudwego/eino-ext/components/embedding/\n//\n// # Output Format\n//\n// [Embedder.EmbedStrings] returns `[][]float64` where:\n//   - outer index corresponds to the input text at the same position\n//   - inner slice is the embedding vector; its length (dimensions) is fixed by\n//     the model and is the same for every text\n//\n// # Consistency Requirement\n//\n// The same model must be used for both indexing and retrieval. Mixing models\n// produces vectors in different spaces — similarity scores become meaningless\n// and semantic search breaks silently.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/embedding_guide/\npackage embedding\n"
  },
  {
    "path": "components/embedding/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage embedding\n\nimport \"context\"\n\n// Embedder converts a batch of strings into dense vector representations.\n//\n// EmbedStrings returns one vector per input text, in the same order. The\n// vector length (dimensions) is fixed by the underlying model and identical\n// for every text in the batch.\n//\n// The returned [][]float64 maps as:\n//\n//\tembeddings[i]  →  vector for texts[i]\n//\tlen(embeddings[i])  →  model's embedding dimension (e.g. 1536 for ada-002)\n//\n// Both [Indexer] and [Retriever] use an Embedder to convert documents and\n// queries into vectors. They must share the exact same model — mismatched\n// dimensions or model families break semantic similarity.\n//\n//go:generate  mockgen -destination ../../internal/mock/components/embedding/Embedding_mock.go --package embedding -source interface.go\ntype Embedder interface {\n\tEmbedStrings(ctx context.Context, texts []string, opts ...Option) ([][]float64, error) // invoke\n}\n"
  },
  {
    "path": "components/embedding/option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage embedding\n\n// Options is the options for the embedding.\ntype Options struct {\n\t// Model is the model name for the embedding.\n\tModel *string\n}\n\n// Option is a call-time option for an Embedder.\ntype Option struct {\n\tapply func(opts *Options)\n\n\timplSpecificOptFn any\n}\n\n// WithModel is the option to set the model for the embedding.\nfunc WithModel(model string) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.Model = &model\n\t\t},\n\t}\n}\n\n// GetCommonOptions extract embedding Options from Option list, optionally providing a base Options with default values.\n// eg.\n//\n//\tdefaultModelName := \"default_model\"\n//\tembeddingOption := &embedding.Options{\n//\t\tModel: &defaultModelName,\n//\t}\n//\tembeddingOption := embedding.GetCommonOptions(embeddingOption, opts...)\nfunc GetCommonOptions(base *Options, opts ...Option) *Options {\n\tif base == nil {\n\t\tbase = &Options{}\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.apply != nil {\n\t\t\topt.apply(base)\n\t\t}\n\t}\n\n\treturn base\n}\n\n// WrapImplSpecificOptFn wraps an implementation-specific option function so it\n// can be passed alongside standard options. For use by Embedder implementors:\n//\n//\tfunc WithMyParam(v string) embedding.Option {\n//\t    return embedding.WrapImplSpecificOptFn(func(o *MyOptions) {\n//\t        o.MyParam = v\n//\t    })\n//\t}\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) Option {\n\treturn Option{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetImplSpecificOptions extracts implementation-specific options from opts,\n// merging them onto base. Call alongside [GetCommonOptions] inside EmbedStrings:\n//\n//\tfunc (e *MyEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) {\n//\t    common := embedding.GetCommonOptions(nil, opts...)\n//\t    mine  := embedding.GetImplSpecificOptions(&MyOptions{}, opts...)\n//\t    // use common.Model, mine.MyParam, etc.\n//\t}\nfunc GetImplSpecificOptions[T any](base *T, opts ...Option) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\toptFn, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\toptFn(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "components/embedding/option_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage embedding\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestOptions(t *testing.T) {\n\tdefaultModel := \"default_model\"\n\topts := GetCommonOptions(&Options{Model: &defaultModel}, WithModel(\"test_model\"))\n\tassert.NotNil(t, opts.Model)\n\tassert.Equal(t, *opts.Model, \"test_model\")\n}\n"
  },
  {
    "path": "components/indexer/callback_extra.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage indexer\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// CallbackInput is the input for the indexer callback.\ntype CallbackInput struct {\n\t// Docs is the documents to be indexed.\n\tDocs []*schema.Document\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// CallbackOutput is the output for the indexer callback.\ntype CallbackOutput struct {\n\t// IDs is the ids of the indexed documents returned by the indexer.\n\tIDs []string\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// ConvCallbackInput converts the callback input to the indexer callback input.\nfunc ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput {\n\tswitch t := src.(type) {\n\tcase *CallbackInput:\n\t\treturn t\n\tcase []*schema.Document:\n\t\treturn &CallbackInput{\n\t\t\tDocs: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// ConvCallbackOutput converts the callback output to the indexer callback output.\nfunc ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {\n\tswitch t := src.(type) {\n\tcase *CallbackOutput:\n\t\treturn t\n\tcase []string:\n\t\treturn &CallbackOutput{\n\t\t\tIDs: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n"
  },
  {
    "path": "components/indexer/callback_extra_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage indexer\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestConvIndexer(t *testing.T) {\n\tassert.NotNil(t, ConvCallbackInput(&CallbackInput{}))\n\tassert.NotNil(t, ConvCallbackInput([]*schema.Document{}))\n\tassert.Nil(t, ConvCallbackInput(\"asd\"))\n\n\tassert.NotNil(t, ConvCallbackOutput(&CallbackOutput{}))\n\tassert.NotNil(t, ConvCallbackOutput([]string{}))\n\tassert.Nil(t, ConvCallbackOutput(\"asd\"))\n}\n"
  },
  {
    "path": "components/indexer/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package indexer defines the Indexer component interface for storing documents\n// and their vector representations in a backend store.\n//\n// # Overview\n//\n// An Indexer is the write path of a RAG pipeline. It takes [schema.Document]\n// values, optionally generates vector embeddings, and persists them in a\n// backend (vector DB, search engine, etc.) for later retrieval.\n//\n// Concrete implementations (VikingDB, Milvus, Elasticsearch, …) live in\n// eino-ext:\n//\n//\tgithub.com/cloudwego/eino-ext/components/indexer/\n//\n// # Vector Dimension Consistency\n//\n// When using the [Options.Embedding] option, the embedding model must be\n// identical to the one used by the paired [retriever.Retriever]. Mismatched\n// models produce vectors in different spaces — queries will not match stored\n// documents.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/indexer_guide/\npackage indexer\n"
  },
  {
    "path": "components/indexer/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage indexer\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Indexer stores documents (and optionally their vector embeddings) in a\n// backend for later retrieval.\n//\n// Store accepts a batch of [schema.Document] values and returns the IDs\n// assigned to them by the backend. When [Options.Embedding] is provided,\n// the implementation generates vectors before storing — the same embedder\n// must be used by the paired [retriever.Retriever].\n//\n// Use [Options.SubIndexes] to write documents into logical sub-partitions\n// within the same store.\n//\n//go:generate  mockgen -destination ../../internal/mock/components/indexer/indexer_mock.go --package indexer -source interface.go\ntype Indexer interface {\n\t// Store stores the documents and returns their assigned IDs.\n\tStore(ctx context.Context, docs []*schema.Document, opts ...Option) (ids []string, err error) // invoke\n}\n"
  },
  {
    "path": "components/indexer/option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage indexer\n\nimport \"github.com/cloudwego/eino/components/embedding\"\n\n// Options is the options for the indexer.\ntype Options struct {\n\t// SubIndexes is the sub indexes to be indexed.\n\tSubIndexes []string\n\t// Embedding is the embedding component.\n\tEmbedding embedding.Embedder\n}\n\n// WithSubIndexes is the option to set the sub indexes for the indexer.\nfunc WithSubIndexes(subIndexes []string) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.SubIndexes = subIndexes\n\t\t},\n\t}\n}\n\n// WithEmbedding is the option to set the embedder for the indexer, which convert document to embeddings.\nfunc WithEmbedding(emb embedding.Embedder) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.Embedding = emb\n\t\t},\n\t}\n}\n\n// Option is a call-time option for an Indexer.\ntype Option struct {\n\tapply func(opts *Options)\n\n\timplSpecificOptFn any\n}\n\n// GetCommonOptions extracts standard [Options] from opts, merging onto base.\n// Implementors must call this inside Store:\n//\n//\tfunc (idx *MyIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) ([]string, error) {\n//\t    options := indexer.GetCommonOptions(nil, opts...)\n//\t    // use options.Embedding to generate vectors before storage\n//\t}\nfunc GetCommonOptions(base *Options, opts ...Option) *Options {\n\tif base == nil {\n\t\tbase = &Options{}\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.apply != nil {\n\t\t\topt.apply(base)\n\t\t}\n\t}\n\n\treturn base\n}\n\n// WrapImplSpecificOptFn wraps an implementation-specific option function so it\n// can be passed alongside standard options. For use by Indexer implementors.\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) Option {\n\treturn Option{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetImplSpecificOptions extracts implementation-specific options from opts,\n// merging onto base. Call alongside [GetCommonOptions] inside Store.\nfunc GetImplSpecificOptions[T any](base *T, opts ...Option) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\toptFn, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\toptFn(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "components/indexer/option_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage indexer\n\nimport (\n\t\"testing\"\n\n\t\"github.com/smartystreets/goconvey/convey\"\n\n\t\"github.com/cloudwego/eino/internal/mock/components/embedding\"\n)\n\nfunc TestOptions(t *testing.T) {\n\tconvey.Convey(\"test options\", t, func() {\n\t\tvar (\n\t\t\tsubIndexes = []string{\"index_1\", \"index_2\"}\n\t\t\te          = &embedding.MockEmbedder{}\n\t\t)\n\n\t\topts := GetCommonOptions(\n\t\t\t&Options{},\n\t\t\tWithSubIndexes(subIndexes),\n\t\t\tWithEmbedding(e),\n\t\t)\n\n\t\tconvey.So(opts, convey.ShouldResemble, &Options{\n\t\t\tSubIndexes: subIndexes,\n\t\t\tEmbedding:  e,\n\t\t})\n\t})\n}\n"
  },
  {
    "path": "components/model/callback_extra.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage model\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// TokenUsage is the token usage for the model.\ntype TokenUsage struct {\n\t// PromptTokens is the number of prompt tokens, including all the input tokens of this request.\n\tPromptTokens int\n\t// PromptTokenDetails is a breakdown of the prompt tokens.\n\tPromptTokenDetails PromptTokenDetails\n\t// CompletionTokens is the number of completion tokens.\n\tCompletionTokens int\n\t// TotalTokens is the total number of tokens.\n\tTotalTokens int\n\t// CompletionTokensDetails is breakdown of completion tokens.\n\tCompletionTokensDetails CompletionTokensDetails `json:\"completion_token_details\"`\n}\n\ntype CompletionTokensDetails struct {\n\t// ReasoningTokens tokens generated by the model for reasoning.\n\t// This is currently supported by OpenAI, Gemini, ARK and Qwen  chat models.\n\t// For other models, this field will be 0.\n\tReasoningTokens int `json:\"reasoning_tokens,omitempty\"`\n}\n\n// PromptTokenDetails provides a breakdown of prompt token usage.\ntype PromptTokenDetails struct {\n\t// Cached tokens present in the prompt.\n\tCachedTokens int\n}\n\n// Config is the config for the model.\ntype Config struct {\n\t// Model is the model name.\n\tModel string\n\t// MaxTokens is the max number of tokens, if reached the max tokens, the model will stop generating, and mostly return an finish reason of \"length\".\n\tMaxTokens int\n\t// Temperature is the temperature, which controls the randomness of the model.\n\tTemperature float32\n\t// TopP is the top p, which controls the diversity of the model.\n\tTopP float32\n\t// Stop is the stop words, which controls the stopping condition of the model.\n\tStop []string\n}\n\n// CallbackInput is the input for the model callback.\ntype CallbackInput struct {\n\t// Messages is the messages to be sent to the model.\n\tMessages []*schema.Message\n\t// Tools is the tools to be used in the model.\n\tTools []*schema.ToolInfo\n\t// ToolChoice is the tool choice, which controls the tool to be used in the model.\n\tToolChoice *schema.ToolChoice\n\t// Config is the config for the model.\n\tConfig *Config\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// CallbackOutput is the output for the model callback.\ntype CallbackOutput struct {\n\t// Message is the message generated by the model.\n\tMessage *schema.Message\n\t// Config is the config for the model.\n\tConfig *Config\n\t// TokenUsage is the token usage of this request.\n\tTokenUsage *TokenUsage\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// ConvCallbackInput converts the callback input to the model callback input.\nfunc ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput {\n\tswitch t := src.(type) {\n\tcase *CallbackInput: // when callback is triggered within component implementation, the input is usually already a typed *model.CallbackInput\n\t\treturn t\n\tcase []*schema.Message: // when callback is injected by graph node, not the component implementation itself, the input is the input of Chat Model interface, which is []*schema.Message\n\t\treturn &CallbackInput{\n\t\t\tMessages: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// ConvCallbackOutput converts the callback output to the model callback output.\nfunc ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {\n\tswitch t := src.(type) {\n\tcase *CallbackOutput: // when callback is triggered within component implementation, the output is usually already a typed *model.CallbackOutput\n\t\treturn t\n\tcase *schema.Message: // when callback is injected by graph node, not the component implementation itself, the output is the output of Chat Model interface, which is *schema.Message\n\t\treturn &CallbackOutput{\n\t\t\tMessage: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n"
  },
  {
    "path": "components/model/callback_extra_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestConvModel(t *testing.T) {\n\tassert.NotNil(t, ConvCallbackInput(&CallbackInput{}))\n\tassert.NotNil(t, ConvCallbackInput([]*schema.Message{}))\n\tassert.Nil(t, ConvCallbackInput(\"asd\"))\n\n\tassert.NotNil(t, ConvCallbackOutput(&CallbackOutput{}))\n\tassert.NotNil(t, ConvCallbackOutput(&schema.Message{}))\n\tassert.Nil(t, ConvCallbackOutput(\"asd\"))\n}\n"
  },
  {
    "path": "components/model/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package model defines the ChatModel component interface for interacting with\n// large language models (LLMs).\n//\n// # Overview\n//\n// A ChatModel takes a slice of [schema.Message] as input and returns a response\n// message — either in full ([BaseChatModel.Generate]) or incrementally as a\n// stream ([BaseChatModel.Stream]). It is the most fundamental building block in\n// an eino pipeline: every application that talks to an LLM goes through this\n// interface.\n//\n// Concrete implementations (OpenAI, Ark, Ollama, …) live in eino-ext:\n//\n//\tgithub.com/cloudwego/eino-ext/components/model/\n//\n// # Interface Hierarchy\n//\n//\tBaseChatModel         — Generate + Stream (all implementations)\n//\t├── ToolCallingChatModel  — preferred; WithTools returns a new instance (concurrency-safe)\n//\t└── ChatModel             — deprecated; BindTools mutates state (avoid in new code)\n//\n// # Choosing Generate vs Stream\n//\n// Use [BaseChatModel.Generate] when the full response is needed before\n// proceeding (e.g. structured extraction, classification).\n// Use [BaseChatModel.Stream] when output should be forwarded to the caller\n// incrementally (e.g. chat UI, long-form generation). Always close the\n// [schema.StreamReader] returned by Stream — failing to do so leaks the\n// underlying connection:\n//\n//\treader, err := model.Stream(ctx, messages)\n//\tif err != nil { ... }\n//\tdefer reader.Close()\n//\n// # Implementing a ChatModel\n//\n// Implementations must call [GetCommonOptions] to extract standard options and\n// [GetImplSpecificOptions] to extract their own options from the Option list.\n// Expose implementation-specific options via [WrapImplSpecificOptFn].\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/chat_model_guide/\n// for the full component guide.\npackage model\n"
  },
  {
    "path": "components/model/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage model\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// BaseChatModel defines the core interface for all chat model implementations.\n//\n// It exposes two modes of interaction:\n//   - [BaseChatModel.Generate]: blocks until the model returns a complete response.\n//   - [BaseChatModel.Stream]: returns a [schema.StreamReader] that yields message\n//     chunks incrementally as the model generates them.\n//\n// The input is a slice of [schema.Message] representing the conversation history.\n// Messages carry a role (system, user, assistant, tool) and support multimodal\n// content (text, images, audio, video). Tool messages must include a ToolCallID\n// that correlates them with a prior assistant tool-call message.\n//\n// Stream usage — the caller is responsible for closing the reader:\n//\n//\treader, err := m.Stream(ctx, messages)\n//\tif err != nil { ... }\n//\tdefer reader.Close()\n//\tfor {\n//\t    chunk, err := reader.Recv()\n//\t    if errors.Is(err, io.EOF) { break }\n//\t    if err != nil { ... }\n//\t    // handle chunk\n//\t}\n//\n// Note: a [schema.StreamReader] can only be read once. If multiple consumers\n// need the stream, it must be copied before reading.\n//\n//go:generate  mockgen -destination ../../internal/mock/components/model/ChatModel_mock.go --package model -source interface.go\ntype BaseChatModel interface {\n\tGenerate(ctx context.Context, input []*schema.Message, opts ...Option) (*schema.Message, error)\n\tStream(ctx context.Context, input []*schema.Message, opts ...Option) (\n\t\t*schema.StreamReader[*schema.Message], error)\n}\n\n// Deprecated: Use [ToolCallingChatModel] instead.\n//\n// ChatModel extends [BaseChatModel] with tool binding via [ChatModel.BindTools].\n// BindTools mutates the instance in place, which causes a race condition when\n// the same instance is used concurrently: one goroutine's tool list can\n// overwrite another's. Prefer [ToolCallingChatModel.WithTools], which returns\n// a new immutable instance and is safe for concurrent use.\ntype ChatModel interface {\n\tBaseChatModel\n\n\t// BindTools bind tools to the model.\n\t// BindTools before requesting ChatModel generally.\n\t// notice the non-atomic problem of BindTools and Generate.\n\tBindTools(tools []*schema.ToolInfo) error\n}\n\n// ToolCallingChatModel extends [BaseChatModel] with safe tool binding.\n//\n// Unlike the deprecated [ChatModel.BindTools], [ToolCallingChatModel.WithTools]\n// does not mutate the receiver — it returns a new instance with the given tools\n// attached. This makes it safe to share a base model instance across goroutines\n// and derive per-request variants with different tool sets:\n//\n//\tbase, _ := openai.NewChatModel(ctx, cfg)           // shared, no tools\n//\twithSearch, _ := base.WithTools([]*schema.ToolInfo{searchTool})\n//\twithCalc, _  := base.WithTools([]*schema.ToolInfo{calcTool})\ntype ToolCallingChatModel interface {\n\tBaseChatModel\n\n\t// WithTools returns a new ToolCallingChatModel instance with the specified tools bound.\n\t// This method does not modify the current instance, making it safer for concurrent use.\n\tWithTools(tools []*schema.ToolInfo) (ToolCallingChatModel, error)\n}\n"
  },
  {
    "path": "components/model/option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage model\n\nimport \"github.com/cloudwego/eino/schema\"\n\n// Options is the common options for the model.\ntype Options struct {\n\t// Temperature is the temperature for the model, which controls the randomness of the model.\n\tTemperature *float32\n\t// MaxTokens is the max number of tokens, if reached the max tokens, the model will stop generating, and mostly return an finish reason of \"length\".\n\tMaxTokens *int\n\t// Model is the model name.\n\tModel *string\n\t// TopP is the top p for the model, which controls the diversity of the model.\n\tTopP *float32\n\t// Stop is the stop words for the model, which controls the stopping condition of the model.\n\tStop []string\n\t// Tools is a list of tools the model may call.\n\tTools []*schema.ToolInfo\n\t// ToolChoice controls which tool is called by the model.\n\tToolChoice *schema.ToolChoice\n\t// AllowedToolNames specifies a list of tool names that the model is allowed to call.\n\t// This allows for constraining the model to a specific subset of the available tools.\n\tAllowedToolNames []string\n}\n\n// Option is a call-time option for a ChatModel. Options are immutable and\n// composable: each Option carries either a common-option setter (applied via\n// [GetCommonOptions]) or an implementation-specific setter (applied via\n// [GetImplSpecificOptions]), never both.\ntype Option struct {\n\tapply func(opts *Options)\n\n\timplSpecificOptFn any\n}\n\n// WithTemperature is the option to set the temperature for the model.\nfunc WithTemperature(temperature float32) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.Temperature = &temperature\n\t\t},\n\t}\n}\n\n// WithMaxTokens is the option to set the max tokens for the model.\nfunc WithMaxTokens(maxTokens int) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.MaxTokens = &maxTokens\n\t\t},\n\t}\n}\n\n// WithModel is the option to set the model name.\nfunc WithModel(name string) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.Model = &name\n\t\t},\n\t}\n}\n\n// WithTopP is the option to set the top p for the model.\nfunc WithTopP(topP float32) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.TopP = &topP\n\t\t},\n\t}\n}\n\n// WithStop is the option to set the stop words for the model.\nfunc WithStop(stop []string) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.Stop = stop\n\t\t},\n\t}\n}\n\n// WithTools is the option to set tools for the model.\nfunc WithTools(tools []*schema.ToolInfo) Option {\n\tif tools == nil {\n\t\ttools = []*schema.ToolInfo{}\n\t}\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.Tools = tools\n\t\t},\n\t}\n}\n\n// WithToolChoice sets the tool choice for the model. It also allows for providing a list of\n// tool names to constrain the model to a specific subset of the available tools.\nfunc WithToolChoice(toolChoice schema.ToolChoice, allowedToolNames ...string) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.ToolChoice = &toolChoice\n\t\t\topts.AllowedToolNames = allowedToolNames\n\t\t},\n\t}\n}\n\n// WrapImplSpecificOptFn wraps an implementation-specific option function into\n// an [Option] so it can be passed alongside standard options.\n//\n// This is intended for ChatModel implementors, not callers. Define a typed\n// setter for your own config struct and expose it as an Option:\n//\n//\t// In your implementation package:\n//\tfunc WithMyParam(v string) model.Option {\n//\t    return model.WrapImplSpecificOptFn(func(o *MyOptions) {\n//\t        o.MyParam = v\n//\t    })\n//\t}\n//\n// Callers can then mix standard and implementation-specific options freely:\n//\n//\tmodel.Generate(ctx, msgs,\n//\t    model.WithTemperature(0.7),\n//\t    mypkg.WithMyParam(\"value\"),\n//\t)\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) Option {\n\treturn Option{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetCommonOptions extracts standard [Options] from an Option list, merging\n// them onto base. If base is nil, a zero-value Options is used.\n//\n// Implementors must call this to honour options passed by callers:\n//\n//\tfunc (m *MyModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n//\t    options := model.GetCommonOptions(&model.Options{Temperature: &m.defaultTemp}, opts...)\n//\t    // use options.Temperature, options.Tools, etc.\n//\t}\nfunc GetCommonOptions(base *Options, opts ...Option) *Options {\n\tif base == nil {\n\t\tbase = &Options{}\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.apply != nil {\n\t\t\topt.apply(base)\n\t\t}\n\t}\n\n\treturn base\n}\n\n// GetImplSpecificOptions extracts implementation-specific options from an\n// Option list, merging them onto base. If base is nil, a zero-value T is used.\n//\n// Call this alongside [GetCommonOptions] to support both standard and custom\n// options in your implementation:\n//\n//\ttype MyOptions struct { MyParam string }\n//\n//\tfunc (m *MyModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n//\t    common  := model.GetCommonOptions(nil, opts...)\n//\t    myOpts  := model.GetImplSpecificOptions(&MyOptions{MyParam: \"default\"}, opts...)\n//\t    // use common.Temperature, myOpts.MyParam, etc.\n//\t}\nfunc GetImplSpecificOptions[T any](base *T, opts ...Option) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\toptFn, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\toptFn(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "components/model/option_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/smartystreets/goconvey/convey\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestOptions(t *testing.T) {\n\tconvey.Convey(\"test options\", t, func() {\n\t\tvar (\n\t\t\tmodelName                  = \"model\"\n\t\t\ttemperature        float32 = 0.9\n\t\t\tmaxToken                   = 5000\n\t\t\ttopP               float32 = 0.8\n\t\t\tdefaultModel               = \"default_model\"\n\t\t\tdefaultTemperature float32 = 1.0\n\t\t\tdefaultMaxTokens           = 1000\n\t\t\tdefaultTopP        float32 = 0.5\n\t\t\ttools                      = []*schema.ToolInfo{{Name: \"asd\"}, {Name: \"qwe\"}}\n\t\t\ttoolChoice                 = schema.ToolChoiceForced\n\t\t\tallowedToolNames           = []string{\"web_search\"}\n\t\t)\n\n\t\topts := GetCommonOptions(\n\t\t\t&Options{\n\t\t\t\tModel:       &defaultModel,\n\t\t\t\tTemperature: &defaultTemperature,\n\t\t\t\tMaxTokens:   &defaultMaxTokens,\n\t\t\t\tTopP:        &defaultTopP,\n\t\t\t},\n\t\t\tWithModel(modelName),\n\t\t\tWithTemperature(temperature),\n\t\t\tWithMaxTokens(maxToken),\n\t\t\tWithTopP(topP),\n\t\t\tWithStop([]string{\"hello\", \"bye\"}),\n\t\t\tWithTools(tools),\n\t\t\tWithToolChoice(toolChoice, allowedToolNames...),\n\t\t)\n\n\t\tconvey.So(opts, convey.ShouldResemble, &Options{\n\t\t\tModel:            &modelName,\n\t\t\tTemperature:      &temperature,\n\t\t\tMaxTokens:        &maxToken,\n\t\t\tTopP:             &topP,\n\t\t\tStop:             []string{\"hello\", \"bye\"},\n\t\t\tTools:            tools,\n\t\t\tToolChoice:       &toolChoice,\n\t\t\tAllowedToolNames: allowedToolNames,\n\t\t})\n\t})\n\n\tconvey.Convey(\"test nil tools option\", t, func() {\n\t\topts := GetCommonOptions(\n\t\t\t&Options{\n\t\t\t\tTools: []*schema.ToolInfo{\n\t\t\t\t\t{Name: \"asd\"},\n\t\t\t\t\t{Name: \"qwe\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\tWithTools(nil),\n\t\t)\n\n\t\tconvey.So(opts.Tools, convey.ShouldNotBeNil)\n\t\tconvey.So(len(opts.Tools), convey.ShouldEqual, 0)\n\t})\n}\n\ntype implOption struct {\n\tuserID int64\n\tname   string\n}\n\nfunc WithUserID(uid int64) Option {\n\treturn WrapImplSpecificOptFn[implOption](func(i *implOption) {\n\t\ti.userID = uid\n\t})\n}\n\nfunc WithName(n string) Option {\n\treturn WrapImplSpecificOptFn[implOption](func(i *implOption) {\n\t\ti.name = n\n\t})\n}\n\nfunc TestImplSpecificOption(t *testing.T) {\n\tconvey.Convey(\"impl_specific_option\", t, func() {\n\t\topt := GetImplSpecificOptions(&implOption{}, WithUserID(101), WithName(\"Wang\"))\n\n\t\tconvey.So(opt, convey.ShouldEqual, &implOption{\n\t\t\tuserID: 101,\n\t\t\tname:   \"Wang\",\n\t\t})\n\t})\n}\n"
  },
  {
    "path": "components/prompt/callback_extra.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage prompt\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// CallbackInput is the input for the callback.\ntype CallbackInput struct {\n\t// Variables is the variables for the callback.\n\tVariables map[string]any\n\t// Templates is the templates for the callback.\n\tTemplates []schema.MessagesTemplate\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// CallbackOutput is the output for the callback.\ntype CallbackOutput struct {\n\t// Result is the result for the callback.\n\tResult []*schema.Message\n\t// Templates is the templates for the callback.\n\tTemplates []schema.MessagesTemplate\n\t// Extra is the extra information for the callback.\n\tExtra map[string]any\n}\n\n// ConvCallbackInput converts the callback input to the prompt callback input.\nfunc ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput {\n\tswitch t := src.(type) {\n\tcase *CallbackInput:\n\t\treturn t\n\tcase map[string]any:\n\t\treturn &CallbackInput{\n\t\t\tVariables: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// ConvCallbackOutput converts the callback output to the prompt callback output.\nfunc ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {\n\tswitch t := src.(type) {\n\tcase *CallbackOutput:\n\t\treturn t\n\tcase []*schema.Message:\n\t\treturn &CallbackOutput{\n\t\t\tResult: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n"
  },
  {
    "path": "components/prompt/callback_extra_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage prompt\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestConvPrompt(t *testing.T) {\n\tassert.NotNil(t, ConvCallbackInput(&CallbackInput{}))\n\tassert.NotNil(t, ConvCallbackInput(map[string]any{}))\n\tassert.Nil(t, ConvCallbackInput(\"asd\"))\n\n\tassert.NotNil(t, ConvCallbackOutput(&CallbackOutput{}))\n\tassert.NotNil(t, ConvCallbackOutput([]*schema.Message{}))\n\tassert.Nil(t, ConvCallbackOutput(\"asd\"))\n}\n"
  },
  {
    "path": "components/prompt/chat_template.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage prompt\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// DefaultChatTemplate is the default chat template implementation.\ntype DefaultChatTemplate struct {\n\t// templates is the templates for the chat template.\n\ttemplates []schema.MessagesTemplate\n\t// formatType is the format type for the chat template.\n\tformatType schema.FormatType\n}\n\n// FromMessages creates a new DefaultChatTemplate from the given templates and format type.\n// eg.\n//\n//\ttemplate := prompt.FromMessages(schema.FString, &schema.Message{Content: \"Hello, {name}!\"}, &schema.Message{Content: \"how are you?\"})\n//\t// in chain, or graph\n//\tchain := compose.NewChain[map[string]any, []*schema.Message]()\n//\tchain.AppendChatTemplate(template)\nfunc FromMessages(formatType schema.FormatType, templates ...schema.MessagesTemplate) *DefaultChatTemplate {\n\treturn &DefaultChatTemplate{\n\t\ttemplates:  templates,\n\t\tformatType: formatType,\n\t}\n}\n\n// Format formats the chat template with the given context and variables.\nfunc (t *DefaultChatTemplate) Format(ctx context.Context,\n\tvs map[string]any, _ ...Option) (result []*schema.Message, err error) {\n\tctx = callbacks.EnsureRunInfo(ctx, t.GetType(), components.ComponentOfPrompt)\n\tctx = callbacks.OnStart(ctx, &CallbackInput{\n\t\tVariables: vs,\n\t\tTemplates: t.templates,\n\t})\n\tdefer func() {\n\t\tif err != nil {\n\t\t\t_ = callbacks.OnError(ctx, err)\n\t\t}\n\t}()\n\n\tresult = make([]*schema.Message, 0, len(t.templates))\n\tfor _, template := range t.templates {\n\t\tmsgs, err := template.Format(ctx, vs, t.formatType)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tresult = append(result, msgs...)\n\t}\n\n\t_ = callbacks.OnEnd(ctx, &CallbackOutput{\n\t\tResult:    result,\n\t\tTemplates: t.templates,\n\t})\n\n\treturn result, nil\n}\n\n// GetType returns the type of the chat template (Default).\nfunc (t *DefaultChatTemplate) GetType() string {\n\treturn \"Default\"\n}\n\n// IsCallbacksEnabled checks if the callbacks are enabled for the chat template.\nfunc (t *DefaultChatTemplate) IsCallbacksEnabled() bool {\n\treturn true\n}\n"
  },
  {
    "path": "components/prompt/chat_template_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage prompt\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestFormat(t *testing.T) {\n\tpyFmtTestTemplate := []schema.MessagesTemplate{\n\t\tschema.SystemMessage(\n\t\t\t\"you are a helpful assistant.\\n\" +\n\t\t\t\t\"here is the context: {context}\"),\n\t\tschema.MessagesPlaceholder(\"chat_history\", true),\n\t\tschema.UserMessage(\"question: {question}\"),\n\t}\n\tjinja2TestTemplate := []schema.MessagesTemplate{\n\t\tschema.SystemMessage(\n\t\t\t\"you are a helpful assistant.\\n\" +\n\t\t\t\t\"here is the context: {{context}}\"),\n\t\tschema.MessagesPlaceholder(\"chat_history\", true),\n\t\tschema.UserMessage(\"question: {{question}}\"),\n\t}\n\tgoFmtTestTemplate := []schema.MessagesTemplate{\n\t\tschema.SystemMessage(\n\t\t\t\"you are a helpful assistant.\\n\" +\n\t\t\t\t\"here is the context: {{.context}}\"),\n\t\tschema.MessagesPlaceholder(\"chat_history\", true),\n\t\tschema.UserMessage(\"question: {{.question}}\"),\n\t}\n\ttestValues := map[string]any{\n\t\t\"context\":  \"it's beautiful day\",\n\t\t\"question\": \"how is the day today\",\n\t\t\"chat_history\": []*schema.Message{\n\t\t\tschema.UserMessage(\"who are you\"),\n\t\t\tschema.AssistantMessage(\"I'm a helpful assistant\", nil),\n\t\t},\n\t}\n\texpected := []*schema.Message{\n\t\tschema.SystemMessage(\n\t\t\t\"you are a helpful assistant.\\n\" +\n\t\t\t\t\"here is the context: it's beautiful day\"),\n\t\tschema.UserMessage(\"who are you\"),\n\t\tschema.AssistantMessage(\"I'm a helpful assistant\", nil),\n\t\tschema.UserMessage(\"question: how is the day today\"),\n\t}\n\n\t// FString\n\tchatTemplate := FromMessages(schema.FString, pyFmtTestTemplate...)\n\tmsgs, err := chatTemplate.Format(context.Background(), testValues)\n\tassert.Nil(t, err)\n\tassert.Equal(t, expected, msgs)\n\n\t// Jinja2\n\tchatTemplate = FromMessages(schema.Jinja2, jinja2TestTemplate...)\n\tmsgs, err = chatTemplate.Format(context.Background(), testValues)\n\tassert.Nil(t, err)\n\tassert.Equal(t, expected, msgs)\n\n\t// GoTemplate\n\tchatTemplate = FromMessages(schema.GoTemplate, goFmtTestTemplate...)\n\tmsgs, err = chatTemplate.Format(context.Background(), testValues)\n\tassert.Nil(t, err)\n\tassert.Equal(t, expected, msgs)\n}\n\nfunc TestDocumentFormat(t *testing.T) {\n\tdocs := []*schema.Document{\n\t\t{\n\t\t\tID:      \"1\",\n\t\t\tContent: \"qwe\",\n\t\t\tMetaData: map[string]any{\n\t\t\t\t\"hello\": 888,\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tID:      \"2\",\n\t\t\tContent: \"asd\",\n\t\t\tMetaData: map[string]any{\n\t\t\t\t\"bye\": 111,\n\t\t\t},\n\t\t},\n\t}\n\n\ttemplate := FromMessages(schema.FString,\n\t\tschema.SystemMessage(\"all:{all_docs}\\nsingle:{single_doc}\"),\n\t)\n\n\tmsgs, err := template.Format(context.Background(), map[string]any{\n\t\t\"all_docs\":   docs,\n\t\t\"single_doc\": docs[0],\n\t})\n\n\tassert.Nil(t, err)\n\tt.Log(msgs)\n}\n\nfunc TestMultiContentFormat(t *testing.T) {\n\tmtpl := []schema.MessagesTemplate{\n\t\t&schema.Message{\n\t\t\tContent: \"{a}\",\n\t\t\tMultiContent: []schema.ChatMessagePart{\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeText,\n\t\t\t\t\tText: \"{b}\",\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeImageURL,\n\t\t\t\t\tImageURL: &schema.ChatMessageImageURL{\n\t\t\t\t\t\tURL: \"{c}\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeAudioURL,\n\t\t\t\t\tAudioURL: &schema.ChatMessageAudioURL{\n\t\t\t\t\t\tURL: \"{d}\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeVideoURL,\n\t\t\t\t\tVideoURL: &schema.ChatMessageVideoURL{\n\t\t\t\t\t\tURL: \"{e}\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeFileURL,\n\t\t\t\t\tFileURL: &schema.ChatMessageFileURL{\n\t\t\t\t\t\tURL: \"{f}\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\tinput := map[string]any{\n\t\t\"a\": \"content\",\n\t\t\"b\": \"text\",\n\t\t\"c\": \"image url\",\n\t\t\"d\": \"audio url\",\n\t\t\"e\": \"video url\",\n\t\t\"f\": \"file url\",\n\t}\n\texpected := []*schema.Message{\n\t\t{\n\t\t\tContent: \"content\",\n\t\t\tMultiContent: []schema.ChatMessagePart{\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeText,\n\t\t\t\t\tText: \"text\",\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeImageURL,\n\t\t\t\t\tImageURL: &schema.ChatMessageImageURL{\n\t\t\t\t\t\tURL: \"image url\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeAudioURL,\n\t\t\t\t\tAudioURL: &schema.ChatMessageAudioURL{\n\t\t\t\t\t\tURL: \"audio url\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeVideoURL,\n\t\t\t\t\tVideoURL: &schema.ChatMessageVideoURL{\n\t\t\t\t\t\tURL: \"video url\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: schema.ChatMessagePartTypeFileURL,\n\t\t\t\t\tFileURL: &schema.ChatMessageFileURL{\n\t\t\t\t\t\tURL: \"file url\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\n\ttpl := FromMessages(schema.FString, mtpl...)\n\tresult, err := tpl.Format(context.Background(), input)\n\tassert.Nil(t, err)\n\tassert.Equal(t, expected, result)\n}\n"
  },
  {
    "path": "components/prompt/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package prompt defines the ChatTemplate component interface for building\n// structured message lists from templates and runtime variables.\n//\n// # Overview\n//\n// A ChatTemplate takes a variables map and produces a []*schema.Message slice\n// ready to pass to a [model.BaseChatModel]. It is typically the first node in\n// a pipeline, sitting before the ChatModel.\n//\n// The built-in [DefaultChatTemplate] supports three template syntaxes:\n//   - FString: {variable} substitution\n//   - GoTemplate: Go's text/template with conditionals and loops\n//   - Jinja2: Jinja2 template syntax\n//\n// # Construction\n//\n// Use [FromMessages] to build a template from a list of message templates:\n//\n//\ttmpl := prompt.FromMessages(schema.FString,\n//\t    schema.SystemMessage(\"You are a helpful assistant.\"),\n//\t    schema.UserMessage(\"Answer this: {question}\"),\n//\t)\n//\tmsgs, err := tmpl.Format(ctx, map[string]any{\"question\": \"What is eino?\"})\n//\n// Use [schema.MessagesPlaceholder] to insert a dynamic list of messages\n// (e.g. conversation history) at a fixed position in the template:\n//\n//\ttmpl := prompt.FromMessages(schema.FString,\n//\t    schema.SystemMessage(\"You are a helpful assistant.\"),\n//\t    schema.MessagesPlaceholder(\"history\", true),\n//\t    schema.UserMessage(\"{question}\"),\n//\t)\n//\n// # Common Pitfall\n//\n// Variable mismatches (a key present in the template but missing from the\n// variables map) produce a runtime error — there is no compile-time check.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/chat_template_guide/\npackage prompt\n"
  },
  {
    "path": "components/prompt/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage prompt\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nvar _ ChatTemplate = &DefaultChatTemplate{}\n\n// ChatTemplate formats a variables map into a list of messages for a ChatModel.\n//\n// Format substitutes the values from vs into the template's message list and\n// returns the resulting []*schema.Message. The exact substitution syntax\n// (FString, GoTemplate, Jinja2) is determined at construction time.\n//\n// Variable keys present in the template but absent from vs produce a runtime\n// error — there is no compile-time safety. Prefer consistent variable naming\n// across templates and callers.\n//\n// In a Graph or Chain, ChatTemplate typically precedes ChatModel. Use\n// compose.WithOutputKey to convert the prior node's output into the map[string]any\n// that Format expects.\n//\n// See [FromMessages] and [schema.MessagesPlaceholder] for construction helpers.\ntype ChatTemplate interface {\n\tFormat(ctx context.Context, vs map[string]any, opts ...Option) ([]*schema.Message, error)\n}\n"
  },
  {
    "path": "components/prompt/option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage prompt\n\n// Option is a call-time option for a ChatTemplate. The built-in\n// [DefaultChatTemplate] has no common options — this type exists primarily for\n// custom ChatTemplate implementations that need per-call configuration.\ntype Option struct {\n\timplSpecificOptFn any\n}\n\n// WrapImplSpecificOptFn wraps an implementation-specific option function so it\n// can be passed alongside any future standard options. For use by custom\n// ChatTemplate implementors.\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) Option {\n\treturn Option{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetImplSpecificOptions extracts the implementation specific options from Option list, optionally providing a base options with default values.\nfunc GetImplSpecificOptions[T any](base *T, opts ...Option) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\ts, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\ts(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "components/prompt/option_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage prompt\n\nimport (\n\t\"testing\"\n\n\t\"github.com/smartystreets/goconvey/convey\"\n)\n\ntype implOption struct {\n\tuserID int64\n\tname   string\n}\n\nfunc WithUserID(uid int64) Option {\n\treturn WrapImplSpecificOptFn[implOption](func(i *implOption) {\n\t\ti.userID = uid\n\t})\n}\n\nfunc WithName(n string) Option {\n\treturn WrapImplSpecificOptFn[implOption](func(i *implOption) {\n\t\ti.name = n\n\t})\n}\n\nfunc TestImplSpecificOption(t *testing.T) {\n\tconvey.Convey(\"impl_specific_option\", t, func() {\n\t\topt := GetImplSpecificOptions(&implOption{}, WithUserID(101), WithName(\"Wang\"))\n\n\t\tconvey.So(opt, convey.ShouldEqual, &implOption{\n\t\t\tuserID: 101,\n\t\t\tname:   \"Wang\",\n\t\t})\n\t})\n}\n"
  },
  {
    "path": "components/retriever/callback_extra.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage retriever\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// CallbackInput is the input for the retriever callback.\ntype CallbackInput struct {\n\t// Query is the query for the retriever.\n\tQuery string\n\n\t// TopK is the top k for the retriever, which means the top number of documents to retrieve.\n\tTopK int\n\t// Filter is the filter for the retriever.\n\tFilter string\n\t// ScoreThreshold is the score threshold for the retriever, eg 0.5 means the score of the document must be greater than 0.5.\n\tScoreThreshold *float64\n\n\t// Extra is the extra information for the retriever.\n\tExtra map[string]any\n}\n\n// CallbackOutput is the output for the retriever callback.\ntype CallbackOutput struct {\n\t// Docs is the documents for the retriever.\n\tDocs []*schema.Document\n\t// Extra is the extra information for the retriever.\n\tExtra map[string]any\n}\n\n// ConvCallbackInput converts the callback input to the retriever callback input.\nfunc ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput {\n\tswitch t := src.(type) {\n\tcase *CallbackInput:\n\t\treturn t\n\tcase string:\n\t\treturn &CallbackInput{\n\t\t\tQuery: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// ConvCallbackOutput converts the callback output to the retriever callback output.\nfunc ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {\n\tswitch t := src.(type) {\n\tcase *CallbackOutput:\n\t\treturn t\n\tcase []*schema.Document:\n\t\treturn &CallbackOutput{\n\t\t\tDocs: t,\n\t\t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n"
  },
  {
    "path": "components/retriever/callback_extra_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage retriever\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestConvRetriever(t *testing.T) {\n\tassert.NotNil(t, ConvCallbackInput(&CallbackInput{}))\n\tassert.NotNil(t, ConvCallbackInput(\"asd\"))\n\tassert.Nil(t, ConvCallbackInput([]string{}))\n\n\tassert.NotNil(t, ConvCallbackOutput(&CallbackOutput{}))\n\tassert.NotNil(t, ConvCallbackOutput([]*schema.Document{}))\n\tassert.Nil(t, ConvCallbackOutput(\"asd\"))\n}\n"
  },
  {
    "path": "components/retriever/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package retriever defines the Retriever component interface for fetching\n// relevant documents from a document store given a query.\n//\n// # Overview\n//\n// A Retriever is the read path of a RAG (Retrieval-Augmented Generation)\n// pipeline. Given a query string it returns the most relevant [schema.Document]\n// values from an underlying store (vector DB, keyword index, etc.).\n//\n// Concrete implementations (VikingDB, Milvus, Elasticsearch, …) live in\n// eino-ext:\n//\n//\tgithub.com/cloudwego/eino-ext/components/retriever/\n//\n// # Relationship to Indexer\n//\n// [Indexer] and Retriever are complementary:\n//   - Indexer writes documents (and their vectors) to the store\n//   - Retriever reads them back\n//\n// When both use an [embedding.Embedder], it must be the same model — vector\n// dimensions must match or similarity scores will be meaningless.\n//\n// # Result Ordering\n//\n// Results are ordered by relevance score (descending). Scores and other\n// backend metadata are available via [schema.Document].MetaData.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/retriever_guide/\npackage retriever\n"
  },
  {
    "path": "components/retriever/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage retriever\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n//go:generate mockgen -destination ../../internal/mock/components/retriever/retriever_mock.go --package retriever -source interface.go\n\n// Retriever fetches the most relevant documents from a store for a given query.\n//\n// Retrieve accepts a natural-language query string and returns matching\n// [schema.Document] values ordered by relevance (most relevant first).\n// Relevance scores and backend-specific metadata are available in\n// [schema.Document].MetaData.\n//\n// When [Options.Embedding] is set, the implementation converts the query to a\n// vector before searching. The embedder must be the same model used at index\n// time — see [indexer.Options.Embedding].\n//\n// [Options.ScoreThreshold] is a filter, not a sort: documents scoring below\n// the threshold are excluded entirely. [Options.TopK] caps the number of\n// results returned.\n//\n// Retrieve can be used standalone or added to a Graph via AddRetrieverNode:\n//\n//\tretriever, _ := redis.NewRetriever(ctx, cfg)\n//\tdocs, _ := retriever.Retrieve(ctx, \"what is eino?\", retriever.WithTopK(5))\n//\n//\tgraph.AddRetrieverNode(\"retriever\", retriever)\ntype Retriever interface {\n\tRetrieve(ctx context.Context, query string, opts ...Option) ([]*schema.Document, error)\n}\n"
  },
  {
    "path": "components/retriever/option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage retriever\n\nimport \"github.com/cloudwego/eino/components/embedding\"\n\n// Options is the options for the retriever.\ntype Options struct {\n\t// Index is the index for the retriever, index in different retriever may be different.\n\tIndex *string\n\t// SubIndex is the sub index for the retriever, sub index in different retriever may be different.\n\tSubIndex *string\n\t// TopK is the top k for the retriever, which means the top number of documents to retrieve.\n\tTopK *int\n\t// ScoreThreshold is the score threshold for the retriever, eg 0.5 means the score of the document must be greater than 0.5.\n\tScoreThreshold *float64\n\t// Embedding is the embedder for the retriever, which is used to embed the query for retrieval\t.\n\tEmbedding embedding.Embedder\n\n\t// DSLInfo carries backend-specific filter/query expressions. The structure and\n\t// semantics are defined by the underlying store implementation.\n\tDSLInfo map[string]any\n}\n\n// WithIndex wraps the index option.\nfunc WithIndex(index string) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.Index = &index\n\t\t},\n\t}\n}\n\n// WithSubIndex wraps the sub index option.\nfunc WithSubIndex(subIndex string) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.SubIndex = &subIndex\n\t\t},\n\t}\n}\n\n// WithTopK wraps the top k option.\nfunc WithTopK(topK int) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.TopK = &topK\n\t\t},\n\t}\n}\n\n// WithScoreThreshold wraps the score threshold option.\nfunc WithScoreThreshold(threshold float64) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.ScoreThreshold = &threshold\n\t\t},\n\t}\n}\n\n// WithEmbedding wraps the embedder option.\nfunc WithEmbedding(emb embedding.Embedder) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.Embedding = emb\n\t\t},\n\t}\n}\n\n// WithDSLInfo wraps the dsl info option.\nfunc WithDSLInfo(dsl map[string]any) Option {\n\treturn Option{\n\t\tapply: func(opts *Options) {\n\t\t\topts.DSLInfo = dsl\n\t\t},\n\t}\n}\n\n// Option is a call-time option for a Retriever.\ntype Option struct {\n\tapply func(opts *Options)\n\n\timplSpecificOptFn any\n}\n\n// GetCommonOptions extracts standard [Options] from opts, merging onto base.\n// Implementors must call this to honour caller-provided options:\n//\n//\tfunc (r *MyRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n//\t    options := retriever.GetCommonOptions(&retriever.Options{TopK: &r.defaultTopK}, opts...)\n//\t    // use options.TopK, options.ScoreThreshold, options.Embedding, etc.\n//\t}\nfunc GetCommonOptions(base *Options, opts ...Option) *Options {\n\tif base == nil {\n\t\tbase = &Options{}\n\t}\n\n\tfor i := range opts {\n\t\tif opts[i].apply != nil {\n\t\t\topts[i].apply(base)\n\t\t}\n\t}\n\n\treturn base\n}\n\n// WrapImplSpecificOptFn wraps an implementation-specific option function so it\n// can be passed alongside standard options. For use by Retriever implementors.\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) Option {\n\treturn Option{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetImplSpecificOptions extracts implementation-specific options from opts,\n// merging onto base. Call alongside [GetCommonOptions] inside Retrieve.\nfunc GetImplSpecificOptions[T any](base *T, opts ...Option) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\toptFn, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\toptFn(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "components/retriever/option_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage retriever\n\nimport (\n\t\"testing\"\n\n\t\"github.com/smartystreets/goconvey/convey\"\n\n\t\"github.com/cloudwego/eino/internal/mock/components/embedding\"\n)\n\nfunc TestOptions(t *testing.T) {\n\tconvey.Convey(\"test options\", t, func() {\n\t\tvar (\n\t\t\tindex          = \"index\"\n\t\t\ttopK           = 2\n\t\t\tscoreThreshold = 4.0\n\t\t\tsubIndex       = \"sub_index\"\n\t\t\tdslInfo        = map[string]any{\"dsl\": \"dsl\"}\n\t\t\te              = &embedding.MockEmbedder{}\n\t\t\tdefaultTopK    = 1\n\t\t)\n\n\t\topts := GetCommonOptions(\n\t\t\t&Options{\n\t\t\t\tTopK: &defaultTopK,\n\t\t\t},\n\t\t\tWithIndex(index),\n\t\t\tWithTopK(topK),\n\t\t\tWithScoreThreshold(scoreThreshold),\n\t\t\tWithSubIndex(subIndex),\n\t\t\tWithDSLInfo(dslInfo),\n\t\t\tWithEmbedding(e),\n\t\t)\n\n\t\tconvey.So(opts, convey.ShouldResemble, &Options{\n\t\t\tIndex:          &index,\n\t\t\tTopK:           &topK,\n\t\t\tScoreThreshold: &scoreThreshold,\n\t\t\tSubIndex:       &subIndex,\n\t\t\tDSLInfo:        dslInfo,\n\t\t\tEmbedding:      e,\n\t\t})\n\t})\n}\n"
  },
  {
    "path": "components/tool/callback_extra.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage tool\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// CallbackInput is the input for the tool callback.\ntype CallbackInput struct {\n\t// ArgumentsInJSON is the arguments in json format for the tool.\n\tArgumentsInJSON string\n\t// Extra is the extra information for the tool.\n\tExtra map[string]any\n}\n\n// CallbackOutput is the output for the tool callback.\ntype CallbackOutput struct {\n\t// Response is the response for the tool.\n\tResponse string\n\t// ToolOutput is the multimodal output for the tool. Used when the tool returns structured data.\n\tToolOutput *schema.ToolResult\n\t// Extra is the extra information for the tool.\n\tExtra map[string]any\n}\n\n// ConvCallbackInput converts the callback input to the tool callback input.\nfunc ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput {\n\tswitch t := src.(type) {\n\tcase *CallbackInput:\n\t\treturn t\n\tcase string:\n\t\treturn &CallbackInput{ArgumentsInJSON: t}\n\tcase *schema.ToolArgument:\n\t\treturn &CallbackInput{ArgumentsInJSON: t.Text}\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// ConvCallbackOutput converts the callback output to the tool callback output.\nfunc ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {\n\tswitch t := src.(type) {\n\tcase *CallbackOutput:\n\t\treturn t\n\tcase string:\n\t\treturn &CallbackOutput{Response: t}\n\tcase *schema.ToolResult:\n\t\treturn &CallbackOutput{ToolOutput: t}\n\tdefault:\n\t\treturn nil\n\t}\n}\n"
  },
  {
    "path": "components/tool/callback_extra_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage tool\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestConvCallbackInput(t *testing.T) {\n\tassert.NotNil(t, ConvCallbackInput(&CallbackInput{}))\n\tassert.NotNil(t, ConvCallbackInput(\"asd\"))\n\tassert.Nil(t, ConvCallbackInput(123))\n\tassert.Nil(t, ConvCallbackInput(nil))\n}\n\nfunc TestConvCallbackOutput(t *testing.T) {\n\tassert.NotNil(t, ConvCallbackOutput(&CallbackOutput{}))\n\tassert.NotNil(t, ConvCallbackOutput(\"asd\"))\n\tassert.Nil(t, ConvCallbackOutput(123))\n\tassert.Nil(t, ConvCallbackOutput(nil))\n}\n"
  },
  {
    "path": "components/tool/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package tool defines the tool component interfaces that allow language models\n// to invoke external capabilities, and helpers for interrupt/resume within tools.\n//\n// # Interface Hierarchy\n//\n//\tBaseTool                  — Info() only; for passing tool metadata to a ChatModel\n//\t├── InvokableTool         — standard: args as JSON string, returns string\n//\t├── StreamableTool        — standard streaming: args as JSON string, returns StreamReader[string]\n//\t├── EnhancedInvokableTool — multimodal: args as *schema.ToolArgument, returns *schema.ToolResult\n//\t└── EnhancedStreamableTool— multimodal streaming\n//\n// # Choosing an Interface\n//\n// Implement [InvokableTool] for most tools — arguments arrive as a JSON string\n// automatically decoded from the model's tool call, and the result is a string\n// sent back to the model.\n//\n// Implement [EnhancedInvokableTool] when the tool needs to return structured\n// multimodal content (images, audio, files) rather than plain text. When a\n// tool implements both a standard and an enhanced interface, ToolsNode\n// prioritises the enhanced interface.\n//\n// # Creating Tools\n//\n// The [utils] sub-package provides constructors that eliminate boilerplate:\n//   - [utils.InferTool] / [utils.InferStreamTool] — infer parameter schema from Go struct tags\n//   - [utils.NewTool] / [utils.NewStreamTool] — manual ToolInfo + typed function\n//\n// # Interrupt / Resume\n//\n// Tools can pause execution and wait for external input using [Interrupt],\n// [StatefulInterrupt], and [CompositeInterrupt]. Use [GetInterruptState] and\n// [GetResumeContext] inside the tool to distinguish first-run from resumed-run.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/tools_node_guide/\n// See https://www.cloudwego.io/docs/eino/core_modules/components/tools_node_guide/how_to_create_a_tool/\npackage tool\n"
  },
  {
    "path": "components/tool/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage tool\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// BaseTool provides the metadata that a ChatModel uses to decide whether and\n// how to call a tool. Info returns a [schema.ToolInfo] containing the tool\n// name, description, and parameter JSON schema.\n//\n// BaseTool alone is sufficient when passing tool definitions to a ChatModel\n// via WithTools — the model only needs the schema to generate tool calls.\n// To also execute the tool, implement [InvokableTool] or [StreamableTool].\ntype BaseTool interface {\n\tInfo(ctx context.Context) (*schema.ToolInfo, error)\n}\n\n// InvokableTool is a tool that can be executed by ToolsNode.\n//\n// InvokableRun receives the model's tool call arguments as a JSON-encoded\n// string and returns a plain string result that is sent back to the model as\n// a tool message. The framework handles JSON decoding automatically when using\n// the [utils.InferTool] or [utils.NewTool] constructors.\ntype InvokableTool interface {\n\tBaseTool\n\n\t// InvokableRun executes the tool with arguments encoded as a JSON string.\n\tInvokableRun(ctx context.Context, argumentsInJSON string, opts ...Option) (string, error)\n}\n\n// StreamableTool is a streaming variant of [InvokableTool].\n//\n// StreamableRun returns a [schema.StreamReader] that yields string chunks\n// incrementally. The caller (ToolsNode) is responsible for closing the reader.\ntype StreamableTool interface {\n\tBaseTool\n\n\tStreamableRun(ctx context.Context, argumentsInJSON string, opts ...Option) (*schema.StreamReader[string], error)\n}\n\n// EnhancedInvokableTool is a tool that returns structured multimodal results.\n//\n// Unlike [InvokableTool], arguments arrive as a [schema.ToolArgument] (not a\n// raw JSON string) and the result is a [schema.ToolResult] which can carry\n// text, images, audio, video, and file content.\n//\n// When a tool implements both a standard and an enhanced interface, ToolsNode\n// prioritises the enhanced interface.\ntype EnhancedInvokableTool interface {\n\tBaseTool\n\tInvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...Option) (*schema.ToolResult, error)\n}\n\n// EnhancedStreamableTool is the streaming variant of [EnhancedInvokableTool].\n//\n// It streams [schema.ToolResult] chunks, enabling incremental multimodal\n// output. The caller is responsible for closing the returned [schema.StreamReader].\ntype EnhancedStreamableTool interface {\n\tBaseTool\n\tStreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...Option) (*schema.StreamReader[*schema.ToolResult], error)\n}\n"
  },
  {
    "path": "components/tool/interrupt.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage tool\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/internal/core\"\n)\n\n// Interrupt pauses tool execution and signals the orchestration layer to checkpoint.\n// The tool can be resumed later with optional data.\n//\n// Parameters:\n//   - ctx: The context passed to InvokableRun/StreamableRun\n//   - info: User-facing information about why the tool is interrupting (e.g., \"needs user confirmation\")\n//\n// Returns an error that should be returned from InvokableRun/StreamableRun.\n//\n// Example:\n//\n//\tfunc (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) {\n//\t    if needsConfirmation(args) {\n//\t        return \"\", tool.Interrupt(ctx, \"Please confirm this action\")\n//\t    }\n//\t    return doWork(args), nil\n//\t}\nfunc Interrupt(ctx context.Context, info any) error {\n\tis, err := core.Interrupt(ctx, info, nil, nil)\n\tif err != nil {\n\t\treturn err\n\t}\n\treturn is\n}\n\n// StatefulInterrupt pauses tool execution with state preservation.\n// Use this when the tool has internal state that must be restored on resume.\n//\n// Parameters:\n//   - ctx: The context passed to InvokableRun/StreamableRun\n//   - info: User-facing information about the interrupt\n//   - state: Internal state to persist (must be gob-serializable)\n//\n// Example:\n//\n//\tfunc (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) {\n//\t    wasInterrupted, hasState, state := tool.GetInterruptState[MyState](ctx)\n//\t    if !wasInterrupted {\n//\t        // First run - interrupt with state\n//\t        return \"\", tool.StatefulInterrupt(ctx, \"processing\", MyState{Step: 1})\n//\t    }\n//\t    // Resumed - continue from saved state\n//\t    return continueFrom(state), nil\n//\t}\nfunc StatefulInterrupt(ctx context.Context, info any, state any) error {\n\tis, err := core.Interrupt(ctx, info, state, nil)\n\tif err != nil {\n\t\treturn err\n\t}\n\treturn is\n}\n\n// CompositeInterrupt creates an interrupt that aggregates multiple sub-interrupts.\n// Use this when a tool internally executes a graph or other interruptible components.\n//\n// Parameters:\n//   - ctx: The context passed to InvokableRun/StreamableRun\n//   - info: User-facing information for this tool's interrupt\n//   - state: Internal state to persist for this tool\n//   - errs: Interrupt errors from sub-components (graphs, other tools, etc.)\n//\n// Example:\n//\n//\tfunc (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) {\n//\t    result, err := t.internalGraph.Invoke(ctx, input)\n//\t    if err != nil {\n//\t        if _, ok := tool.IsInterruptError(err); ok {\n//\t            return \"\", tool.CompositeInterrupt(ctx, \"graph interrupted\", myState, err)\n//\t        }\n//\t        return \"\", err\n//\t    }\n//\t    return result, nil\n//\t}\nfunc CompositeInterrupt(ctx context.Context, info any, state any, errs ...error) error {\n\tif len(errs) == 0 {\n\t\treturn StatefulInterrupt(ctx, info, state)\n\t}\n\n\tvar cErrs []*core.InterruptSignal\n\tfor _, err := range errs {\n\t\tire := &core.InterruptSignal{}\n\t\tif errors.As(err, &ire) {\n\t\t\tcErrs = append(cErrs, ire)\n\t\t\tcontinue\n\t\t}\n\n\t\tvar provider core.InterruptContextsProvider\n\t\tif errors.As(err, &provider) {\n\t\t\tis := core.FromInterruptContexts(provider.GetInterruptContexts())\n\t\t\tif is != nil {\n\t\t\t\tcErrs = append(cErrs, is)\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\n\t\treturn fmt.Errorf(\"composite interrupt but one of the sub error is not interrupt error: %w\", err)\n\t}\n\n\tis, err := core.Interrupt(ctx, info, state, cErrs)\n\tif err != nil {\n\t\treturn err\n\t}\n\treturn is\n}\n\n// GetInterruptState checks if the tool was previously interrupted and retrieves saved state.\n//\n// Returns:\n//   - wasInterrupted: true if this tool was part of a previous interruption\n//   - hasState: true if state was saved and successfully cast to type T\n//   - state: the saved state (zero value if hasState is false)\n//\n// Example:\n//\n//\tfunc (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) {\n//\t    wasInterrupted, hasState, state := tool.GetInterruptState[MyState](ctx)\n//\t    if wasInterrupted && hasState {\n//\t        // Continue from saved state\n//\t        return continueFrom(state), nil\n//\t    }\n//\t    // First run\n//\t    return \"\", tool.StatefulInterrupt(ctx, \"need input\", MyState{Step: 1})\n//\t}\nfunc GetInterruptState[T any](ctx context.Context) (wasInterrupted bool, hasState bool, state T) {\n\treturn core.GetInterruptState[T](ctx)\n}\n\n// GetResumeContext checks if this tool is the explicit target of a resume operation.\n//\n// Returns:\n//   - isResumeTarget: true if this tool was explicitly targeted for resume\n//   - hasData: true if resume data was provided\n//   - data: the resume data (zero value if hasData is false)\n//\n// Use this to differentiate between:\n//   - Being resumed as the target (should proceed with work)\n//   - Being re-executed because a sibling was resumed (should re-interrupt)\n//\n// Example:\n//\n//\tfunc (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) {\n//\t    wasInterrupted, _, _ := tool.GetInterruptState[any](ctx)\n//\t    if !wasInterrupted {\n//\t        return \"\", tool.Interrupt(ctx, \"need confirmation\")\n//\t    }\n//\n//\t    isTarget, hasData, data := tool.GetResumeContext[string](ctx)\n//\t    if !isTarget {\n//\t        // Not our turn - re-interrupt\n//\t        return \"\", tool.Interrupt(ctx, nil)\n//\t    }\n//\t    if hasData {\n//\t        return data, nil\n//\t    }\n//\t    return \"default result\", nil\n//\t}\nfunc GetResumeContext[T any](ctx context.Context) (isResumeTarget bool, hasData bool, data T) {\n\treturn core.GetResumeContext[T](ctx)\n}\n"
  },
  {
    "path": "components/tool/interrupt_test.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage tool\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/internal/core\"\n)\n\nfunc TestInterrupt(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"basic interrupt\", func(t *testing.T) {\n\t\terr := Interrupt(ctx, \"test info\")\n\t\tassert.Error(t, err)\n\n\t\tvar signal *core.InterruptSignal\n\t\tassert.True(t, errors.As(err, &signal))\n\t\tassert.Equal(t, \"test info\", signal.Info)\n\t\tassert.True(t, signal.IsRootCause)\n\t})\n}\n\nfunc TestStatefulInterrupt(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"stateful interrupt\", func(t *testing.T) {\n\t\ttype myState struct {\n\t\t\tValue int\n\t\t}\n\t\tstate := &myState{Value: 42}\n\n\t\terr := StatefulInterrupt(ctx, \"test info\", state)\n\t\tassert.Error(t, err)\n\n\t\tvar signal *core.InterruptSignal\n\t\tassert.True(t, errors.As(err, &signal))\n\t\tassert.Equal(t, \"test info\", signal.Info)\n\t\tassert.Equal(t, state, signal.State)\n\t\tassert.True(t, signal.IsRootCause)\n\t})\n}\n\nfunc TestCompositeInterrupt(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"no sub errors falls back to StatefulInterrupt\", func(t *testing.T) {\n\t\terr := CompositeInterrupt(ctx, \"composite info\", \"my state\")\n\t\tassert.Error(t, err)\n\n\t\tvar signal *core.InterruptSignal\n\t\tassert.True(t, errors.As(err, &signal))\n\t\tassert.Equal(t, \"composite info\", signal.Info)\n\t\tassert.Equal(t, \"my state\", signal.State)\n\t\tassert.True(t, signal.IsRootCause)\n\t\tassert.Empty(t, signal.Subs)\n\t})\n\n\tt.Run(\"with InterruptSignal sub error\", func(t *testing.T) {\n\t\tsubSignal, _ := core.Interrupt(ctx, \"sub info\", \"sub state\", nil)\n\n\t\terr := CompositeInterrupt(ctx, \"composite info\", \"my state\", subSignal)\n\t\tassert.Error(t, err)\n\n\t\tvar signal *core.InterruptSignal\n\t\tassert.True(t, errors.As(err, &signal))\n\t\tassert.Equal(t, \"composite info\", signal.Info)\n\t\tassert.Equal(t, \"my state\", signal.State)\n\t\tassert.Len(t, signal.Subs, 1)\n\t\tassert.Equal(t, \"sub info\", signal.Subs[0].Info)\n\t})\n\n\tt.Run(\"with non-interrupt error returns error\", func(t *testing.T) {\n\t\tnonInterruptErr := errors.New(\"regular error\")\n\n\t\terr := CompositeInterrupt(ctx, \"composite info\", \"my state\", nonInterruptErr)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"composite interrupt but one of the sub error is not interrupt error\")\n\n\t\tvar signal *core.InterruptSignal\n\t\tassert.False(t, errors.As(err, &signal))\n\t})\n\n\tt.Run(\"with multiple sub errors\", func(t *testing.T) {\n\t\tsubSignal1, _ := core.Interrupt(ctx, \"sub1 info\", nil, nil)\n\t\tsubSignal2, _ := core.Interrupt(ctx, \"sub2 info\", nil, nil)\n\n\t\terr := CompositeInterrupt(ctx, \"composite info\", nil, subSignal1, subSignal2)\n\t\tassert.Error(t, err)\n\n\t\tvar signal *core.InterruptSignal\n\t\tassert.True(t, errors.As(err, &signal))\n\t\tassert.Len(t, signal.Subs, 2)\n\t})\n}\n\nfunc TestGetInterruptState(t *testing.T) {\n\tt.Run(\"not interrupted returns false\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\twasInterrupted, hasState, state := GetInterruptState[string](ctx)\n\t\tassert.False(t, wasInterrupted)\n\t\tassert.False(t, hasState)\n\t\tassert.Empty(t, state)\n\t})\n}\n\nfunc TestGetResumeContext(t *testing.T) {\n\tt.Run(\"not resume target returns false\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tisResumeTarget, hasData, data := GetResumeContext[string](ctx)\n\t\tassert.False(t, isResumeTarget)\n\t\tassert.False(t, hasData)\n\t\tassert.Empty(t, data)\n\t})\n}\n"
  },
  {
    "path": "components/tool/option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage tool\n\n// Option defines call option for InvokableTool or StreamableTool component, which is part of component interface signature.\n// Each tool implementation could define its own options struct and option funcs within its own package,\n// then wrap the impl specific option funcs into this type, before passing to InvokableRun or StreamableRun.\ntype Option struct {\n\timplSpecificOptFn any\n}\n\n// WrapImplSpecificOptFn wraps the impl specific option functions into Option type.\n// T: the type of the impl specific options struct.\n// Tool implementations are required to use this function to convert its own option functions into the unified Option type.\n// For example, if the tool defines its own options struct:\n//\n//\ttype customOptions struct {\n//\t    conf string\n//\t}\n//\n// Then the tool needs to provide an option function as such:\n//\n//\tfunc WithConf(conf string) Option {\n//\t    return WrapImplSpecificOptFn(func(o *customOptions) {\n//\t\t\to.conf = conf\n//\t\t}\n//\t}\n//\n// .\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) Option {\n\treturn Option{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetImplSpecificOptions provides tool author the ability to extract their own custom options from the unified Option type.\n// T: the type of the impl specific options struct.\n// This function should be used within the tool implementation's InvokableRun or StreamableRun functions.\n// It is recommended to provide a base T as the first argument, within which the tool author can provide default values for the impl specific options.\n// eg.\n//\n//\ttype customOptions struct {\n//\t    conf string\n//\t}\n//\tdefaultOptions := &customOptions{}\n//\n//\tcustomOptions := tool.GetImplSpecificOptions(defaultOptions, opts...)\nfunc GetImplSpecificOptions[T any](base *T, opts ...Option) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\toptFn, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\toptFn(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "components/tool/option_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage tool\n\nimport (\n\t\"testing\"\n\n\t\"github.com/smartystreets/goconvey/convey\"\n)\n\nfunc TestImplSpecificOpts(t *testing.T) {\n\tconvey.Convey(\"TestImplSpecificOpts\", t, func() {\n\t\ttype implSpecificOptions struct {\n\t\t\tconf  string\n\t\t\tindex int\n\t\t}\n\n\t\twithConf := func(conf string) func(o *implSpecificOptions) {\n\t\t\treturn func(o *implSpecificOptions) {\n\t\t\t\to.conf = conf\n\t\t\t}\n\t\t}\n\n\t\twithIndex := func(index int) func(o *implSpecificOptions) {\n\t\t\treturn func(o *implSpecificOptions) {\n\t\t\t\to.index = index\n\t\t\t}\n\t\t}\n\n\t\ttoolOption1 := WrapImplSpecificOptFn(withConf(\"test_conf\"))\n\t\ttoolOption2 := WrapImplSpecificOptFn(withIndex(1))\n\n\t\timplSpecificOpts := GetImplSpecificOptions(&implSpecificOptions{}, toolOption1, toolOption2)\n\n\t\tconvey.So(implSpecificOpts, convey.ShouldResemble, &implSpecificOptions{\n\t\t\tconf:  \"test_conf\",\n\t\t\tindex: 1,\n\t\t})\n\t})\n}\n"
  },
  {
    "path": "components/tool/utils/common.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"github.com/bytedance/sonic\"\n)\n\nfunc marshalString(resp any) (string, error) {\n\tif rs, ok := resp.(string); ok {\n\t\treturn rs, nil\n\t}\n\treturn sonic.MarshalString(resp)\n}\n"
  },
  {
    "path": "components/tool/utils/common_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"fmt\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestMarshalString(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\tinput    interface{}\n\t\texpected string\n\t\thasError bool\n\t}{\n\t\t{\n\t\t\tname:     \"string input should return as-is\",\n\t\t\tinput:    \"hello world\",\n\t\t\texpected: \"hello world\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"empty string should return empty\",\n\t\t\tinput:    \"\",\n\t\t\texpected: \"\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"string with special characters\",\n\t\t\tinput:    \"hello\\nworld\\t\\\"test\\\"\",\n\t\t\texpected: \"hello\\nworld\\t\\\"test\\\"\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"string with unicode\",\n\t\t\tinput:    \"你好世界\",\n\t\t\texpected: \"你好世界\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"integer should be marshaled to JSON\",\n\t\t\tinput:    42,\n\t\t\texpected: \"42\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"float should be marshaled to JSON\",\n\t\t\tinput:    3.14,\n\t\t\texpected: \"3.14\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"boolean true should be marshaled to JSON\",\n\t\t\tinput:    true,\n\t\t\texpected: \"true\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"boolean false should be marshaled to JSON\",\n\t\t\tinput:    false,\n\t\t\texpected: \"false\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"nil should be marshaled to JSON null\",\n\t\t\tinput:    nil,\n\t\t\texpected: \"null\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"slice should be marshaled to JSON array\",\n\t\t\tinput:    []int{1, 2, 3},\n\t\t\texpected: \"[1,2,3]\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"empty slice should be marshaled to JSON empty array\",\n\t\t\tinput:    []int{},\n\t\t\texpected: \"[]\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"empty map should be marshaled to JSON empty object\",\n\t\t\tinput:    map[string]int{},\n\t\t\texpected: \"{}\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"struct should be marshaled to JSON\",\n\t\t\tinput:    struct{ Name string }{Name: \"test\"},\n\t\t\texpected: `{\"Name\":\"test\"}`,\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"pointer to string should be handled as non-string\",\n\t\t\tinput:    func() *string { s := \"test\"; return &s }(),\n\t\t\texpected: `\"test\"`,\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"interface{} containing string should return as-is\",\n\t\t\tinput:    interface{}(\"test string\"),\n\t\t\texpected: \"test string\",\n\t\t\thasError: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"interface{} containing int should be marshaled\",\n\t\t\tinput:    interface{}(123),\n\t\t\texpected: \"123\",\n\t\t\thasError: false,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := marshalString(tt.input)\n\n\t\t\tif tt.hasError {\n\t\t\t\tassert.Error(t, err)\n\t\t\t} else {\n\t\t\t\tassert.NoError(t, err)\n\t\t\t\tassert.Equal(t, tt.expected, result)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestMarshalStringEdgeCases(t *testing.T) {\n\tt.Run(\"complex nested structure\", func(t *testing.T) {\n\t\tcomplex := map[string]interface{}{\n\t\t\t\"string\": \"value\",\n\t\t\t\"number\": 42,\n\t\t\t\"nested\": map[string]interface{}{\n\t\t\t\t\"array\": []string{\"a\", \"b\", \"c\"},\n\t\t\t\t\"bool\":  true,\n\t\t\t},\n\t\t}\n\n\t\tresult, err := marshalString(complex)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, result, `\"string\":\"value\"`)\n\t\tassert.Contains(t, result, `\"number\":42`)\n\t\tassert.Contains(t, result, `\"nested\"`)\n\t})\n\n\tt.Run(\"string type assertion priority\", func(t *testing.T) {\n\t\t// Test that string type assertion has priority over JSON marshaling\n\t\tvar input interface{} = \"direct string\"\n\t\tresult, err := marshalString(input)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"direct string\", result)\n\n\t\t// Verify it's not JSON encoded\n\t\tassert.NotEqual(t, `\"direct string\"`, result)\n\t})\n}\n\nfunc TestMarshalStringConsistency(t *testing.T) {\n\tt.Run(\"string vs JSON marshaling difference\", func(t *testing.T) {\n\t\tinput := `{\"key\": \"value\"}`\n\n\t\t// Direct string should return as-is\n\t\tresult, err := marshalString(input)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, input, result)\n\n\t\t// Should not be double-encoded\n\t\tassert.NotEqual(t, fmt.Sprintf(`\"%s\"`, input), result)\n\t})\n}\n"
  },
  {
    "path": "components/tool/utils/create_options.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"context\"\n\t\"reflect\"\n\n\t\"github.com/eino-contrib/jsonschema\"\n)\n\n// UnmarshalArguments is the function type for unmarshalling the arguments.\ntype UnmarshalArguments func(ctx context.Context, arguments string) (any, error)\n\n// MarshalOutput is the function type for marshalling the output.\ntype MarshalOutput func(ctx context.Context, output any) (string, error)\n\ntype toolOptions struct {\n\tum         UnmarshalArguments\n\tm          MarshalOutput\n\tscModifier SchemaModifierFn\n}\n\n// Option is the option func for the tool.\ntype Option func(o *toolOptions)\n\n// WithUnmarshalArguments wraps the unmarshal arguments option.\n// when you want to unmarshal the arguments by yourself, you can use this option.\nfunc WithUnmarshalArguments(um UnmarshalArguments) Option {\n\treturn func(o *toolOptions) {\n\t\to.um = um\n\t}\n}\n\n// WithMarshalOutput wraps the marshal output option.\n// when you want to marshal the output by yourself, you can use this option.\nfunc WithMarshalOutput(m MarshalOutput) Option {\n\treturn func(o *toolOptions) {\n\t\to.m = m\n\t}\n}\n\n// SchemaModifierFn is the schema modifier function for inferring tool parameter from tagged go struct.\n// Within this function, end-user can parse custom go struct tags into corresponding json schema field.\n// Parameters:\n// 1. jsonTagName: the name defined in the json tag. Specifically, the last 'jsonTagName' visited is fixed to be '_root', which represents the entire go struct. Also, for array field, both the field itself and the element within the array will trigger this function.\n// 2. t: the type of current schema, usually the field type of the go struct.\n// 3. tag: the struct tag of current schema, usually the field tag of the go struct. Note that the element within an array field will use the same go struct tag as the array field itself.\n// 4. schema: the current json schema object to be modified.\ntype SchemaModifierFn func(jsonTagName string, t reflect.Type, tag reflect.StructTag, schema *jsonschema.Schema)\n\n// WithSchemaModifier sets a user-defined schema modifier for inferring tool parameter from tagged go struct.\nfunc WithSchemaModifier(modifier SchemaModifierFn) Option {\n\treturn func(o *toolOptions) {\n\t\to.scModifier = modifier\n\t}\n}\n\nfunc getToolOptions(opt ...Option) *toolOptions {\n\topts := &toolOptions{\n\t\tum: nil,\n\t\tm:  nil,\n\t}\n\tfor _, o := range opt {\n\t\to(opts)\n\t}\n\treturn opts\n}\n"
  },
  {
    "path": "components/tool/utils/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package utils provides constructors for building tool implementations without\n// writing boilerplate JSON serialization code.\n//\n// # Choosing a Constructor\n//\n// There are two main strategies:\n//\n//  1. Infer from struct tags (recommended): [InferTool], [InferStreamTool],\n//     [InferEnhancedTool], [InferEnhancedStreamTool].\n//     The parameter JSON schema is derived automatically from the input struct's\n//     field names and tags. Requires a typed input struct.\n//\n//  2. Manual ToolInfo: [NewTool], [NewStreamTool], [NewEnhancedTool],\n//     [NewEnhancedStreamTool].\n//     You supply a [schema.ToolInfo] directly. Useful when the schema cannot\n//     be expressed as a Go struct, or must be dynamically constructed.\n//\n// # Struct Tag Convention\n//\n// InferTool and friends use the following tags on the input struct fields:\n//\n//\ttype Input struct {\n//\t    Query    string `json:\"query\"     jsonschema:\"required\"             jsonschema_description:\"The search query\"`\n//\t    MaxItems int    `json:\"max_items\"                                   jsonschema_description:\"Maximum results to return\"`\n//\t}\n//\n// Key rules:\n//   - Use a separate jsonschema_description tag for field descriptions —\n//     embedding descriptions inside the jsonschema tag causes comma-parsing\n//     issues.\n//   - Use jsonschema:\"required\" to mark mandatory parameters.\n//   - The json tag controls the parameter name visible to the model.\n//\n// # Schema Utilities\n//\n// [GoStruct2ToolInfo] and [GoStruct2ParamsOneOf] convert a Go struct to schema\n// types without creating a tool — useful for ChatModel structured output via\n// ResponseFormat or BindTools.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/components/tools_node_guide/how_to_create_a_tool/\npackage utils\n"
  },
  {
    "path": "components/tool/utils/error_handler.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// ErrorHandler converts a tool error into a string response.\ntype ErrorHandler func(context.Context, error) string\n\n// WrapToolWithErrorHandler wraps any BaseTool with custom error handling.\n// This function detects the tool type (InvokableTool, StreamableTool, or both)\n// and applies the appropriate error handling wrapper.\n// When the wrapped tool returns an error, the error handler function 'h' will be called\n// to convert the error into a string result, and no error will be returned from the wrapper.\n//\n// Parameters:\n//   - t: The original BaseTool to be wrapped\n//   - h: A function that converts an error to a string\n//\n// Returns:\n//   - A wrapped BaseTool that handles errors internally based on its capabilities\nfunc WrapToolWithErrorHandler(t tool.BaseTool, h ErrorHandler) tool.BaseTool {\n\tih := &infoHelper{info: t.Info}\n\tvar s tool.StreamableTool\n\tif st, ok := t.(tool.StreamableTool); ok {\n\t\ts = st\n\t}\n\tif it, ok := t.(tool.InvokableTool); ok {\n\t\tif s == nil {\n\t\t\treturn WrapInvokableToolWithErrorHandler(it, h)\n\t\t} else {\n\t\t\treturn &combinedErrorWrapper{\n\t\t\t\tinfoHelper: ih,\n\t\t\t\terrorHelper: &errorHelper{\n\t\t\t\t\ti: it.InvokableRun,\n\t\t\t\t\th: h,\n\t\t\t\t},\n\t\t\t\tstreamErrorHelper: &streamErrorHelper{\n\t\t\t\t\ts: s.StreamableRun,\n\t\t\t\t\th: h,\n\t\t\t\t},\n\t\t\t}\n\t\t}\n\t}\n\tif s != nil {\n\t\treturn WrapStreamableToolWithErrorHandler(s, h)\n\t}\n\treturn t\n}\n\n// WrapInvokableToolWithErrorHandler wraps an InvokableTool with custom error handling.\n// When the wrapped tool returns an error, the error handler function 'h' will be called\n// to convert the error into a string result, and no error will be returned from the wrapper.\n//\n// Parameters:\n//   - tool: The original InvokableTool to be wrapped\n//   - h: A function that converts an error to a string\n//\n// Returns:\n//   - A wrapped InvokableTool that handles errors internally\nfunc WrapInvokableToolWithErrorHandler(t tool.InvokableTool, h ErrorHandler) tool.InvokableTool {\n\treturn &errorWrapper{\n\t\tinfoHelper: &infoHelper{info: t.Info},\n\t\terrorHelper: &errorHelper{\n\t\t\ti: t.InvokableRun,\n\t\t\th: h,\n\t\t},\n\t}\n}\n\n// WrapStreamableToolWithErrorHandler wraps a StreamableTool with custom error handling.\n// When the wrapped tool returns an error, the error handler function 'h' will be called\n// to convert the error into a string result, which will be returned as a single-item stream,\n// and no error will be returned from the wrapper.\n//\n// Parameters:\n//   - tool: The original StreamableTool to be wrapped\n//   - h: A function that converts an error to a string\n//\n// Returns:\n//   - A wrapped StreamableTool that handles errors internally\nfunc WrapStreamableToolWithErrorHandler(t tool.StreamableTool, h ErrorHandler) tool.StreamableTool {\n\treturn &streamErrorWrapper{\n\t\tinfoHelper: &infoHelper{info: t.Info},\n\t\tstreamErrorHelper: &streamErrorHelper{\n\t\t\ts: t.StreamableRun,\n\t\t\th: h,\n\t\t},\n\t}\n}\n\ntype errorWrapper struct {\n\t*infoHelper\n\t*errorHelper\n}\n\ntype streamErrorWrapper struct {\n\t*infoHelper\n\t*streamErrorHelper\n}\n\ntype combinedErrorWrapper struct {\n\t*infoHelper\n\t*errorHelper\n\t*streamErrorHelper\n}\n\ntype infoHelper struct {\n\tinfo func(ctx context.Context) (*schema.ToolInfo, error)\n}\n\nfunc (i *infoHelper) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn i.info(ctx)\n}\n\ntype errorHelper struct {\n\ti func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error)\n\th ErrorHandler\n}\n\nfunc (s *errorHelper) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tresult, err := s.i(ctx, argumentsInJSON, opts...)\n\tif _, ok := compose.IsInterruptRerunError(err); ok {\n\t\treturn result, err\n\t}\n\tif err != nil {\n\t\treturn s.h(ctx, err), nil\n\t}\n\treturn result, nil\n}\n\ntype streamErrorHelper struct {\n\ts func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error)\n\th ErrorHandler\n}\n\nfunc (s *streamErrorHelper) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\tresult, err := s.s(ctx, argumentsInJSON, opts...)\n\tif _, ok := compose.IsInterruptRerunError(err); ok {\n\t\treturn result, err\n\t}\n\tif err != nil {\n\t\treturn schema.StreamReaderFromArray([]string{s.h(ctx, err)}), nil\n\t}\n\treturn result, nil\n}\n"
  },
  {
    "path": "components/tool/utils/error_handler_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"io\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype testErrorTool struct{}\n\nfunc (t *testErrorTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn nil, nil\n}\n\nfunc (t *testErrorTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\treturn \"\", errors.New(\"test error\")\n}\n\nfunc (t *testErrorTool) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\treturn nil, errors.New(\"test stream error\")\n}\n\nfunc TestErrorWrapper(t *testing.T) {\n\tctx := context.Background()\n\n\tnt := WrapToolWithErrorHandler(&testErrorTool{}, func(_ context.Context, err error) string {\n\t\treturn err.Error()\n\t})\n\tresult, err := nt.(tool.InvokableTool).InvokableRun(ctx, \"\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"test error\", result)\n\tstreamResult, err := nt.(tool.StreamableTool).StreamableRun(ctx, \"\")\n\tassert.NoError(t, err)\n\tchunk, err := streamResult.Recv()\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"test stream error\", chunk)\n\t_, err = streamResult.Recv()\n\tassert.True(t, errors.Is(err, io.EOF))\n\n\twrappedTool := WrapInvokableToolWithErrorHandler(&testErrorTool{}, func(_ context.Context, err error) string {\n\t\treturn err.Error()\n\t})\n\tresult, err = wrappedTool.InvokableRun(ctx, \"\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"test error\", result)\n\n\twrappedStreamTool := WrapStreamableToolWithErrorHandler(&testErrorTool{}, func(_ context.Context, err error) string {\n\t\treturn err.Error()\n\t})\n\n\tstreamResult, err = wrappedStreamTool.StreamableRun(ctx, \"\")\n\tassert.NoError(t, err)\n\n\tchunk, err = streamResult.Recv()\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"test stream error\", chunk)\n\t_, err = streamResult.Recv()\n\tassert.True(t, errors.Is(err, io.EOF))\n}\n"
  },
  {
    "path": "components/tool/utils/invokable_func.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/eino-contrib/jsonschema\"\n\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// InvokeFunc is the function type for the tool.\ntype InvokeFunc[T, D any] func(ctx context.Context, input T) (output D, err error)\n\n// OptionableInvokeFunc is the function type for the tool with tool option.\ntype OptionableInvokeFunc[T, D any] func(ctx context.Context, input T, opts ...tool.Option) (output D, err error)\n\n// InferTool creates an [tool.InvokableTool] by inferring the parameter JSON\n// schema from the fields and tags of the input type T.\n//\n// The tool automatically JSON-decodes the model's argument string into T before\n// calling fn, and JSON-encodes the D return value into the result string.\n//\n// Use [WithSchemaModifier] in opts to customise how struct tags are mapped to\n// JSON schema fields.\nfunc InferTool[T, D any](toolName, toolDesc string, i InvokeFunc[T, D], opts ...Option) (tool.InvokableTool, error) {\n\tti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn NewTool(ti, i, opts...), nil\n}\n\n// InferOptionableTool is like [InferTool] but the function also receives\n// [tool.Option] values passed by ToolsNode at call time.\nfunc InferOptionableTool[T, D any](toolName, toolDesc string, i OptionableInvokeFunc[T, D], opts ...Option) (tool.InvokableTool, error) {\n\tti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn newOptionableTool(ti, i, opts...), nil\n}\n\n// EnhancedInvokeFunc is the function type for the enhanced tool.\ntype EnhancedInvokeFunc[T any] func(ctx context.Context, input T) (output *schema.ToolResult, err error)\n\n// OptionableEnhancedInvokeFunc is the function type for the enhanced tool with tool option.\ntype OptionableEnhancedInvokeFunc[T any] func(ctx context.Context, input T, opts ...tool.Option) (output *schema.ToolResult, err error)\n\n// InferEnhancedTool creates an [tool.EnhancedInvokableTool] by inferring the\n// parameter JSON schema from type T. The function returns a [schema.ToolResult]\n// for multimodal output (text, images, audio, video, files).\nfunc InferEnhancedTool[T any](toolName, toolDesc string, i EnhancedInvokeFunc[T], opts ...Option) (tool.EnhancedInvokableTool, error) {\n\tti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn NewEnhancedTool(ti, i, opts...), nil\n}\n\n// InferOptionableEnhancedTool creates an EnhancedInvokableTool from a given function by inferring the ToolInfo from the function's request parameters, with tool option.\nfunc InferOptionableEnhancedTool[T any](toolName, toolDesc string, i OptionableEnhancedInvokeFunc[T], opts ...Option) (tool.EnhancedInvokableTool, error) {\n\tti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn newOptionableEnhancedTool(ti, i, opts...), nil\n}\n\n// GoStruct2ParamsOneOf converts a Go struct's fields and tags into a\n// [schema.ParamsOneOf] (JSON Schema 2020-12). Useful for ChatModel structured\n// output via ResponseFormat without creating a full tool.\nfunc GoStruct2ParamsOneOf[T any](opts ...Option) (*schema.ParamsOneOf, error) {\n\treturn goStruct2ParamsOneOf[T](opts...)\n}\n\n// GoStruct2ToolInfo converts a Go struct into a [schema.ToolInfo]. Useful for\n// binding a typed schema to a ChatModel via BindTools for structured output,\n// when you do not need a full executable tool.\nfunc GoStruct2ToolInfo[T any](toolName, toolDesc string, opts ...Option) (*schema.ToolInfo, error) {\n\treturn goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n}\n\nfunc goStruct2ToolInfo[T any](toolName, toolDesc string, opts ...Option) (*schema.ToolInfo, error) {\n\tparamsOneOf, err := goStruct2ParamsOneOf[T](opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn &schema.ToolInfo{\n\t\tName:        toolName,\n\t\tDesc:        toolDesc,\n\t\tParamsOneOf: paramsOneOf,\n\t}, nil\n}\n\nfunc goStruct2ParamsOneOf[T any](opts ...Option) (*schema.ParamsOneOf, error) {\n\toptions := getToolOptions(opts...)\n\n\tr := &jsonschema.Reflector{\n\t\tAnonymous:      true,\n\t\tDoNotReference: true,\n\t\tSchemaModifier: jsonschema.SchemaModifierFn(options.scModifier),\n\t}\n\n\tjs := r.Reflect(generic.NewInstance[T]())\n\tjs.Version = \"\"\n\n\tparamsOneOf := schema.NewParamsOneOfByJSONSchema(js)\n\n\treturn paramsOneOf, nil\n}\n\n// NewTool creates an [tool.InvokableTool] from an explicit [schema.ToolInfo]\n// and a typed function. Use this when the schema cannot be inferred from struct\n// tags (e.g. dynamic or complex parameter schemas).\n//\n// Note: you are responsible for keeping desc.ParamsOneOf consistent with the\n// actual fields of T — there is no compile-time check.\nfunc NewTool[T, D any](desc *schema.ToolInfo, i InvokeFunc[T, D], opts ...Option) tool.InvokableTool {\n\treturn newOptionableTool(desc, func(ctx context.Context, input T, _ ...tool.Option) (D, error) {\n\t\treturn i(ctx, input)\n\t}, opts...)\n}\n\nfunc newOptionableTool[T, D any](desc *schema.ToolInfo, i OptionableInvokeFunc[T, D], opts ...Option) tool.InvokableTool {\n\tto := getToolOptions(opts...)\n\n\treturn &invokableTool[T, D]{\n\t\tinfo: desc,\n\t\tum:   to.um,\n\t\tm:    to.m,\n\t\tFn:   i,\n\t}\n}\n\ntype invokableTool[T, D any] struct {\n\tinfo *schema.ToolInfo\n\n\tum UnmarshalArguments\n\tm  MarshalOutput\n\n\tFn OptionableInvokeFunc[T, D]\n}\n\nfunc (i *invokableTool[T, D]) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn i.info, nil\n}\n\n// InvokableRun invokes the tool with the given arguments.\nfunc (i *invokableTool[T, D]) InvokableRun(ctx context.Context, arguments string, opts ...tool.Option) (output string, err error) {\n\n\tvar inst T\n\tif i.um != nil {\n\t\tvar val any\n\t\tval, err = i.um(ctx, arguments)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"[LocalFunc] failed to unmarshal arguments, toolName=%s, err=%w\", i.getToolName(), err)\n\t\t}\n\t\tgt, ok := val.(T)\n\t\tif !ok {\n\t\t\treturn \"\", fmt.Errorf(\"[LocalFunc] invalid type, toolName=%s, expected=%T, given=%T\", i.getToolName(), inst, val)\n\t\t}\n\t\tinst = gt\n\t} else {\n\t\tinst = generic.NewInstance[T]()\n\n\t\terr = sonic.UnmarshalString(arguments, &inst)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"[LocalFunc] failed to unmarshal arguments in json, toolName=%s, err=%w\", i.getToolName(), err)\n\t\t}\n\t}\n\n\tresp, err := i.Fn(ctx, inst, opts...)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"[LocalFunc] failed to invoke tool, toolName=%s, err=%w\", i.getToolName(), err)\n\t}\n\n\tif i.m != nil {\n\t\toutput, err = i.m(ctx, resp)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"[LocalFunc] failed to marshal output, toolName=%s, err=%w\", i.getToolName(), err)\n\t\t}\n\t} else {\n\t\toutput, err = marshalString(resp)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"[LocalFunc] failed to marshal output in json, toolName=%s, err=%w\", i.getToolName(), err)\n\t\t}\n\t}\n\n\treturn output, nil\n}\n\nfunc (i *invokableTool[T, D]) GetType() string {\n\treturn snakeToCamel(i.getToolName())\n}\n\nfunc (i *invokableTool[T, D]) getToolName() string {\n\tif i.info == nil {\n\t\treturn \"\"\n\t}\n\n\treturn i.info.Name\n}\n\n// snakeToCamel converts a snake_case string to CamelCase.\nfunc snakeToCamel(s string) string {\n\tif s == \"\" {\n\t\treturn \"\"\n\t}\n\n\tparts := strings.Split(s, \"_\")\n\n\tfor i := 0; i < len(parts); i++ {\n\t\tif len(parts[i]) > 0 {\n\t\t\tparts[i] = strings.ToUpper(string(parts[i][0])) + strings.ToLower(parts[i][1:])\n\t\t}\n\t}\n\n\treturn strings.Join(parts, \"\")\n}\n\n// NewEnhancedTool creates an [tool.EnhancedInvokableTool] from an explicit\n// [schema.ToolInfo] and a function that returns [schema.ToolResult].\nfunc NewEnhancedTool[T any](desc *schema.ToolInfo, i EnhancedInvokeFunc[T], opts ...Option) tool.EnhancedInvokableTool {\n\treturn newOptionableEnhancedTool(desc, func(ctx context.Context, input T, _ ...tool.Option) (*schema.ToolResult, error) {\n\t\treturn i(ctx, input)\n\t}, opts...)\n}\n\nfunc newOptionableEnhancedTool[T any](desc *schema.ToolInfo, i OptionableEnhancedInvokeFunc[T], opts ...Option) tool.EnhancedInvokableTool {\n\tto := getToolOptions(opts...)\n\n\treturn &enhancedInvokableTool[T]{\n\t\tinfo: desc,\n\t\tum:   to.um,\n\t\tFn:   i,\n\t}\n}\n\ntype enhancedInvokableTool[T any] struct {\n\tinfo *schema.ToolInfo\n\n\tum UnmarshalArguments\n\n\tFn OptionableEnhancedInvokeFunc[T]\n}\n\nfunc (e *enhancedInvokableTool[T]) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn e.info, nil\n}\n\nfunc (e *enhancedInvokableTool[T]) InvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\tvar inst T\n\tvar err error\n\n\tif e.um != nil {\n\t\tvar val any\n\t\tval, err = e.um(ctx, toolArgument.Text)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"[EnhancedLocalFunc] failed to unmarshal arguments, toolName=%s, err=%w\", e.getToolName(), err)\n\t\t}\n\t\tgt, ok := val.(T)\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"[EnhancedLocalFunc] invalid type, toolName=%s, expected=%T, given=%T\", e.getToolName(), inst, val)\n\t\t}\n\t\tinst = gt\n\t} else {\n\t\tinst = generic.NewInstance[T]()\n\n\t\terr = sonic.UnmarshalString(toolArgument.Text, &inst)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"[EnhancedLocalFunc] failed to unmarshal arguments in json, toolName=%s, err=%w\", e.getToolName(), err)\n\t\t}\n\t}\n\n\tresp, err := e.Fn(ctx, inst, opts...)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"[EnhancedLocalFunc] failed to invoke tool, toolName=%s, err=%w\", e.getToolName(), err)\n\t}\n\n\treturn resp, nil\n}\n\nfunc (e *enhancedInvokableTool[T]) GetType() string {\n\treturn snakeToCamel(e.getToolName())\n}\n\nfunc (e *enhancedInvokableTool[T]) getToolName() string {\n\tif e.info == nil {\n\t\treturn \"\"\n\t}\n\n\treturn e.info.Name\n}\n"
  },
  {
    "path": "components/tool/utils/invokable_func_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"testing\"\n\n\t\"github.com/eino-contrib/jsonschema\"\n\t\"github.com/stretchr/testify/assert\"\n\torderedmap \"github.com/wk8/go-ordered-map/v2\"\n\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype Job struct {\n\tCompany       string  `json:\"company\" jsonschema:\"description=the company where the user works\"`\n\tPosition      string  `json:\"position,omitempty\" jsonschema:\"description=the position of the user's job\"`\n\tServiceLength float32 `json:\"service_length,omitempty\" jsonschema:\"description=the year of user's service\"` // 司龄，年\n}\n\ntype Income struct {\n\tSource    string `json:\"source\" jsonschema:\"description=the source of income\"`\n\tAmount    int    `json:\"amount\" jsonschema:\"description=the amount of income\"`\n\tHasPayTax bool   `json:\"has_pay_tax\" jsonschema:\"description=whether the user has paid tax\"`\n\tJob       *Job   `json:\"job,omitempty\" jsonschema:\"description=the job of the user when earning this income\"`\n}\n\ntype User struct {\n\tName string `json:\"name\" jsonschema:\"required,description=the name of the user\"`\n\tAge  int    `json:\"age\" jsonschema:\"required,description=the age of the user\"`\n\n\tJob *Job `json:\"job,omitempty\" jsonschema:\"description=the job of the user\"`\n\n\tIncomes []*Income `json:\"incomes\" jsonschema:\"description=the incomes of the user\"`\n}\n\ntype UserResult struct {\n\tCode int    `json:\"code\"`\n\tMsg  string `json:\"msg\"`\n}\n\nvar toolInfo = &schema.ToolInfo{\n\tName: \"update_user_info\",\n\tDesc: \"full update user info\",\n\tParamsOneOf: schema.NewParamsOneOfByJSONSchema(\n\t\t&jsonschema.Schema{\n\t\t\tType:                 \"object\",\n\t\t\tRequired:             []string{\"name\", \"age\", \"incomes\"},\n\t\t\tAdditionalProperties: jsonschema.FalseSchema,\n\t\t\tProperties: orderedmap.New[string, *jsonschema.Schema](\n\t\t\t\torderedmap.WithInitialData(\n\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\tKey: \"name\",\n\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\tType:        \"string\",\n\t\t\t\t\t\t\tDescription: \"the name of the user\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\tKey: \"age\",\n\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\tType:        \"integer\",\n\t\t\t\t\t\t\tDescription: \"the age of the user\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\tKey: \"job\",\n\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\tType:                 \"object\",\n\t\t\t\t\t\t\tRequired:             []string{\"company\"},\n\t\t\t\t\t\t\tAdditionalProperties: jsonschema.FalseSchema,\n\t\t\t\t\t\t\tDescription:          \"the job of the user\",\n\t\t\t\t\t\t\tProperties: orderedmap.New[string, *jsonschema.Schema](\n\t\t\t\t\t\t\t\torderedmap.WithInitialData(\n\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\tKey: \"company\",\n\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\tType:        \"string\",\n\t\t\t\t\t\t\t\t\t\t\tDescription: \"the company where the user works\",\n\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\tKey: \"position\",\n\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\tType:        \"string\",\n\t\t\t\t\t\t\t\t\t\t\tDescription: \"the position of the user's job\",\n\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\tKey: \"service_length\",\n\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\tType:        \"number\",\n\t\t\t\t\t\t\t\t\t\t\tDescription: \"the year of user's service\",\n\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t),\n\t\t\t\t\t\t\t),\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\tKey: \"incomes\",\n\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\tType:        \"array\",\n\t\t\t\t\t\t\tDescription: \"the incomes of the user\",\n\t\t\t\t\t\t\tItems: &jsonschema.Schema{\n\t\t\t\t\t\t\t\tType:                 \"object\",\n\t\t\t\t\t\t\t\tAdditionalProperties: jsonschema.FalseSchema,\n\t\t\t\t\t\t\t\tRequired:             []string{\"source\", \"amount\", \"has_pay_tax\"},\n\t\t\t\t\t\t\t\tProperties: orderedmap.New[string, *jsonschema.Schema](\n\t\t\t\t\t\t\t\t\torderedmap.WithInitialData(\n\t\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\t\tKey: \"source\",\n\t\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\t\tType:        \"string\",\n\t\t\t\t\t\t\t\t\t\t\t\tDescription: \"the source of income\",\n\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\t\tKey: \"amount\",\n\t\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\t\tType:        \"integer\",\n\t\t\t\t\t\t\t\t\t\t\t\tDescription: \"the amount of income\",\n\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\t\tKey: \"has_pay_tax\",\n\t\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\t\tType:        \"boolean\",\n\t\t\t\t\t\t\t\t\t\t\t\tDescription: \"whether the user has paid tax\",\n\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\t\tKey: \"job\",\n\t\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\t\tType:                 \"object\",\n\t\t\t\t\t\t\t\t\t\t\t\tAdditionalProperties: jsonschema.FalseSchema,\n\t\t\t\t\t\t\t\t\t\t\t\tRequired:             []string{\"company\"},\n\t\t\t\t\t\t\t\t\t\t\t\tDescription:          \"the job of the user when earning this income\",\n\t\t\t\t\t\t\t\t\t\t\t\tProperties: orderedmap.New[string, *jsonschema.Schema](\n\t\t\t\t\t\t\t\t\t\t\t\t\torderedmap.WithInitialData(\n\t\t\t\t\t\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tKey: \"company\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tType:        \"string\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tDescription: \"the company where the user works\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tKey: \"position\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tType:        \"string\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tDescription: \"the position of the user's job\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tKey: \"service_length\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tType:        \"number\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tDescription: \"the year of user's service\",\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t\t\t\t),\n\t\t\t\t\t\t\t\t\t\t\t\t),\n\t\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t),\n\t\t\t\t\t\t\t\t),\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t),\n\t\t\t),\n\t\t}),\n}\n\nfunc updateUserInfo(ctx context.Context, input *User) (output *UserResult, err error) {\n\treturn &UserResult{\n\t\tCode: 200,\n\t\tMsg:  fmt.Sprintf(\"update %v success\", input.Name),\n\t}, nil\n}\n\ntype UserInfoOption struct {\n\tField1 string\n}\n\nfunc WithUserInfoOption(s string) tool.Option {\n\treturn tool.WrapImplSpecificOptFn(func(t *UserInfoOption) {\n\t\tt.Field1 = s\n\t})\n}\n\nfunc updateUserInfoWithOption(_ context.Context, input *User, opts ...tool.Option) (output *UserResult, err error) {\n\tbaseOption := &UserInfoOption{\n\t\tField1: \"test_origin\",\n\t}\n\n\toption := tool.GetImplSpecificOptions(baseOption, opts...)\n\treturn &UserResult{\n\t\tCode: 200,\n\t\tMsg:  option.Field1,\n\t}, nil\n}\n\nfunc TestInferTool(t *testing.T) {\n\tt.Run(\"invoke_infer_tool\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\ttl, err := InferTool(\"update_user_info\", \"full update user info\", updateUserInfo)\n\t\tassert.NoError(t, err)\n\n\t\tinfo, err := tl.Info(context.Background())\n\t\tassert.NoError(t, err)\n\n\t\tactual, err := info.ToJSONSchema()\n\t\tassert.NoError(t, err)\n\t\tactualStr, err := json.Marshal(actual)\n\t\tassert.NoError(t, err)\n\n\t\texpect, err := toolInfo.ToJSONSchema()\n\t\tassert.NoError(t, err)\n\t\texpectStr, err := json.Marshal(expect)\n\t\tassert.NoError(t, err)\n\n\t\tassert.Equal(t, string(expectStr), string(actualStr))\n\n\t\tcontent, err := tl.InvokableRun(ctx, `{\"name\": \"bruce lee\"}`)\n\t\tassert.NoError(t, err)\n\t\tassert.JSONEq(t, `{\"code\":200,\"msg\":\"update bruce lee success\"}`, content)\n\t})\n}\n\nfunc TestInferOptionableTool(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"invoke_infer_optionable_tool\", func(t *testing.T) {\n\n\t\ttl, err := InferOptionableTool(\"invoke_infer_optionable_tool\", \"full update user info\", updateUserInfoWithOption)\n\t\tassert.NoError(t, err)\n\n\t\tcontent, err := tl.InvokableRun(ctx, `{\"name\": \"bruce lee\"}`, WithUserInfoOption(\"hello world\"))\n\t\tassert.NoError(t, err)\n\t\tassert.JSONEq(t, `{\"code\":200,\"msg\":\"hello world\"}`, content)\n\t})\n}\n\nfunc TestNewTool(t *testing.T) {\n\tctx := context.Background()\n\ttype Input struct {\n\t\tName string `json:\"name\"`\n\t}\n\ttype Output struct {\n\t\tName string `json:\"name\"`\n\t}\n\n\tt.Run(\"struct_input_struct_output\", func(t *testing.T) {\n\n\t\ttl := NewTool[Input, Output](nil, func(ctx context.Context, input Input) (output Output, err error) {\n\t\t\treturn Output{\n\t\t\t\tName: input.Name,\n\t\t\t}, nil\n\t\t})\n\n\t\t_, err := tl.InvokableRun(ctx, `{\"name\":\"test\"}`)\n\t\tassert.Nil(t, err)\n\t})\n\n\tt.Run(\"pointer_input_pointer_output\", func(t *testing.T) {\n\t\ttl := NewTool[*Input, *Output](nil, func(ctx context.Context, input *Input) (output *Output, err error) {\n\t\t\treturn &Output{\n\t\t\t\tName: input.Name,\n\t\t\t}, nil\n\t\t})\n\n\t\tcontent, err := tl.InvokableRun(ctx, `{\"name\":\"test\"}`)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, `{\"name\":\"test\"}`, content)\n\t})\n\n\tt.Run(\"string_input_int64_output\", func(t *testing.T) {\n\t\ttl := NewTool(nil, func(ctx context.Context, input string) (output int64, err error) {\n\t\t\treturn 10, nil\n\t\t})\n\n\t\tcontent, err := tl.InvokableRun(ctx, `100`) // json unmarshal must contains double quote if is not json string.\n\t\tassert.Error(t, err)\n\t\tassert.Equal(t, \"\", content)\n\t})\n\n\tt.Run(\"string_pointer_input_int64_pointer_output\", func(t *testing.T) {\n\t\ttl := NewTool[*string, *int64](nil, func(ctx context.Context, input *string) (output *int64, err error) {\n\t\t\tn := int64(10)\n\t\t\treturn &n, nil\n\t\t})\n\n\t\tcontent, err := tl.InvokableRun(ctx, `\"100\"`)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, `10`, content)\n\t})\n}\n\nfunc TestSnakeToCamel(t *testing.T) {\n\tt.Run(\"normal_case\", func(t *testing.T) {\n\t\tassert.Equal(t, \"GoogleSearch3\", snakeToCamel(\"google_search_3\"))\n\t})\n\n\tt.Run(\"empty_case\", func(t *testing.T) {\n\t\tassert.Equal(t, \"\", snakeToCamel(\"\"))\n\t})\n\n\tt.Run(\"single_word_case\", func(t *testing.T) {\n\t\tassert.Equal(t, \"Google\", snakeToCamel(\"google\"))\n\t})\n\n\tt.Run(\"upper_case\", func(t *testing.T) {\n\t\tassert.Equal(t, \"HttpHost\", snakeToCamel(\"_HTTP_HOST_\"))\n\t})\n\n\tt.Run(\"underscore_case\", func(t *testing.T) {\n\t\tassert.Equal(t, \"\", snakeToCamel(\"_\"))\n\t})\n}\n\ntype stringAlias string\ntype integerAlias uint32\ntype floatAlias float64\ntype boolAlias bool\n\ntype testEnumStruct struct {\n\tField1 string       `json:\"field1\" jsonschema:\"enum=a,enum=b\"`\n\tField2 int          `json:\"field2\" jsonschema:\"enum=1,enum=2\"`\n\tField3 float32      `json:\"field3\" jsonschema:\"enum=1.1,enum=2.2\"`\n\tField4 bool         `json:\"field4\" jsonschema:\"default=true\"`\n\tField5 stringAlias  `json:\"field5\" jsonschema:\"enum=a,enum=c\"`\n\tField6 integerAlias `json:\"field6\" jsonschema:\"enum=3,enum=4\"`\n\tField7 floatAlias   `json:\"field7\" jsonschema:\"enum=3.3,enum=4.4\"`\n\tField8 boolAlias    `json:\"field8\" jsonschema:\"enum=false\"`\n}\n\ntype testEnumStruct2 struct {\n\tField1 int8 `json:\"field1\" jsonschema:\"enum=1.1\"`\n}\n\ntype testEnumStruct3 struct {\n\tField1 float64 `json:\"field1\" jsonschema:\"enum=a\"`\n}\n\nfunc TestEnumTag(t *testing.T) {\n\tinfo, err := goStruct2ParamsOneOf[testEnumStruct]()\n\tassert.NoError(t, err)\n\ts, err := info.ToJSONSchema()\n\tassert.NoError(t, err)\n\n\tenum, ok := s.Properties.Get(\"field1\")\n\tassert.True(t, ok)\n\tassert.Equal(t, []any{\"a\", \"b\"}, enum.Enum)\n\n\tenum, ok = s.Properties.Get(\"field2\")\n\tassert.True(t, ok)\n\tassert.Equal(t, []any{json.Number(\"1\"), json.Number(\"2\")}, enum.Enum)\n\n\tenum, ok = s.Properties.Get(\"field3\")\n\tassert.True(t, ok)\n\tassert.Equal(t, []any{json.Number(\"1.1\"), json.Number(\"2.2\")}, enum.Enum)\n\n\tenum, ok = s.Properties.Get(\"field4\")\n\tassert.True(t, ok)\n\tassert.Equal(t, true, enum.Default)\n\n\tenum, ok = s.Properties.Get(\"field5\")\n\tassert.True(t, ok)\n\tassert.Equal(t, []any{\"a\", \"c\"}, enum.Enum)\n\n\tenum, ok = s.Properties.Get(\"field6\")\n\tassert.True(t, ok)\n\tassert.Equal(t, []any{json.Number(\"3\"), json.Number(\"4\")}, enum.Enum)\n\n\tenum, ok = s.Properties.Get(\"field7\")\n\tassert.True(t, ok)\n\tassert.Equal(t, []any{json.Number(\"3.3\"), json.Number(\"4.4\")}, enum.Enum)\n\n\t_, err = goStruct2ParamsOneOf[testEnumStruct2]()\n\tassert.NoError(t, err)\n\n\t_, err = goStruct2ParamsOneOf[testEnumStruct3]()\n\tassert.NoError(t, err)\n}\n"
  },
  {
    "path": "components/tool/utils/streamable_func.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/bytedance/sonic\"\n\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// StreamFunc is the function type for the streamable tool.\ntype StreamFunc[T, D any] func(ctx context.Context, input T) (output *schema.StreamReader[D], err error)\n\n// OptionableStreamFunc is the function type for the streamable tool with tool option.\ntype OptionableStreamFunc[T, D any] func(ctx context.Context, input T, opts ...tool.Option) (output *schema.StreamReader[D], err error)\n\n// InferStreamTool creates a [tool.StreamableTool] by inferring the parameter\n// JSON schema from type T. The function returns a [schema.StreamReader] of D\n// values which the framework serialises to a string stream.\nfunc InferStreamTool[T, D any](toolName, toolDesc string, s StreamFunc[T, D], opts ...Option) (tool.StreamableTool, error) {\n\tti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn NewStreamTool(ti, s, opts...), nil\n}\n\n// InferOptionableStreamTool is like [InferStreamTool] but the function also\n// receives [tool.Option] values passed by ToolsNode at call time.\nfunc InferOptionableStreamTool[T, D any](toolName, toolDesc string, s OptionableStreamFunc[T, D], opts ...Option) (tool.StreamableTool, error) {\n\tti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn newOptionableStreamTool(ti, s, opts...), nil\n}\n\n// NewStreamTool creates a [tool.StreamableTool] from an explicit [schema.ToolInfo]\n// and a typed streaming function.\nfunc NewStreamTool[T, D any](desc *schema.ToolInfo, s StreamFunc[T, D], opts ...Option) tool.StreamableTool {\n\treturn newOptionableStreamTool(desc,\n\t\tfunc(ctx context.Context, input T, _ ...tool.Option) (output *schema.StreamReader[D], err error) {\n\t\t\treturn s(ctx, input)\n\t\t},\n\t\topts...)\n}\n\nfunc newOptionableStreamTool[T, D any](desc *schema.ToolInfo, s OptionableStreamFunc[T, D], opts ...Option) tool.StreamableTool {\n\n\tto := getToolOptions(opts...)\n\n\treturn &streamableTool[T, D]{\n\t\tinfo: desc,\n\n\t\tum: to.um,\n\t\tm:  to.m,\n\t\tFn: s,\n\t}\n}\n\ntype streamableTool[T, D any] struct {\n\tinfo *schema.ToolInfo\n\n\tum UnmarshalArguments\n\tm  MarshalOutput\n\n\tFn OptionableStreamFunc[T, D]\n}\n\n// Info returns the tool info, implement the BaseTool interface.\nfunc (s *streamableTool[T, D]) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn s.info, nil\n}\n\n// StreamableRun invokes the tool with the given arguments, implement the StreamableTool interface.\nfunc (s *streamableTool[T, D]) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (\n\toutStream *schema.StreamReader[string], err error) {\n\n\tvar inst T\n\tif s.um != nil {\n\t\tvar val any\n\t\tval, err = s.um(ctx, argumentsInJSON)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"[LocalStreamFunc] failed to unmarshal arguments, toolName=%s, err=%w\", s.getToolName(), err)\n\t\t}\n\n\t\tgt, ok := val.(T)\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"[LocalStreamFunc] type err, toolName=%s, expected=%T, given=%T\", s.getToolName(), inst, val)\n\t\t}\n\t\tinst = gt\n\t} else {\n\n\t\tinst = generic.NewInstance[T]()\n\n\t\terr = sonic.UnmarshalString(argumentsInJSON, &inst)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"[LocalStreamFunc] failed to unmarshal arguments in json, toolName=%s, err=%w\", s.getToolName(), err)\n\t\t}\n\t}\n\n\tstreamD, err := s.Fn(ctx, inst, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\toutStream = schema.StreamReaderWithConvert(streamD, func(d D) (string, error) {\n\t\tvar out string\n\t\tvar e error\n\t\tif s.m != nil {\n\t\t\tout, e = s.m(ctx, d)\n\t\t\tif e != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"[LocalStreamFunc] failed to marshal output, toolName=%s, err=%w\", s.getToolName(), e)\n\t\t\t}\n\t\t} else {\n\t\t\tout, e = marshalString(d)\n\t\t\tif e != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"[LocalStreamFunc] failed to marshal output in json, toolName=%s, err=%w\", s.getToolName(), e)\n\t\t\t}\n\t\t}\n\n\t\treturn out, nil\n\t})\n\n\treturn outStream, nil\n}\n\nfunc (s *streamableTool[T, D]) GetType() string {\n\treturn snakeToCamel(s.getToolName())\n}\n\nfunc (s *streamableTool[T, D]) getToolName() string {\n\tif s.info == nil {\n\t\treturn \"\"\n\t}\n\n\treturn s.info.Name\n}\n\n// EnhancedStreamFunc is the function type for the enhanced streamable tool.\ntype EnhancedStreamFunc[T any] func(ctx context.Context, input T) (output *schema.StreamReader[*schema.ToolResult], err error)\n\n// OptionableEnhancedStreamFunc is the function type for the enhanced streamable tool with tool option.\ntype OptionableEnhancedStreamFunc[T any] func(ctx context.Context, input T, opts ...tool.Option) (output *schema.StreamReader[*schema.ToolResult], err error)\n\n// InferEnhancedStreamTool creates an [tool.EnhancedStreamableTool] by inferring\n// the parameter JSON schema from type T. The function streams [schema.ToolResult]\n// values for multimodal output.\nfunc InferEnhancedStreamTool[T any](toolName, toolDesc string, s EnhancedStreamFunc[T], opts ...Option) (tool.EnhancedStreamableTool, error) {\n\tti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn NewEnhancedStreamTool(ti, s, opts...), nil\n}\n\n// InferOptionableEnhancedStreamTool creates an EnhancedStreamableTool from a given function by inferring the ToolInfo from the function's request parameters, with tool option.\nfunc InferOptionableEnhancedStreamTool[T any](toolName, toolDesc string, s OptionableEnhancedStreamFunc[T], opts ...Option) (tool.EnhancedStreamableTool, error) {\n\tti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn newOptionableEnhancedStreamTool(ti, s, opts...), nil\n}\n\n// NewEnhancedStreamTool creates an [tool.EnhancedStreamableTool] from an\n// explicit [schema.ToolInfo] and a typed streaming function.\nfunc NewEnhancedStreamTool[T any](desc *schema.ToolInfo, s EnhancedStreamFunc[T], opts ...Option) tool.EnhancedStreamableTool {\n\treturn newOptionableEnhancedStreamTool(desc,\n\t\tfunc(ctx context.Context, input T, _ ...tool.Option) (output *schema.StreamReader[*schema.ToolResult], err error) {\n\t\t\treturn s(ctx, input)\n\t\t},\n\t\topts...)\n}\n\nfunc newOptionableEnhancedStreamTool[T any](desc *schema.ToolInfo, s OptionableEnhancedStreamFunc[T], opts ...Option) tool.EnhancedStreamableTool {\n\tto := getToolOptions(opts...)\n\n\treturn &enhancedStreamableTool[T]{\n\t\tinfo: desc,\n\t\tum:   to.um,\n\t\tFn:   s,\n\t}\n}\n\ntype enhancedStreamableTool[T any] struct {\n\tinfo *schema.ToolInfo\n\n\tum UnmarshalArguments\n\n\tFn OptionableEnhancedStreamFunc[T]\n}\n\nfunc (s *enhancedStreamableTool[T]) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn s.info, nil\n}\n\nfunc (s *enhancedStreamableTool[T]) StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (\n\toutStream *schema.StreamReader[*schema.ToolResult], err error) {\n\n\tvar inst T\n\tif s.um != nil {\n\t\tvar val any\n\t\tval, err = s.um(ctx, toolArgument.Text)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"[EnhancedLocalStreamFunc] failed to unmarshal arguments, toolName=%s, err=%w\", s.getToolName(), err)\n\t\t}\n\n\t\tgt, ok := val.(T)\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"[EnhancedLocalStreamFunc] type err, toolName=%s, expected=%T, given=%T\", s.getToolName(), inst, val)\n\t\t}\n\t\tinst = gt\n\t} else {\n\t\tinst = generic.NewInstance[T]()\n\n\t\terr = sonic.UnmarshalString(toolArgument.Text, &inst)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"[EnhancedLocalStreamFunc] failed to unmarshal arguments in json, toolName=%s, err=%w\", s.getToolName(), err)\n\t\t}\n\t}\n\n\treturn s.Fn(ctx, inst, opts...)\n}\n\nfunc (s *enhancedStreamableTool[T]) GetType() string {\n\treturn snakeToCamel(s.getToolName())\n}\n\nfunc (s *enhancedStreamableTool[T]) getToolName() string {\n\tif s.info == nil {\n\t\treturn \"\"\n\t}\n\n\treturn s.info.Name\n}\n"
  },
  {
    "path": "components/tool/utils/streamable_func_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage utils\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"io\"\n\t\"testing\"\n\n\t\"github.com/eino-contrib/jsonschema\"\n\t\"github.com/stretchr/testify/assert\"\n\torderedmap \"github.com/wk8/go-ordered-map/v2\"\n\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestNewStreamableTool(t *testing.T) {\n\tctx := context.Background()\n\ttype Input struct {\n\t\tName string `json:\"name\"`\n\t}\n\ttype Output struct {\n\t\tName string `json:\"name\"`\n\t}\n\n\tt.Run(\"simple_case\", func(t *testing.T) {\n\t\ttl := NewStreamTool[*Input, *Output](\n\t\t\t&schema.ToolInfo{\n\t\t\t\tName: \"search_user\",\n\t\t\t\tDesc: \"search user info\",\n\t\t\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\t\t\"name\": {\n\t\t\t\t\t\t\tType: \"string\",\n\t\t\t\t\t\t\tDesc: \"user name\",\n\t\t\t\t\t\t},\n\t\t\t\t\t}),\n\t\t\t},\n\t\t\tfunc(ctx context.Context, input *Input) (output *schema.StreamReader[*Output], err error) {\n\t\t\t\tsr, sw := schema.Pipe[*Output](2)\n\t\t\t\tsw.Send(&Output{\n\t\t\t\t\tName: input.Name,\n\t\t\t\t}, nil)\n\t\t\t\tsw.Send(&Output{\n\t\t\t\t\tName: \"lee\",\n\t\t\t\t}, nil)\n\t\t\t\tsw.Close()\n\n\t\t\t\treturn sr, nil\n\t\t\t},\n\t\t)\n\n\t\tinfo, err := tl.Info(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"search_user\", info.Name)\n\n\t\tjs, err := info.ToJSONSchema()\n\t\tassert.NoError(t, err)\n\n\t\tassert.Equal(t, &jsonschema.Schema{\n\t\t\tType: \"object\",\n\t\t\tProperties: orderedmap.New[string, *jsonschema.Schema](\n\t\t\t\torderedmap.WithInitialData[string, *jsonschema.Schema](\n\t\t\t\t\torderedmap.Pair[string, *jsonschema.Schema]{\n\t\t\t\t\t\tKey: \"name\",\n\t\t\t\t\t\tValue: &jsonschema.Schema{\n\t\t\t\t\t\t\tType:        \"string\",\n\t\t\t\t\t\t\tDescription: \"user name\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t),\n\t\t\t),\n\t\t\tRequired: make([]string, 0),\n\t\t}, js)\n\n\t\tsr, err := tl.StreamableRun(ctx, `{\"name\":\"xxx\"}`)\n\t\tassert.NoError(t, err)\n\n\t\tdefer sr.Close()\n\n\t\tidx := 0\n\t\tfor {\n\t\t\tm, err := sr.Recv()\n\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\n\t\t\tif idx == 0 {\n\t\t\t\tassert.Equal(t, `{\"name\":\"xxx\"}`, m)\n\t\t\t} else {\n\t\t\t\tassert.Equal(t, `{\"name\":\"lee\"}`, m)\n\t\t\t}\n\t\t\tidx++\n\t\t}\n\n\t\tassert.Equal(t, 2, idx)\n\t})\n}\n\ntype FakeStreamOption struct {\n\tField string\n}\n\ntype FakeStreamInferToolInput struct {\n\tField string `json:\"field\"`\n}\n\ntype FakeStreamInferToolOutput struct {\n\tField string `json:\"field\"`\n}\n\nfunc FakeWithToolOption(s string) tool.Option {\n\treturn tool.WrapImplSpecificOptFn(func(t *FakeStreamOption) {\n\t\tt.Field = s\n\t})\n}\n\nfunc fakeStreamFunc(ctx context.Context, input FakeStreamInferToolInput, opts ...tool.Option) (output *schema.StreamReader[*FakeStreamInferToolOutput], err error) {\n\tbaseOpt := &FakeStreamOption{\n\t\tField: \"default_field_value\",\n\t}\n\toption := tool.GetImplSpecificOptions(baseOpt, opts...)\n\n\treturn schema.StreamReaderFromArray([]*FakeStreamInferToolOutput{\n\t\t{\n\t\t\tField: option.Field,\n\t\t},\n\t}), nil\n}\n\nfunc TestInferStreamTool(t *testing.T) {\n\tst, err := InferOptionableStreamTool(\"infer_optionable_stream_tool\", \"test infer stream tool with option\", fakeStreamFunc)\n\tassert.Nil(t, err)\n\n\tsr, err := st.StreamableRun(context.Background(), `{\"field\": \"value\"}`, FakeWithToolOption(\"hello world\"))\n\tassert.Nil(t, err)\n\n\tdefer sr.Close()\n\n\tidx := 0\n\tfor {\n\t\tm, err := sr.Recv()\n\t\tif errors.Is(err, io.EOF) {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\n\t\tif idx == 0 {\n\t\t\tassert.JSONEq(t, `{\"field\":\"hello world\"}`, m)\n\t\t}\n\t}\n}\n\ntype EnhancedStreamInput struct {\n\tQuery string `json:\"query\" jsonschema:\"description=the search query\"`\n}\n\nfunc TestNewEnhancedStreamTool(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"simple_case\", func(t *testing.T) {\n\t\ttl := NewEnhancedStreamTool[*EnhancedStreamInput](\n\t\t\t&schema.ToolInfo{\n\t\t\t\tName: \"enhanced_stream_search\",\n\t\t\t\tDesc: \"search with enhanced stream output\",\n\t\t\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\t\t\"query\": {\n\t\t\t\t\t\t\tType: \"string\",\n\t\t\t\t\t\t\tDesc: \"the search query\",\n\t\t\t\t\t\t},\n\t\t\t\t\t}),\n\t\t\t},\n\t\t\tfunc(ctx context.Context, input *EnhancedStreamInput) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\t\tsr, sw := schema.Pipe[*schema.ToolResult](2)\n\t\t\t\tsw.Send(&schema.ToolResult{\n\t\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t\t{Type: schema.ToolPartTypeText, Text: \"result for: \" + input.Query},\n\t\t\t\t\t},\n\t\t\t\t}, nil)\n\t\t\t\tsw.Send(&schema.ToolResult{\n\t\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t\t{Type: schema.ToolPartTypeText, Text: \"more results\"},\n\t\t\t\t\t},\n\t\t\t\t}, nil)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t},\n\t\t)\n\n\t\tinfo, err := tl.Info(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"enhanced_stream_search\", info.Name)\n\n\t\tsr, err := tl.StreamableRun(ctx, &schema.ToolArgument{Text: `{\"query\":\"test\"}`})\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tidx := 0\n\t\tfor {\n\t\t\tm, err := sr.Recv()\n\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\n\t\t\tif idx == 0 {\n\t\t\t\tassert.Len(t, m.Parts, 1)\n\t\t\t\tassert.Equal(t, schema.ToolPartTypeText, m.Parts[0].Type)\n\t\t\t\tassert.Equal(t, \"result for: test\", m.Parts[0].Text)\n\t\t\t} else {\n\t\t\t\tassert.Len(t, m.Parts, 1)\n\t\t\t\tassert.Equal(t, \"more results\", m.Parts[0].Text)\n\t\t\t}\n\t\t\tidx++\n\t\t}\n\t\tassert.Equal(t, 2, idx)\n\t})\n}\n\ntype FakeEnhancedStreamOption struct {\n\tPrefix string\n}\n\nfunc FakeWithEnhancedStreamOption(prefix string) tool.Option {\n\treturn tool.WrapImplSpecificOptFn(func(t *FakeEnhancedStreamOption) {\n\t\tt.Prefix = prefix\n\t})\n}\n\nfunc fakeEnhancedStreamFunc(ctx context.Context, input EnhancedStreamInput) (*schema.StreamReader[*schema.ToolResult], error) {\n\treturn schema.StreamReaderFromArray([]*schema.ToolResult{\n\t\t{\n\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t{Type: schema.ToolPartTypeText, Text: \"result: \" + input.Query},\n\t\t\t},\n\t\t},\n\t}), nil\n}\n\nfunc fakeOptionableEnhancedStreamFunc(ctx context.Context, input EnhancedStreamInput, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\tbaseOpt := &FakeEnhancedStreamOption{\n\t\tPrefix: \"default\",\n\t}\n\toption := tool.GetImplSpecificOptions(baseOpt, opts...)\n\n\treturn schema.StreamReaderFromArray([]*schema.ToolResult{\n\t\t{\n\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t{Type: schema.ToolPartTypeText, Text: option.Prefix + \": \" + input.Query},\n\t\t\t},\n\t\t},\n\t}), nil\n}\n\nfunc TestInferEnhancedStreamTool(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"infer_enhanced_stream_tool\", func(t *testing.T) {\n\t\ttl, err := InferEnhancedStreamTool(\"infer_enhanced_stream\", \"test infer enhanced stream tool\", fakeEnhancedStreamFunc)\n\t\tassert.NoError(t, err)\n\n\t\tinfo, err := tl.Info(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"infer_enhanced_stream\", info.Name)\n\n\t\tsr, err := tl.StreamableRun(ctx, &schema.ToolArgument{Text: `{\"query\":\"hello\"}`})\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tm, err := sr.Recv()\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, m.Parts, 1)\n\t\tassert.Equal(t, \"result: hello\", m.Parts[0].Text)\n\t})\n}\n\nfunc TestInferOptionableEnhancedStreamTool(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"infer_optionable_enhanced_stream_tool\", func(t *testing.T) {\n\t\ttl, err := InferOptionableEnhancedStreamTool(\"infer_optionable_enhanced_stream\", \"test infer optionable enhanced stream tool\", fakeOptionableEnhancedStreamFunc)\n\t\tassert.NoError(t, err)\n\n\t\tinfo, err := tl.Info(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"infer_optionable_enhanced_stream\", info.Name)\n\n\t\tsr, err := tl.StreamableRun(ctx, &schema.ToolArgument{Text: `{\"query\":\"world\"}`}, FakeWithEnhancedStreamOption(\"custom\"))\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tm, err := sr.Recv()\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, m.Parts, 1)\n\t\tassert.Equal(t, \"custom: world\", m.Parts[0].Text)\n\t})\n\n\tt.Run(\"infer_optionable_enhanced_stream_tool_default_option\", func(t *testing.T) {\n\t\ttl, err := InferOptionableEnhancedStreamTool(\"infer_optionable_enhanced_stream\", \"test infer optionable enhanced stream tool\", fakeOptionableEnhancedStreamFunc)\n\t\tassert.NoError(t, err)\n\n\t\tsr, err := tl.StreamableRun(ctx, &schema.ToolArgument{Text: `{\"query\":\"test\"}`})\n\t\tassert.NoError(t, err)\n\t\tdefer sr.Close()\n\n\t\tm, err := sr.Recv()\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, m.Parts, 1)\n\t\tassert.Equal(t, \"default: test\", m.Parts[0].Text)\n\t})\n}\n"
  },
  {
    "path": "components/types.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package components defines common interfaces that describe component\n// types and callback capabilities used across Eino.\npackage components\n\n// Typer provides a human-readable type name for a component implementation.\n//\n// When implemented, the component's full display name in DevOps tooling\n// (visual debugger, IDE plugin, dashboards) becomes \"{GetType()}{ComponentKind}\"\n// — e.g. \"OpenAIChatModel\". Use CamelCase naming.\n//\n// Also used by [utils.InferTool] and similar constructors to set the display\n// name of tool instances.\ntype Typer interface {\n\tGetType() string\n}\n\n// GetType returns the type name for a component that implements Typer.\nfunc GetType(component any) (string, bool) {\n\tif typer, ok := component.(Typer); ok {\n\t\treturn typer.GetType(), true\n\t}\n\n\treturn \"\", false\n}\n\n// Checker controls whether the framework's automatic callback instrumentation\n// is active for a component.\n//\n// When IsCallbacksEnabled returns true, the framework skips its default\n// OnStart/OnEnd wrapping and trusts the component to invoke callbacks itself\n// at the correct points. Implement this when your component needs precise\n// control over callback timing or content — for example, when streaming\n// requires callbacks to fire mid-stream rather than only at completion.\ntype Checker interface {\n\tIsCallbacksEnabled() bool\n}\n\n// IsCallbacksEnabled reports whether a component implements Checker and enables callbacks.\nfunc IsCallbacksEnabled(i any) bool {\n\tif checker, ok := i.(Checker); ok {\n\t\treturn checker.IsCallbacksEnabled()\n\t}\n\n\treturn false\n}\n\n// Component names representing the different categories of components.\ntype Component string\n\nconst (\n\t// ComponentOfPrompt identifies chat template components.\n\tComponentOfPrompt Component = \"ChatTemplate\"\n\t// ComponentOfChatModel identifies chat model components.\n\tComponentOfChatModel Component = \"ChatModel\"\n\t// ComponentOfEmbedding identifies embedding components.\n\tComponentOfEmbedding Component = \"Embedding\"\n\t// ComponentOfIndexer identifies indexer components.\n\tComponentOfIndexer Component = \"Indexer\"\n\t// ComponentOfRetriever identifies retriever components.\n\tComponentOfRetriever Component = \"Retriever\"\n\t// ComponentOfLoader identifies loader components.\n\tComponentOfLoader Component = \"Loader\"\n\t// ComponentOfTransformer identifies document transformer components.\n\tComponentOfTransformer Component = \"DocumentTransformer\"\n\t// ComponentOfTool identifies tool components.\n\tComponentOfTool Component = \"Tool\"\n)\n"
  },
  {
    "path": "compose/branch.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// GraphBranchCondition is the condition type for the branch.\ntype GraphBranchCondition[T any] func(ctx context.Context, in T) (endNode string, err error)\n\n// StreamGraphBranchCondition is the condition type for the stream branch.\ntype StreamGraphBranchCondition[T any] func(ctx context.Context, in *schema.StreamReader[T]) (endNode string, err error)\n\n// GraphMultiBranchCondition is the condition type for the multi choice branch.\ntype GraphMultiBranchCondition[T any] func(ctx context.Context, in T) (endNode map[string]bool, err error)\n\n// StreamGraphMultiBranchCondition is the condition type for the stream multi choice branch.\ntype StreamGraphMultiBranchCondition[T any] func(ctx context.Context, in *schema.StreamReader[T]) (endNodes map[string]bool, err error)\n\n// GraphBranch is the branch type for the graph.\n// It is used to determine the next node based on the condition.\ntype GraphBranch struct {\n\tinvoke    func(ctx context.Context, input any) (output []string, err error)\n\tcollect   func(ctx context.Context, input streamReader) (output []string, err error)\n\tinputType reflect.Type\n\t*genericHelper\n\tendNodes   map[string]bool\n\tidx        int // used to distinguish branches in parallel\n\tnoDataFlow bool\n}\n\n// GetEndNode returns the all end nodes of the branch.\nfunc (gb *GraphBranch) GetEndNode() map[string]bool {\n\treturn gb.endNodes\n}\n\nfunc newGraphBranch[T any](r *runnablePacker[T, []string, any], endNodes map[string]bool) *GraphBranch {\n\treturn &GraphBranch{\n\t\tinvoke: func(ctx context.Context, input any) (output []string, err error) {\n\t\t\tin, ok := input.(T)\n\t\t\tif !ok {\n\t\t\t\t// When a nil is passed as an 'any' type, its original type information is lost,\n\t\t\t\t// becoming an untyped nil. This would cause type assertions to fail.\n\t\t\t\t// So if the input is nil and the target type T is an interface, we need to explicitly create a nil of type T.\n\t\t\t\tif input == nil && generic.TypeOf[T]().Kind() == reflect.Interface {\n\t\t\t\t\tvar i T\n\t\t\t\t\tin = i\n\t\t\t\t} else {\n\t\t\t\t\tpanic(newUnexpectedInputTypeErr(generic.TypeOf[T](), reflect.TypeOf(input)))\n\t\t\t\t}\n\t\t\t}\n\t\t\treturn r.Invoke(ctx, in)\n\t\t},\n\t\tcollect: func(ctx context.Context, input streamReader) (output []string, err error) {\n\t\t\tin, ok := unpackStreamReader[T](input)\n\t\t\tif !ok {\n\t\t\t\tpanic(newUnexpectedInputTypeErr(generic.TypeOf[T](), input.getType()))\n\t\t\t}\n\t\t\treturn r.Collect(ctx, in)\n\t\t},\n\t\tinputType:     generic.TypeOf[T](),\n\t\tgenericHelper: newGenericHelper[T, T](),\n\t\tendNodes:      endNodes,\n\t}\n}\n\n// NewGraphMultiBranch creates a branch for graphs where a condition selects\n// multiple end nodes; only keys present in endNodes are allowed.\nfunc NewGraphMultiBranch[T any](condition GraphMultiBranchCondition[T], endNodes map[string]bool) *GraphBranch {\n\tcondRun := func(ctx context.Context, in T, opts ...any) ([]string, error) {\n\t\tends, err := condition(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tret := make([]string, 0, len(ends))\n\t\tfor end := range ends {\n\t\t\tif !endNodes[end] {\n\t\t\t\treturn nil, fmt.Errorf(\"branch invocation returns unintended end node: %s\", end)\n\t\t\t}\n\t\t\tret = append(ret, end)\n\t\t}\n\n\t\treturn ret, nil\n\t}\n\n\treturn newGraphBranch(newRunnablePacker(condRun, nil, nil, nil, false), endNodes)\n}\n\n// NewStreamGraphMultiBranch creates a streaming branch where a condition on\n// the input stream selects multiple end nodes.\nfunc NewStreamGraphMultiBranch[T any](condition StreamGraphMultiBranchCondition[T],\n\tendNodes map[string]bool) *GraphBranch {\n\n\tcondRun := func(ctx context.Context, in *schema.StreamReader[T], opts ...any) ([]string, error) {\n\t\tends, err := condition(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tret := make([]string, 0, len(ends))\n\t\tfor end := range ends {\n\t\t\tif !endNodes[end] {\n\t\t\t\treturn nil, fmt.Errorf(\"branch invocation returns unintended end node: %s\", end)\n\t\t\t}\n\t\t\tret = append(ret, end)\n\t\t}\n\t\treturn ret, nil\n\t}\n\n\treturn newGraphBranch(newRunnablePacker(nil, nil, condRun, nil, false), endNodes)\n}\n\n// NewGraphBranch creates a new graph branch.\n// It is used to determine the next node based on the condition.\n// e.g.\n//\n//\tcondition := func(ctx context.Context, in string) (string, error) {\n//\t\t// logic to determine the next node\n//\t\treturn \"next_node_key\", nil\n//\t}\n//\tendNodes := map[string]bool{\"path01\": true, \"path02\": true}\n//\tbranch := compose.NewGraphBranch(condition, endNodes)\n//\n//\tgraph.AddBranch(\"key_of_node_before_branch\", branch)\nfunc NewGraphBranch[T any](condition GraphBranchCondition[T], endNodes map[string]bool) *GraphBranch {\n\treturn NewGraphMultiBranch(func(ctx context.Context, in T) (endNode map[string]bool, err error) {\n\t\tret, err := condition(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn map[string]bool{ret: true}, nil\n\t}, endNodes)\n}\n\n// NewStreamGraphBranch creates a new stream graph branch.\n// It is used to determine the next node based on the condition of stream input.\n// e.g.\n//\n//\tcondition := func(ctx context.Context, in *schema.StreamReader[T]) (string, error) {\n//\t\t// logic to determine the next node.\n//\t\t// to use the feature of stream, you can use the first chunk to determine the next node.\n//\t\treturn \"next_node_key\", nil\n//\t}\n//\tendNodes := map[string]bool{\"path01\": true, \"path02\": true}\n//\tbranch := compose.NewStreamGraphBranch(condition, endNodes)\n//\n//\tgraph.AddBranch(\"key_of_node_before_branch\", branch)\nfunc NewStreamGraphBranch[T any](condition StreamGraphBranchCondition[T], endNodes map[string]bool) *GraphBranch {\n\treturn NewStreamGraphMultiBranch(func(ctx context.Context, in *schema.StreamReader[T]) (endNode map[string]bool, err error) {\n\t\tret, err := condition(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn map[string]bool{ret: true}, nil\n\t}, endNodes)\n}\n"
  },
  {
    "path": "compose/branch_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"io\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestMultiBranch(t *testing.T) {\n\tg := NewGraph[string, map[string]any]()\n\temptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil })\n\terr := g.AddLambdaNode(\"1\", emptyLambda, WithOutputKey(\"1\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", emptyLambda, WithOutputKey(\"2\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"3\", emptyLambda, WithOutputKey(\"3\"))\n\tassert.NoError(t, err)\n\n\terr = g.AddBranch(START, NewGraphMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) {\n\t\treturn map[string]bool{\"1\": true, \"2\": true}, nil\n\t}, map[string]bool{\"1\": true, \"2\": true, \"3\": true}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx)\n\tassert.NoError(t, err)\n\n\tresult, err := r.Invoke(ctx, \"start\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": \"start\",\n\t\t\"2\": \"start\",\n\t}, result)\n\n\tstreamResult, err := r.Stream(ctx, \"start\")\n\tassert.NoError(t, err)\n\tresult = map[string]any{}\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tfor k, v := range chunk {\n\t\t\tresult[k] = v\n\t\t}\n\t}\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": \"start\",\n\t\t\"2\": \"start\",\n\t}, result)\n}\n\nfunc TestStreamMultiBranch(t *testing.T) {\n\tg := NewGraph[string, map[string]any]()\n\temptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil })\n\terr := g.AddLambdaNode(\"1\", emptyLambda, WithOutputKey(\"1\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", emptyLambda, WithOutputKey(\"2\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"3\", emptyLambda, WithOutputKey(\"3\"))\n\tassert.NoError(t, err)\n\n\terr = g.AddBranch(START, NewStreamGraphMultiBranch(func(ctx context.Context, in *schema.StreamReader[string]) (endNode map[string]bool, err error) {\n\t\tin.Close()\n\t\treturn map[string]bool{\"1\": true, \"2\": true}, nil\n\t}, map[string]bool{\"1\": true, \"2\": true, \"3\": true}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx)\n\tassert.NoError(t, err)\n\n\tresult, err := r.Invoke(ctx, \"start\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": \"start\",\n\t\t\"2\": \"start\",\n\t}, result)\n\n\tstreamResult, err := r.Stream(ctx, \"start\")\n\tassert.NoError(t, err)\n\tresult = map[string]any{}\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tfor k, v := range chunk {\n\t\t\tresult[k] = v\n\t\t}\n\t}\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": \"start\",\n\t\t\"2\": \"start\",\n\t}, result)\n}\n"
  },
  {
    "path": "compose/chain.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/internal/gmap\"\n\t\"github.com/cloudwego/eino/internal/gslice\"\n)\n\n// NewChain create a chain with input/output type.\nfunc NewChain[I, O any](opts ...NewGraphOption) *Chain[I, O] {\n\tch := &Chain[I, O]{\n\t\tgg: NewGraph[I, O](opts...),\n\t}\n\n\tch.gg.cmp = ComponentOfChain\n\n\treturn ch\n}\n\n// Chain is a chain of components.\n// Chain nodes can be parallel / branch / sequence components.\n// Chain is designed to be used in a builder pattern (should Compile() before use).\n// And the interface is `Chain style`, you can use it like: `chain.AppendXX(...).AppendXX(...)`\n//\n// Normal usage:\n//  1. create a chain with input/output type: `chain := NewChain[inputType, outputType]()`\n//  2. add components to chainable list:\n//     2.1 add components: `chain.AppendChatTemplate(...).AppendChatModel(...).AppendToolsNode(...)`\n//     2.2 add parallel or branch node if needed: `chain.AppendParallel()`, `chain.AppendBranch()`\n//  3. compile: `r, err := c.Compile()`\n//  4. run:\n//     4.1 `one input & one output` use `r.Invoke(ctx, input)`\n//     4.2 `one input & multi output chunk` use `r.Stream(ctx, input)`\n//     4.3 `multi input chunk & one output` use `r.Collect(ctx, inputReader)`\n//     4.4 `multi input chunk & multi output chunk` use `r.Transform(ctx, inputReader)`\n//\n// Using in graph or other chain:\n// chain1 := NewChain[inputType, outputType]()\n// graph := NewGraph[](runTypePregel)\n// graph.AddGraph(\"key\", chain1) // chain is an AnyGraph implementation\n//\n// // or in another chain:\n// chain2 := NewChain[inputType, outputType]()\n// chain2.AppendGraph(chain1)\ntype Chain[I, O any] struct {\n\terr error\n\n\tgg *Graph[I, O]\n\n\tnodeIdx int\n\n\tpreNodeKeys []string\n\n\thasEnd bool\n}\n\n// ErrChainCompiled is returned when attempting to modify a chain after it has been compiled\nvar ErrChainCompiled = errors.New(\"chain has been compiled, cannot be modified\")\n\n// implements AnyGraph.\nfunc (c *Chain[I, O]) compile(ctx context.Context, option *graphCompileOptions) (*composableRunnable, error) {\n\tif err := c.addEndIfNeeded(); err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn c.gg.compile(ctx, option)\n}\n\n// addEndIfNeeded add END edge of the chain/graph.\n// only run once when compiling.\nfunc (c *Chain[I, O]) addEndIfNeeded() error {\n\tif c.hasEnd {\n\t\treturn nil\n\t}\n\n\tif c.err != nil {\n\t\treturn c.err\n\t}\n\n\tif len(c.preNodeKeys) == 0 {\n\t\treturn fmt.Errorf(\"pre node keys not set, number of nodes in chain= %d\", len(c.gg.nodes))\n\t}\n\n\tfor _, nodeKey := range c.preNodeKeys {\n\t\terr := c.gg.AddEdge(nodeKey, END)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\tc.hasEnd = true\n\n\treturn nil\n}\n\nfunc (c *Chain[I, O]) getGenericHelper() *genericHelper {\n\treturn newGenericHelper[I, O]()\n}\n\n// inputType returns the input type of the chain.\n// implements AnyGraph.\nfunc (c *Chain[I, O]) inputType() reflect.Type {\n\treturn generic.TypeOf[I]()\n}\n\n// outputType returns the output type of the chain.\n// implements AnyGraph.\nfunc (c *Chain[I, O]) outputType() reflect.Type {\n\treturn generic.TypeOf[O]()\n}\n\n// compositeType returns the composite type of the chain.\n// implements AnyGraph.\nfunc (c *Chain[I, O]) component() component {\n\treturn c.gg.component()\n}\n\n// Compile to a Runnable.\n// Runnable can be used directly.\n// e.g.\n//\n//\t\tchain := NewChain[string, string]()\n//\t\tr, err := chain.Compile()\n//\t\tif err != nil {}\n//\n//\t \tr.Invoke(ctx, input) // ping => pong\n//\t\tr.Stream(ctx, input) // ping => stream out\n//\t\tr.Collect(ctx, inputReader) // stream in => pong\n//\t\tr.Transform(ctx, inputReader) // stream in => stream out\nfunc (c *Chain[I, O]) Compile(ctx context.Context, opts ...GraphCompileOption) (Runnable[I, O], error) {\n\tif err := c.addEndIfNeeded(); err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn c.gg.Compile(ctx, opts...)\n}\n\n// AppendChatModel add a ChatModel node to the chain.\n// e.g.\n//\n//\tmodel, err := openai.NewChatModel(ctx, config)\n//\tif err != nil {...}\n//\tchain.AppendChatModel(model)\nfunc (c *Chain[I, O]) AppendChatModel(node model.BaseChatModel, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toChatModelNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendChatTemplate add a ChatTemplate node to the chain.\n// eg.\n//\n//\tchatTemplate, err := prompt.FromMessages(schema.FString, &schema.Message{\n//\t\tRole:    schema.System,\n//\t\tContent: \"You are acting as a {role}.\",\n//\t})\n//\n//\tchain.AppendChatTemplate(chatTemplate)\nfunc (c *Chain[I, O]) AppendChatTemplate(node prompt.ChatTemplate, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toChatTemplateNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendToolsNode add a ToolsNode node to the chain.\n// e.g.\n//\n//\ttoolsNode, err := tools.NewToolNode(ctx, &tools.ToolsNodeConfig{\n//\t\tTools: []tools.Tool{...},\n//\t})\n//\n//\tchain.AppendToolsNode(toolsNode)\nfunc (c *Chain[I, O]) AppendToolsNode(node *ToolsNode, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toToolsNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendDocumentTransformer add a DocumentTransformer node to the chain.\n// e.g.\n//\n//\tmarkdownSplitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderSplitterConfig{})\n//\n//\tchain.AppendDocumentTransformer(markdownSplitter)\nfunc (c *Chain[I, O]) AppendDocumentTransformer(node document.Transformer, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toDocumentTransformerNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendLambda add a Lambda node to the chain.\n// Lambda is a node that can be used to implement custom logic.\n// e.g.\n//\n//\tlambdaNode := compose.InvokableLambda(func(ctx context.Context, docs []*schema.Document) (string, error) {...})\n//\tchain.AppendLambda(lambdaNode)\n//\n// Note:\n// to create a Lambda node, you need to use `compose.AnyLambda` or `compose.InvokableLambda` or `compose.StreamableLambda` or `compose.TransformableLambda`.\n// if you want this node has real stream output, you need to use `compose.StreamableLambda` or `compose.TransformableLambda`, for example.\nfunc (c *Chain[I, O]) AppendLambda(node *Lambda, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toLambdaNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendEmbedding add a Embedding node to the chain.\n// e.g.\n//\n//\tembedder, err := openai.NewEmbedder(ctx, config)\n//\tif err != nil {...}\n//\tchain.AppendEmbedding(embedder)\nfunc (c *Chain[I, O]) AppendEmbedding(node embedding.Embedder, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toEmbeddingNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendRetriever add a Retriever node to the chain.\n// e.g.\n//\n//\t\tretriever, err := vectorstore.NewRetriever(ctx, config)\n//\t\tif err != nil {...}\n//\t\tchain.AppendRetriever(retriever)\n//\n//\t or using fornax knowledge as retriever:\n//\n//\t\tconfig := fornaxknowledge.Config{...}\n//\t\tretriever, err := fornaxknowledge.NewKnowledgeRetriever(ctx, config)\n//\t\tif err != nil {...}\n//\t\tchain.AppendRetriever(retriever)\nfunc (c *Chain[I, O]) AppendRetriever(node retriever.Retriever, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toRetrieverNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendLoader adds a Loader node to the chain.\n// e.g.\n//\n//\tloader, err := file.NewFileLoader(ctx, &file.FileLoaderConfig{})\n//\tif err != nil {...}\n//\tchain.AppendLoader(loader)\nfunc (c *Chain[I, O]) AppendLoader(node document.Loader, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toLoaderNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendIndexer add an Indexer node to the chain.\n// Indexer is a node that can store documents.\n// e.g.\n//\n//\tvectorStoreImpl, err := vikingdb.NewVectorStorer(ctx, vikingdbConfig) // in components/vectorstore/vikingdb/vectorstore.go\n//\tif err != nil {...}\n//\n//\tconfig := vectorstore.IndexerConfig{VectorStore: vectorStoreImpl}\n//\tindexer, err := vectorstore.NewIndexer(ctx, config)\n//\tif err != nil {...}\n//\n//\tchain.AppendIndexer(indexer)\nfunc (c *Chain[I, O]) AppendIndexer(node indexer.Indexer, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toIndexerNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendBranch add a conditional branch to chain.\n// Each branch within the ChainBranch can be an AnyGraph.\n// All branches should either lead to END, or converge to another node within the Chain.\n// e.g.\n//\n//\tcb := compose.NewChainBranch(conditionFunc)\n//\tcb.AddChatTemplate(\"chat_template_key_01\", chatTemplate)\n//\tcb.AddChatTemplate(\"chat_template_key_02\", chatTemplate2)\n//\tchain.AppendBranch(cb)\nfunc (c *Chain[I, O]) AppendBranch(b *ChainBranch) *Chain[I, O] {\n\tif b == nil {\n\t\tc.reportError(fmt.Errorf(\"append branch invalid, branch is nil\"))\n\t\treturn c\n\t}\n\n\tif b.err != nil {\n\t\tc.reportError(fmt.Errorf(\"append branch error: %w\", b.err))\n\t\treturn c\n\t}\n\n\tif len(b.key2BranchNode) == 0 {\n\t\tc.reportError(fmt.Errorf(\"append branch invalid, nodeList is empty\"))\n\t\treturn c\n\t}\n\n\tif len(b.key2BranchNode) == 1 {\n\t\tc.reportError(fmt.Errorf(\"append branch invalid, nodeList length = 1\"))\n\t\treturn c\n\t}\n\n\tvar startNode string\n\tif len(c.preNodeKeys) == 0 { // branch appended directly to START\n\t\tstartNode = START\n\t} else if len(c.preNodeKeys) == 1 {\n\t\tstartNode = c.preNodeKeys[0]\n\t} else {\n\t\tc.reportError(fmt.Errorf(\"append branch invalid, multiple previous nodes: %v \", c.preNodeKeys))\n\t\treturn c\n\t}\n\n\tprefix := c.nextNodeKey()\n\tkey2NodeKey := make(map[string]string, len(b.key2BranchNode))\n\n\tfor key := range b.key2BranchNode {\n\t\tnode := b.key2BranchNode[key]\n\n\t\tvar nodeKey string\n\n\t\tif node.Second != nil && node.Second.nodeOptions != nil && node.Second.nodeOptions.nodeKey != \"\" {\n\t\t\tnodeKey = node.Second.nodeOptions.nodeKey\n\t\t} else {\n\t\t\tnodeKey = fmt.Sprintf(\"%s_branch_%s\", prefix, key)\n\t\t}\n\n\t\tif err := c.gg.addNode(nodeKey, node.First, node.Second); err != nil {\n\t\t\tc.reportError(fmt.Errorf(\"add branch node[%s] to chain failed: %w\", nodeKey, err))\n\t\t\treturn c\n\t\t}\n\n\t\tkey2NodeKey[key] = nodeKey\n\t}\n\n\tgBranch := *b.internalBranch\n\n\tinvokeCon := func(ctx context.Context, in any) (endNode []string, err error) {\n\t\tends, err := b.internalBranch.invoke(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tnodeKeyEnds := make([]string, 0, len(ends))\n\t\tfor _, end := range ends {\n\t\t\tif nodeKey, ok := key2NodeKey[end]; !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"branch invocation returns unintended end node: %s\", end)\n\t\t\t} else {\n\t\t\t\tnodeKeyEnds = append(nodeKeyEnds, nodeKey)\n\t\t\t}\n\t\t}\n\n\t\treturn nodeKeyEnds, nil\n\t}\n\tgBranch.invoke = invokeCon\n\n\tcollectCon := func(ctx context.Context, sr streamReader) ([]string, error) {\n\t\tends, err := b.internalBranch.collect(ctx, sr)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tnodeKeyEnds := make([]string, 0, len(ends))\n\t\tfor _, end := range ends {\n\t\t\tif nodeKey, ok := key2NodeKey[end]; !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"branch invocation returns unintended end node: %s\", end)\n\t\t\t} else {\n\t\t\t\tnodeKeyEnds = append(nodeKeyEnds, nodeKey)\n\t\t\t}\n\t\t}\n\n\t\treturn nodeKeyEnds, nil\n\t}\n\tgBranch.collect = collectCon\n\n\tgBranch.endNodes = gslice.ToMap(gmap.Values(key2NodeKey), func(k string) (string, bool) {\n\t\treturn k, true\n\t})\n\n\tif err := c.gg.AddBranch(startNode, &gBranch); err != nil {\n\t\tc.reportError(fmt.Errorf(\"chain append branch failed: %w\", err))\n\t\treturn c\n\t}\n\n\tc.preNodeKeys = gmap.Values(key2NodeKey)\n\n\treturn c\n}\n\n// AppendParallel add a Parallel structure (multiple concurrent nodes) to the chain.\n// e.g.\n//\n//\tparallel := compose.NewParallel()\n//\tparallel.AddChatModel(\"openai\", model1) // => \"openai\": *schema.Message{}\n//\tparallel.AddChatModel(\"maas\", model2) // => \"maas\": *schema.Message{}\n//\n//\tchain.AppendParallel(parallel) // => multiple concurrent nodes are added to the Chain\n//\n//\tThe next node in the chain is either an END, or a node which accepts a map[string]any, where keys are `openai` `maas` as specified above.\nfunc (c *Chain[I, O]) AppendParallel(p *Parallel) *Chain[I, O] {\n\tif p == nil {\n\t\tc.reportError(fmt.Errorf(\"append parallel invalid, parallel is nil\"))\n\t\treturn c\n\t}\n\n\tif p.err != nil {\n\t\tc.reportError(fmt.Errorf(\"append parallel invalid, parallel error: %w\", p.err))\n\t\treturn c\n\t}\n\n\tif len(p.nodes) <= 1 {\n\t\tc.reportError(fmt.Errorf(\"append parallel invalid, not enough nodes, count = %d\", len(p.nodes)))\n\t\treturn c\n\t}\n\n\tvar startNode string\n\tif len(c.preNodeKeys) == 0 { // parallel appended directly to START\n\t\tstartNode = START\n\t} else if len(c.preNodeKeys) == 1 {\n\t\tstartNode = c.preNodeKeys[0]\n\t} else {\n\t\tc.reportError(fmt.Errorf(\"append parallel invalid, multiple previous nodes: %v \", c.preNodeKeys))\n\t\treturn c\n\t}\n\n\tprefix := c.nextNodeKey()\n\tvar nodeKeys []string\n\n\tfor i := range p.nodes {\n\t\tnode := p.nodes[i]\n\n\t\tvar nodeKey string\n\t\tif node.Second != nil && node.Second.nodeOptions != nil && node.Second.nodeOptions.nodeKey != \"\" {\n\t\t\tnodeKey = node.Second.nodeOptions.nodeKey\n\t\t} else {\n\t\t\tnodeKey = fmt.Sprintf(\"%s_parallel_%d\", prefix, i)\n\t\t}\n\n\t\tif err := c.gg.addNode(nodeKey, node.First, node.Second); err != nil {\n\t\t\tc.reportError(fmt.Errorf(\"add parallel node to chain failed, key=%s, err: %w\", nodeKey, err))\n\t\t\treturn c\n\t\t}\n\n\t\tif err := c.gg.AddEdge(startNode, nodeKey); err != nil {\n\t\t\tc.reportError(fmt.Errorf(\"add parallel edge failed, from=%s, to=%s, err: %w\", startNode, nodeKey, err))\n\t\t\treturn c\n\t\t}\n\n\t\tnodeKeys = append(nodeKeys, nodeKey)\n\t}\n\n\tc.preNodeKeys = nodeKeys\n\n\treturn c\n}\n\n// AppendGraph add a AnyGraph node to the chain.\n// AnyGraph can be a chain or a graph.\n// e.g.\n//\n//\tgraph := compose.NewGraph[string, string]()\n//\tchain.AppendGraph(graph)\nfunc (c *Chain[I, O]) AppendGraph(node AnyGraph, opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toAnyGraphNode(node, opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// AppendPassthrough add a Passthrough node to the chain.\n// Could be used to connect multiple ChainBranch or Parallel.\n// e.g.\n//\n//\tchain.AppendPassthrough()\nfunc (c *Chain[I, O]) AppendPassthrough(opts ...GraphAddNodeOpt) *Chain[I, O] {\n\tgNode, options := toPassthroughNode(opts...)\n\tc.addNode(gNode, options)\n\treturn c\n}\n\n// nextIdx.\n// get the next idx for the chain.\n// chain key is: node_idx => eg: node_0 => represent the first node of the chain (idx start from 0)\n// if has parallel: node_idx_parallel_idx => eg: node_0_parallel_1 => represent the first node of the chain, and is a parallel node, and the second node of the parallel\n// if has branch: node_idx_branch_key => eg: node_1_branch_customkey => represent the second node of the chain, and is a branch node, and the 'customkey' is the key of the branch\nfunc (c *Chain[I, O]) nextNodeKey() string {\n\tidx := c.nodeIdx\n\tc.nodeIdx++\n\treturn fmt.Sprintf(\"node_%d\", idx)\n}\n\n// reportError.\n// save the first error in the chain.\nfunc (c *Chain[I, O]) reportError(err error) {\n\tif c.err == nil {\n\t\tc.err = err\n\t}\n}\n\n// addNode.\n// add a node to the chain.\nfunc (c *Chain[I, O]) addNode(node *graphNode, options *graphAddNodeOpts) {\n\tif c.err != nil {\n\t\treturn\n\t}\n\n\tif c.gg.compiled {\n\t\tc.reportError(ErrChainCompiled)\n\t\treturn\n\t}\n\n\tif node == nil {\n\t\tc.reportError(fmt.Errorf(\"chain add node invalid, node is nil\"))\n\t\treturn\n\t}\n\n\tnodeKey := options.nodeOptions.nodeKey\n\tdefaultNodeKey := c.nextNodeKey()\n\tif nodeKey == \"\" {\n\t\tnodeKey = defaultNodeKey\n\t}\n\n\terr := c.gg.addNode(nodeKey, node, options)\n\tif err != nil {\n\t\tc.reportError(err)\n\t\treturn\n\t}\n\n\tif len(c.preNodeKeys) == 0 {\n\t\tc.preNodeKeys = append(c.preNodeKeys, START)\n\t}\n\n\tfor _, preNodeKey := range c.preNodeKeys {\n\t\te := c.gg.AddEdge(preNodeKey, nodeKey)\n\t\tif e != nil {\n\t\t\tc.reportError(e)\n\t\t\treturn\n\t\t}\n\t}\n\n\tc.preNodeKeys = []string{nodeKey}\n}\n"
  },
  {
    "path": "compose/chain_branch.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype nodeOptionsPair generic.Pair[*graphNode, *graphAddNodeOpts]\n\n// ChainBranch represents a conditional branch in a chain of operations.\n// It allows for dynamic routing of execution based on a condition.\n// All branches within ChainBranch are expected to either end the Chain, or converge to another node in the Chain.\ntype ChainBranch struct {\n\tinternalBranch *GraphBranch\n\tkey2BranchNode map[string]nodeOptionsPair\n\terr            error\n}\n\n// NewChainMultiBranch creates a chain branch where a condition selects\n// multiple end nodes to route execution.\nfunc NewChainMultiBranch[T any](cond GraphMultiBranchCondition[T]) *ChainBranch {\n\tinvokeCond := func(ctx context.Context, in T, opts ...any) (endNodes []string, err error) {\n\t\tends, err := cond(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tendNodes = make([]string, 0, len(ends))\n\t\tfor end := range ends {\n\t\t\tendNodes = append(endNodes, end)\n\t\t}\n\t\treturn endNodes, nil\n\t}\n\n\treturn &ChainBranch{\n\t\tkey2BranchNode: make(map[string]nodeOptionsPair),\n\t\tinternalBranch: newGraphBranch(newRunnablePacker(invokeCond, nil, nil, nil, false), nil),\n\t}\n}\n\n// NewStreamChainMultiBranch creates a chain branch that selects multiple end\n// nodes based on a condition evaluated on the input stream.\nfunc NewStreamChainMultiBranch[T any](cond StreamGraphMultiBranchCondition[T]) *ChainBranch {\n\tcollectCon := func(ctx context.Context, in *schema.StreamReader[T], opts ...any) (endNodes []string, err error) {\n\t\tends, err := cond(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tendNodes = make([]string, 0, len(ends))\n\t\tfor end := range ends {\n\t\t\tendNodes = append(endNodes, end)\n\t\t}\n\t\treturn endNodes, nil\n\t}\n\n\treturn &ChainBranch{\n\t\tkey2BranchNode: make(map[string]nodeOptionsPair),\n\t\tinternalBranch: newGraphBranch(newRunnablePacker(nil, nil, collectCon, nil, false), nil),\n\t}\n}\n\n// NewChainBranch creates a new ChainBranch instance based on a given condition.\n// It takes a generic type T and a GraphBranchCondition function for that type.\n// The returned ChainBranch will have an empty key2BranchNode map and a condition function\n// that wraps the provided cond to handle type assertions and error checking.\n// eg.\n//\n//\tcondition := func(ctx context.Context, in string, opts ...any) (endNode string, err error) {\n//\t\t// logic to determine the next node\n//\t\treturn \"some_next_node_key\", nil\n//\t}\n//\n//\tcb := NewChainBranch[string](condition)\n//\tcb.AddPassthrough(\"next_node_key_01\", xxx) // node in branch, represent one path of branch\n//\tcb.AddPassthrough(\"next_node_key_02\", xxx) // node in branch\nfunc NewChainBranch[T any](cond GraphBranchCondition[T]) *ChainBranch {\n\treturn NewChainMultiBranch(func(ctx context.Context, in T) (endNode map[string]bool, err error) {\n\t\tret, err := cond(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn map[string]bool{ret: true}, nil\n\t})\n}\n\n// NewStreamChainBranch creates a new ChainBranch instance based on a given stream condition.\n// It takes a generic type T and a StreamGraphBranchCondition function for that type.\n// The returned ChainBranch will have an empty key2BranchNode map and a condition function\n// that wraps the provided cond to handle type assertions and error checking.\n// eg.\n//\n//\tcondition := func(ctx context.Context, in *schema.StreamReader[string], opts ...any) (endNode string, err error) {\n//\t\t// logic to determine the next node, you can read the stream and make a decision.\n//\t\t// to save time, usually read the first chunk of stream, then make a decision which path to go.\n//\t\treturn \"some_next_node_key\", nil\n//\t}\n//\n//\tcb := NewStreamChainBranch[string](condition)\nfunc NewStreamChainBranch[T any](cond StreamGraphBranchCondition[T]) *ChainBranch {\n\treturn NewStreamChainMultiBranch(func(ctx context.Context, in *schema.StreamReader[T]) (endNodes map[string]bool, err error) {\n\t\tret, err := cond(ctx, in)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn map[string]bool{ret: true}, nil\n\t})\n}\n\n// AddChatModel adds a ChatModel node to the branch.\n// eg.\n//\n//\tchatModel01, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{\n//\t\tModel: \"gpt-4o\",\n//\t})\n//\tchatModel02, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{\n//\t\tModel: \"gpt-4o-mini\",\n//\t})\n//\tcb.AddChatModel(\"chat_model_key_01\", chatModel01)\n//\tcb.AddChatModel(\"chat_model_key_02\", chatModel02)\nfunc (cb *ChainBranch) AddChatModel(key string, node model.BaseChatModel, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toChatModelNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddChatTemplate adds a ChatTemplate node to the branch.\n// eg.\n//\n//\tchatTemplate, err := prompt.FromMessages(schema.FString, &schema.Message{\n//\t\tRole:    schema.System,\n//\t\tContent: \"You are acting as a {role}.\",\n//\t})\n//\n//\tcb.AddChatTemplate(\"chat_template_key_01\", chatTemplate)\n//\n//\tchatTemplate2, err := prompt.FromMessages(schema.FString, &schema.Message{\n//\t\tRole:    schema.System,\n//\t\tContent: \"You are acting as a {role}, you are not allowed to chat in other topics.\",\n//\t})\n//\n//\tcb.AddChatTemplate(\"chat_template_key_02\", chatTemplate2)\nfunc (cb *ChainBranch) AddChatTemplate(key string, node prompt.ChatTemplate, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toChatTemplateNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddToolsNode adds a ToolsNode to the branch.\n// eg.\n//\n//\ttoolsNode, err := tools.NewToolNode(ctx, &tools.ToolsNodeConfig{\n//\t\tTools: []tools.Tool{...},\n//\t})\n//\n//\tcb.AddToolsNode(\"tools_node_key\", toolsNode)\nfunc (cb *ChainBranch) AddToolsNode(key string, node *ToolsNode, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toToolsNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddLambda adds a Lambda node to the branch.\n// eg.\n//\n//\tlambdaFunc := func(ctx context.Context, in string, opts ...any) (out string, err error) {\n//\t\t// logic to process the input\n//\t\treturn \"processed_output\", nil\n//\t}\n//\n//\tcb.AddLambda(\"lambda_node_key\", compose.InvokeLambda(lambdaFunc))\nfunc (cb *ChainBranch) AddLambda(key string, node *Lambda, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toLambdaNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddEmbedding adds an Embedding node to the branch.\n// eg.\n//\n//\tembeddingNode, err := openai.NewEmbedder(ctx, &openai.EmbeddingConfig{\n//\t\tModel: \"text-embedding-3-small\",\n//\t})\n//\n//\tcb.AddEmbedding(\"embedding_node_key\", embeddingNode)\nfunc (cb *ChainBranch) AddEmbedding(key string, node embedding.Embedder, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toEmbeddingNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddRetriever adds a Retriever node to the branch.\n// eg.\n//\n//\tretriever, err := volc_vikingdb.NewRetriever(ctx, &volc_vikingdb.RetrieverConfig{\n//\t\tCollection: \"my_collection\",\n//\t})\n//\n//\tcb.AddRetriever(\"retriever_node_key\", retriever)\nfunc (cb *ChainBranch) AddRetriever(key string, node retriever.Retriever, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toRetrieverNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddLoader adds a Loader node to the branch.\n// eg.\n//\n//\tpdfParser, err := pdf.NewPDFParser()\n//\tloader, err := file.NewFileLoader(ctx, &file.FileLoaderConfig{\n//\t\tParser: pdfParser,\n//\t})\n//\n//\tcb.AddLoader(\"loader_node_key\", loader)\nfunc (cb *ChainBranch) AddLoader(key string, node document.Loader, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toLoaderNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddIndexer adds an Indexer node to the branch.\n// eg.\n//\n//\tindexer, err := volc_vikingdb.NewIndexer(ctx, &volc_vikingdb.IndexerConfig{\n//\t\tCollection: \"my_collection\",\n//\t})\n//\n//\tcb.AddIndexer(\"indexer_node_key\", indexer)\nfunc (cb *ChainBranch) AddIndexer(key string, node indexer.Indexer, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toIndexerNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddDocumentTransformer adds an Document Transformer node to the branch.\n// eg.\n//\n//\tmarkdownSplitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderSplitterConfig{})\n//\n//\tcb.AddDocumentTransformer(\"document_transformer_node_key\", markdownSplitter)\nfunc (cb *ChainBranch) AddDocumentTransformer(key string, node document.Transformer, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toDocumentTransformerNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddGraph adds a generic Graph node to the branch.\n// eg.\n//\n//\tgraph, err := compose.NewGraph[string, string]()\n//\n//\tcb.AddGraph(\"graph_node_key\", graph)\nfunc (cb *ChainBranch) AddGraph(key string, node AnyGraph, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toAnyGraphNode(node, opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\n// AddPassthrough adds a Passthrough node to the branch.\n// eg.\n//\n//\tcb.AddPassthrough(\"passthrough_node_key\")\nfunc (cb *ChainBranch) AddPassthrough(key string, opts ...GraphAddNodeOpt) *ChainBranch {\n\tgNode, options := toPassthroughNode(opts...)\n\treturn cb.addNode(key, gNode, options)\n}\n\nfunc (cb *ChainBranch) addNode(key string, node *graphNode, options *graphAddNodeOpts) *ChainBranch {\n\tif cb.err != nil {\n\t\treturn cb\n\t}\n\n\tif cb.key2BranchNode == nil {\n\t\tcb.key2BranchNode = make(map[string]nodeOptionsPair)\n\t}\n\n\t_, ok := cb.key2BranchNode[key]\n\tif ok {\n\t\tcb.err = fmt.Errorf(\"chain branch add node, duplicate branch node key= %s\", key)\n\t\treturn cb\n\t}\n\n\tcb.key2BranchNode[key] = nodeOptionsPair{node, options}\n\n\treturn cb\n}\n"
  },
  {
    "path": "compose/chain_branch_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"strconv\"\n\t\"strings\"\n\t\"testing\"\n\t\"unicode/utf8\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestChainBranch(t *testing.T) {\n\tcond := func(ctx context.Context, input string) (key string, err error) {\n\t\tswitch input {\n\t\tcase \"one\":\n\t\t\treturn \"one_key\", nil\n\t\tcase \"two\":\n\t\t\treturn \"two_key\", nil\n\t\tcase \"three\":\n\t\t\treturn \"three_key\", nil\n\t\tdefault:\n\t\t\treturn \"\", fmt.Errorf(\"invalid input= %s\", input)\n\t\t}\n\t}\n\n\tt.Run(\"nested chain\", func(t *testing.T) {\n\t\tinner := NewChain[string, string]()\n\t\tinner.AppendBranch(NewChainBranch(cond).\n\t\t\tAddLambda(\"one_key\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\t\treturn in + in, nil\n\t\t\t})).\n\t\t\tAddLambda(\"two_key\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\t\treturn in + in + in, nil\n\t\t\t})))\n\t\tinner.AppendParallel(NewParallel().\n\t\t\tAddLambda(\"one_key\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\t\treturn in + in, nil\n\t\t\t})).\n\t\t\tAddLambda(\"two_key\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\t\treturn in + in + in, nil\n\t\t\t})))\n\n\t\touter := NewChain[string, string]()\n\t\touter.AppendGraph(inner)\n\t\t_, err := outer.Compile(context.Background())\n\t\tassert.Error(t, err)\n\t})\n\n\tt.Run(\"bad param\", func(t *testing.T) {\n\t\tc := NewChain[string, string]()\n\t\tc.AppendBranch(nil)\n\t\tassert.NotNil(t, c.err)\n\n\t\tc = NewChain[string, string]()\n\t\tc.AppendBranch(NewChainBranch[string](nil))\n\t\tassert.NotNil(t, c.err)\n\n\t\tc = NewChain[string, string]()\n\t\tc.AppendBranch(NewChainBranch(cond).AddChatTemplate(\"template\", prompt.FromMessages(schema.FString, schema.SystemMessage(\"hello\"))))\n\t\tassert.NotNil(t, c.err)\n\n\t\tc = NewChain[string, string]()\n\t\tc.AppendBranch(NewChainBranch(cond).AddChatTemplate(\"1\", prompt.FromMessages(schema.FString)).AddChatTemplate(\"1\", prompt.FromMessages(schema.FString)))\n\t\tassert.NotNil(t, c.err)\n\t})\n\n\tt.Run(\"different Node types in branch\", func(t *testing.T) {\n\t\tc := NewChain[string, string]()\n\t\tc.AppendBranch(NewChainBranch(cond).\n\t\t\tAddChatTemplate(\"t\", prompt.FromMessages(schema.FString)).\n\t\t\tAddGraph(\"c\", NewChain[string, string]()))\n\t\tassert.NotNil(t, c.err)\n\t})\n\n\tt.Run(\"type mismatch\", func(t *testing.T) {\n\t\tc := NewChain[int, string]()\n\t\tc.AppendBranch(NewChainBranch(cond).\n\t\t\tAddLambda(\"one_key\", InvokableLambda(func(ctx context.Context, in int) (output string, err error) {\n\t\t\t\treturn strconv.Itoa(in), nil\n\t\t\t})).\n\t\t\tAddLambda(\"two_key\", InvokableLambda(func(ctx context.Context, in int) (output string, err error) {\n\t\t\t\treturn strconv.Itoa(in), nil\n\t\t\t})))\n\t\t_, err := c.Compile(context.Background())\n\t\tassert.NotNil(t, err)\n\t})\n\n\tt.Run(\"invoke\", func(t *testing.T) {\n\t\tc := NewChain[string, string]()\n\t\tc.AppendBranch(NewChainBranch(cond).\n\t\t\tAddLambda(\"one_key\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\t\treturn in + in, nil\n\t\t\t})).\n\t\t\tAddLambda(\"two_key\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\t\treturn in + in + in, nil\n\t\t\t})))\n\t\tc.AppendLambda(InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in + in, nil\n\t\t}))\n\t\tassert.Nil(t, c.err)\n\t\tcompiledChain, err := c.Compile(context.Background())\n\t\tassert.Nil(t, err)\n\n\t\tout, err := compiledChain.Invoke(context.Background(), \"two\")\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, \"twotwotwotwotwotwo\", out)\n\n\t\t_, err = compiledChain.Invoke(context.Background(), \"three\")\n\t\tassert.NotNil(t, err)\n\n\t\t_, err = compiledChain.Invoke(context.Background(), \"four\")\n\t\tassert.NotNil(t, err)\n\t})\n\n\tt.Run(\"fake stream\", func(t *testing.T) {\n\t\tc := NewChain[string, string]()\n\t\tc.AppendLambda(StreamableLambda(func(ctx context.Context, in string) (output *schema.StreamReader[string], err error) {\n\t\t\tsr, sw := schema.Pipe[string](utf8.RuneCountInString(in))\n\n\t\t\tgo func() {\n\t\t\t\tfor _, field := range strings.Fields(in) {\n\t\t\t\t\tsw.Send(field, nil)\n\t\t\t\t}\n\t\t\t\tsw.Close()\n\t\t\t}()\n\n\t\t\treturn sr, nil\n\t\t}))\n\t\tc.AppendBranch(NewChainBranch[string](cond).AddLambda(\"one_key\", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) {\n\t\t\tdefer in.Close()\n\t\t\tfor {\n\t\t\t\tv, err := in.Recv()\n\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn \"\", err\n\t\t\t\t}\n\n\t\t\t\toutput += v\n\t\t\t}\n\n\t\t\treturn output + output, nil\n\t\t})).\n\t\t\tAddLambda(\"two_key\", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) {\n\t\t\t\tdefer in.Close()\n\t\t\t\tfor {\n\t\t\t\t\tv, err := in.Recv()\n\t\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn \"\", err\n\t\t\t\t\t}\n\n\t\t\t\t\toutput += v\n\t\t\t\t}\n\n\t\t\t\treturn output + output + output, nil\n\t\t\t})))\n\n\t\tassert.Nil(t, c.err)\n\t\tcompiledChain, err := c.Compile(context.Background())\n\t\tassert.Nil(t, err)\n\n\t\tout, err := compiledChain.Invoke(context.Background(), \"one\")\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, \"oneone\", out)\n\t})\n\n\tt.Run(\"real stream\", func(t *testing.T) {\n\t\tstreamCon := func(ctx context.Context, sr *schema.StreamReader[string]) (key string, err error) {\n\t\t\tmsg, err := sr.Recv()\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", err\n\t\t\t}\n\t\t\tdefer sr.Close()\n\n\t\t\tswitch msg {\n\t\t\tcase \"one\":\n\t\t\t\treturn \"one_key\", nil\n\t\t\tcase \"two\":\n\t\t\t\treturn \"two_key\", nil\n\t\t\tcase \"three\":\n\t\t\t\treturn \"three_key\", nil\n\t\t\tdefault:\n\t\t\t\treturn \"\", fmt.Errorf(\"invalid input= %s\", msg)\n\t\t\t}\n\t\t}\n\n\t\tc := NewChain[string, string]()\n\t\tc.AppendLambda(StreamableLambda(func(ctx context.Context, in string) (output *schema.StreamReader[string], err error) {\n\t\t\tsr, sw := schema.Pipe[string](utf8.RuneCountInString(in))\n\n\t\t\tgo func() {\n\t\t\t\tfor _, field := range strings.Fields(in) {\n\t\t\t\t\tsw.Send(field, nil)\n\t\t\t\t}\n\t\t\t\tsw.Close()\n\t\t\t}()\n\n\t\t\treturn sr, nil\n\t\t}))\n\t\tc.AppendBranch(NewStreamChainBranch(streamCon).AddLambda(\"one_key\", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) {\n\t\t\tdefer in.Close()\n\t\t\tfor {\n\t\t\t\tv, err := in.Recv()\n\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn \"\", err\n\t\t\t\t}\n\n\t\t\t\toutput += v\n\t\t\t}\n\n\t\t\treturn output + output, nil\n\t\t})).\n\t\t\tAddLambda(\"two_key\", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) {\n\t\t\t\tdefer in.Close()\n\t\t\t\tfor {\n\t\t\t\t\tv, err := in.Recv()\n\t\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn \"\", err\n\t\t\t\t\t}\n\n\t\t\t\t\toutput += v\n\t\t\t\t}\n\n\t\t\t\treturn output + output + output, nil\n\t\t\t})))\n\n\t\tassert.Nil(t, c.err)\n\t\tcompiledChain, err := c.Compile(context.Background())\n\t\tassert.Nil(t, err)\n\n\t\tout, err := compiledChain.Stream(context.Background(), \"one size fit all\")\n\t\tassert.Nil(t, err)\n\t\tconcat, err := concatStreamReader(out)\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, \"onesizefitallonesizefitall\", concat)\n\t})\n}\n\nfunc TestChainMultiBranch(t *testing.T) {\n\temptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil })\n\n\tctx := context.Background()\n\tr, err := NewChain[string, map[string]any]().\n\t\tAppendBranch(NewChainMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) {\n\t\t\treturn map[string]bool{\"1\": true, \"2\": true}, nil\n\t\t}).AddLambda(\"1\", emptyLambda, WithOutputKey(\"1\")).AddLambda(\"2\", emptyLambda, WithOutputKey(\"2\")).AddLambda(\"3\", emptyLambda, WithOutputKey(\"3\"))).\n\t\tCompile(ctx)\n\tassert.Nil(t, err)\n\n\tresult, err := r.Invoke(ctx, \"start\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": \"start\",\n\t\t\"2\": \"start\",\n\t}, result)\n\n\tstreamResult, err := r.Stream(ctx, \"start\")\n\tassert.NoError(t, err)\n\tresult = map[string]any{}\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tfor k, v := range chunk {\n\t\t\tresult[k] = v\n\t\t}\n\t}\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": \"start\",\n\t\t\"2\": \"start\",\n\t}, result)\n}\n\nfunc TestStreamChainMultiBranch(t *testing.T) {\n\temptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil })\n\n\tctx := context.Background()\n\tr, err := NewChain[string, map[string]any]().\n\t\tAppendBranch(NewStreamChainMultiBranch(func(ctx context.Context, in *schema.StreamReader[string]) (endNode map[string]bool, err error) {\n\t\t\treturn map[string]bool{\"1\": true, \"2\": true}, nil\n\t\t}).AddLambda(\"1\", emptyLambda, WithOutputKey(\"1\")).AddLambda(\"2\", emptyLambda, WithOutputKey(\"2\")).AddLambda(\"3\", emptyLambda, WithOutputKey(\"3\"))).\n\t\tCompile(ctx)\n\tassert.Nil(t, err)\n\n\tresult, err := r.Invoke(ctx, \"start\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": \"start\",\n\t\t\"2\": \"start\",\n\t}, result)\n\n\tstreamResult, err := r.Stream(ctx, \"start\")\n\tassert.NoError(t, err)\n\tresult = map[string]any{}\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tfor k, v := range chunk {\n\t\t\tresult[k] = v\n\t\t}\n\t}\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": \"start\",\n\t\t\"2\": \"start\",\n\t}, result)\n}\n"
  },
  {
    "path": "compose/chain_parallel.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n)\n\n// NewParallel creates a new parallel type.\n// it is useful when you want to run multiple nodes in parallel in a chain.\nfunc NewParallel() *Parallel {\n\treturn &Parallel{\n\t\toutputKeys: make(map[string]bool),\n\t}\n}\n\n// Parallel run multiple nodes in parallel\n//\n// use `NewParallel()` to create a new parallel type\n// Example:\n//\n//\tparallel := NewParallel()\n//\tparallel.AddChatModel(\"output_key01\", chat01)\n//\tparallel.AddChatModel(\"output_key01\", chat02)\n//\n//\tchain := NewChain[any,any]()\n//\tchain.AppendParallel(parallel)\ntype Parallel struct {\n\tnodes      []nodeOptionsPair\n\toutputKeys map[string]bool\n\terr        error\n}\n\n// AddChatModel adds a chat model to the parallel.\n// eg.\n//\n//\tchatModel01, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{\n//\t\tModel: \"gpt-4o\",\n//\t})\n//\n//\tchatModel02, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{\n//\t\tModel: \"gpt-4o\",\n//\t})\n//\n//\tp.AddChatModel(\"output_key01\", chatModel01)\n//\tp.AddChatModel(\"output_key02\", chatModel02)\nfunc (p *Parallel) AddChatModel(outputKey string, node model.BaseChatModel, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toChatModelNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddChatTemplate adds a chat template to the parallel.\n// eg.\n//\n//\tchatTemplate01, err := prompt.FromMessages(schema.FString, &schema.Message{\n//\t\tRole:    schema.System,\n//\t\tContent: \"You are acting as a {role}.\",\n//\t})\n//\n//\tp.AddChatTemplate(\"output_key01\", chatTemplate01)\nfunc (p *Parallel) AddChatTemplate(outputKey string, node prompt.ChatTemplate, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toChatTemplateNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddToolsNode adds a tools node to the parallel.\n// eg.\n//\n//\ttoolsNode, err := compose.NewToolNode(ctx, &compose.ToolsNodeConfig{\n//\t\tTools: []tool.BaseTool{...},\n//\t})\n//\n//\tp.AddToolsNode(\"output_key01\", toolsNode)\nfunc (p *Parallel) AddToolsNode(outputKey string, node *ToolsNode, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toToolsNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddLambda adds a lambda node to the parallel.\n// eg.\n//\n//\tlambdaFunc := func(ctx context.Context, input *schema.Message) ([]*schema.Message, error) {\n//\t\treturn []*schema.Message{input}, nil\n//\t}\n//\n//\tp.AddLambda(\"output_key01\", compose.InvokeLambda(lambdaFunc))\nfunc (p *Parallel) AddLambda(outputKey string, node *Lambda, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toLambdaNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddEmbedding adds an embedding node to the parallel.\n// eg.\n//\n//\tembeddingNode, err := openai.NewEmbedder(ctx, &openai.EmbeddingConfig{\n//\t\tModel: \"text-embedding-3-small\",\n//\t})\n//\n//\tp.AddEmbedding(\"output_key01\", embeddingNode)\nfunc (p *Parallel) AddEmbedding(outputKey string, node embedding.Embedder, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toEmbeddingNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddRetriever adds a retriever node to the parallel.\n// eg.\n//\n// retriever, err := vikingdb.NewRetriever(ctx, &vikingdb.RetrieverConfig{})\n//\n//\tp.AddRetriever(\"output_key01\", retriever)\nfunc (p *Parallel) AddRetriever(outputKey string, node retriever.Retriever, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toRetrieverNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddLoader adds a loader node to the parallel.\n// eg.\n//\n//\tloader, err := file.NewLoader(ctx, &file.LoaderConfig{})\n//\n//\tp.AddLoader(\"output_key01\", loader)\nfunc (p *Parallel) AddLoader(outputKey string, node document.Loader, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toLoaderNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddIndexer adds an indexer node to the parallel.\n// eg.\n//\n//\tindexer, err := volc_vikingdb.NewIndexer(ctx, &volc_vikingdb.IndexerConfig{\n//\t\tCollection: \"my_collection\",\n//\t})\n//\n//\tp.AddIndexer(\"output_key01\", indexer)\nfunc (p *Parallel) AddIndexer(outputKey string, node indexer.Indexer, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toIndexerNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddDocumentTransformer adds an Document Transformer node to the parallel.\n// eg.\n//\n//\tmarkdownSplitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderSplitterConfig{})\n//\n//\tp.AddDocumentTransformer(\"output_key01\", markdownSplitter)\nfunc (p *Parallel) AddDocumentTransformer(outputKey string, node document.Transformer, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toDocumentTransformerNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddGraph adds a graph node to the parallel.\n// It is useful when you want to use a graph or a chain as a node in the parallel.\n// eg.\n//\n//\tgraph, err := compose.NewChain[any,any]()\n//\n//\tp.AddGraph(\"output_key01\", graph)\nfunc (p *Parallel) AddGraph(outputKey string, node AnyGraph, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toAnyGraphNode(node, append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\n// AddPassthrough adds a passthrough node to the parallel.\n// eg.\n//\n//\tp.AddPassthrough(\"output_key01\")\nfunc (p *Parallel) AddPassthrough(outputKey string, opts ...GraphAddNodeOpt) *Parallel {\n\tgNode, options := toPassthroughNode(append(opts, WithOutputKey(outputKey))...)\n\treturn p.addNode(outputKey, gNode, options)\n}\n\nfunc (p *Parallel) addNode(outputKey string, node *graphNode, options *graphAddNodeOpts) *Parallel {\n\tif p.err != nil {\n\t\treturn p\n\t}\n\n\tif node == nil {\n\t\tp.err = fmt.Errorf(\"chain parallel add node invalid, node is nil\")\n\t\treturn p\n\t}\n\n\tif p.outputKeys == nil {\n\t\tp.outputKeys = make(map[string]bool)\n\t}\n\n\tif _, ok := p.outputKeys[outputKey]; ok {\n\t\tp.err = fmt.Errorf(\"parallel add node err, duplicate output key= %s\", outputKey)\n\t\treturn p\n\t}\n\n\tif node.nodeInfo == nil {\n\t\tp.err = fmt.Errorf(\"chain parallel add node invalid, nodeInfo is nil\")\n\t\treturn p\n\t}\n\n\tnode.nodeInfo.outputKey = outputKey\n\tp.nodes = append(p.nodes, nodeOptionsPair{node, options})\n\tp.outputKeys[outputKey] = true\n\treturn p\n}\n"
  },
  {
    "path": "compose/chain_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"math/rand\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/internal/mock/components/document\"\n\t\"github.com/cloudwego/eino/internal/mock/components/embedding\"\n\t\"github.com/cloudwego/eino/internal/mock/components/indexer\"\n\t\"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/internal/mock/components/retriever\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestChain(t *testing.T) {\n\n\tcm := &mockIntentChatModel{}\n\n\t// 构建 branch\n\tbranchCond := func(ctx context.Context, input map[string]any) (string, error) {\n\t\tif rand.Intn(2) == 1 {\n\t\t\treturn \"b1\", nil\n\t\t}\n\t\treturn \"b2\", nil\n\t}\n\n\tb1 := InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\tt.Log(\"hello in branch lambda 01\")\n\t\tkvs[\"role\"] = \"cat\"\n\t\treturn kvs, nil\n\t})\n\tb2 := InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\tt.Log(\"hello in branch lambda 02\")\n\t\tkvs[\"role\"] = \"dog\"\n\t\treturn kvs, nil\n\t})\n\n\t// 并发节点\n\tparallel := NewParallel()\n\tparallel.\n\t\tAddLambda(\"role\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\t// may be change role to others by input kvs, for example (dentist/doctor...)\n\t\t\trole := kvs[\"role\"]\n\t\t\tif role.(string) == \"\" {\n\t\t\t\trole = \"bird\"\n\t\t\t}\n\t\t\treturn role.(string), nil\n\t\t})).\n\t\tAddLambda(\"input\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\treturn \"你的叫声是怎样的？\", nil\n\t\t}))\n\n\t// 顺序节点\n\trolePlayChain := NewChain[map[string]any, *schema.Message]()\n\trolePlayChain.\n\t\tAppendChatTemplate(prompt.FromMessages(schema.FString, schema.SystemMessage(`You are a {role}.`), schema.UserMessage(`{input}`))).\n\t\tAppendChatModel(cm)\n\n\t// 构建 chain\n\n\tchain := NewChain[map[string]any, string]()\n\tchain.\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t// do some logic to prepare kv as variables for next Node\n\t\t\t// just pass through\n\t\t\tt.Log(\"in view lambda: \", kvs)\n\t\t\treturn kvs, nil\n\t\t})).\n\t\tAppendBranch(NewChainBranch[map[string]any](branchCond).AddLambda(\"b1\", b1).AddLambda(\"b2\", b2)).\n\t\tAppendPassthrough().\n\t\tAppendParallel(parallel).\n\t\tAppendGraph(rolePlayChain).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, m *schema.Message) (string, error) {\n\t\t\t// do some logic to check the output or something\n\t\t\tt.Log(\"in view of messages: \", m.Content)\n\n\t\t\treturn m.Content, nil\n\t\t}))\n\n\tr, err := chain.Compile(context.Background())\n\tassert.Nil(t, err)\n\n\tout, err := r.Invoke(context.Background(), map[string]any{})\n\tassert.Nil(t, err)\n\tt.Log(err)\n\n\tt.Log(\"out is : \", out)\n}\n\nfunc TestChainWithException(t *testing.T) {\n\tchain := NewChain[map[string]any, string]()\n\tchain.\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t// do some logic to prepare kv as variables for next Node\n\t\t\t// just pass through\n\t\t\tt.Log(\"in view lambda: \", kvs)\n\t\t\treturn kvs, nil\n\t\t})).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in view lambda 02: \", kvs)\n\t\t\treturn kvs, nil\n\t\t}), WithNodeKey(\"xlam\"))\n\n\t// items with parallels\n\tparallel := NewParallel()\n\tparallel.\n\t\tAddLambda(\"hello\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\tt.Log(\"in parallel item 01\")\n\t\t\treturn \"world\", nil\n\t\t})).\n\t\tAddLambda(\"world\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\tt.Log(\"in parallel item 02\")\n\t\t\treturn \"hello\", nil\n\t\t}))\n\n\t// sequence items\n\tnchain := NewChain[map[string]any, map[string]any]()\n\tnchain.\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in sequence item 01\")\n\t\t\treturn kvs, nil\n\t\t})).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in sequence item 02\")\n\t\t\treturn kvs, nil\n\t\t}))\n\n\tbranchCond := func(ctx context.Context, input map[string]any) (string, error) {\n\t\tif rand.Intn(2) == 1 {\n\t\t\treturn \"b1\", nil\n\t\t}\n\t\treturn \"b2\", nil\n\t}\n\n\tb1 := InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\tt.Log(\"hello in branch lambda 01\")\n\t\tkvs[\"role\"] = \"cat\"\n\t\treturn kvs, nil\n\t})\n\tb2 := InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\treturn kvs, nil\n\t})\n\n\t// sequence with branch\n\tchain.AppendBranch(NewChainBranch[map[string]any](branchCond).AddLambda(\"b1\", b1).AddLambda(\"b2\", b2))\n\n\t// parallel with sequence\n\tparallel.AddGraph(\"test_sequence\", nchain)\n\n\t// parallel with parallel\n\tnpara := NewParallel().\n\t\tAddLambda(\"test_parallel1\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t})).\n\t\tAddLambda(\"test_parallel2\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\n\t// parallel with graph\n\tngraph := NewChain[map[string]any, map[string]any]().\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in graph item 01\")\n\t\t\treturn kvs, nil\n\t\t})).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in graph item 02\")\n\t\t\treturn kvs, nil\n\t\t}))\n\tnc := NewChain[map[string]any, map[string]any]()\n\tnc.AppendGraph(ngraph)\n\tparallel.AddGraph(\"test_graph\", nc)\n\n\tchain.AppendPassthrough()\n\n\t// sequence with parallel\n\tchain.AppendParallel(npara)\n\n\t// 构建 chain\n\tchain.\n\t\tAppendGraph(nchain).\n\t\tAppendParallel(parallel).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\tt.Log(\"in last view lambda: \", kvs)\n\t\t\treturn \"hello last\", nil\n\t\t}))\n\n\tctx := context.Background()\n\n\tr, err := chain.Compile(ctx)\n\tassert.Nil(t, err)\n\n\tout, err := r.Invoke(ctx, map[string]any{\"test\": \"test\"})\n\tassert.Nil(t, err)\n\tt.Log(\"out is : \", out)\n}\n\nfunc TestEmptyList(t *testing.T) {\n\tctx := context.Background()\n\n\t// no nodes in chain\n\tchain := NewChain[map[string]any, map[string]any]()\n\t_, err := chain.Compile(ctx)\n\tassert.Error(t, err)\n\n\t// no nodes in parallel\n\tparallel := NewParallel()\n\tchain = NewChain[map[string]any, map[string]any]()\n\tchain.AppendParallel(parallel)\n\n\t_, err = chain.Compile(ctx)\n\tassert.Error(t, err)\n\n\t// no nodes in sequence\n\temptyChain := NewChain[map[string]any, map[string]any]()\n\tchain = NewChain[map[string]any, map[string]any]()\n\n\tchain.\n\t\tAppendParallel(parallel).\n\t\tAppendGraph(emptyChain).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\n\t_, err = chain.Compile(ctx)\n\tassert.Error(t, err)\n}\n\nfunc TestChainList(t *testing.T) {\n\tchain := NewChain[map[string]any, map[string]any]()\n\tchain.\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in view lambda: \", kvs)\n\t\t\treturn kvs, nil\n\t\t}))\n\n\t// parallel\n\tparallel := NewParallel()\n\tparallel.\n\t\tAddLambda(\"test_parallel1\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in parallel item 01\")\n\t\t\treturn kvs, nil\n\t\t}))\n\n\t// seq in parallel\n\tnchain := NewChain[map[string]any, map[string]any]()\n\tnchain.\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in sequence in parallel item 01\")\n\t\t\treturn kvs, nil\n\t\t})).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in sequence in parallel item 02\")\n\t\t\treturn kvs, nil\n\t\t}))\n\n\t// seq in seq\n\tnchainInChain := NewChain[map[string]any, map[string]any]()\n\tnchainInChain.\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in sequence in sequence item 01\")\n\t\t\treturn kvs, nil\n\t\t})).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in sequence in sequence item 02\")\n\t\t\treturn kvs, nil\n\t\t}))\n\n\tnchain.AppendGraph(nchainInChain)\n\n\tparallel.AddGraph(\"test_seq_in_parallel\", nchain)\n\n\tchain.AppendParallel(parallel)\n\n\tr, err := chain.Compile(context.Background())\n\tassert.Nil(t, err)\n\tout, err := r.Invoke(context.Background(), map[string]any{\"test\": \"test\"})\n\tassert.Nil(t, err)\n\tt.Log(\"out is : \", out)\n}\n\nfunc TestChainSingleNode(t *testing.T) {\n\tchain := NewChain[map[string]any, map[string]any]()\n\tchain.\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in view lambda: \", kvs)\n\t\t\treturn kvs, nil\n\t\t}))\n\n\t// single Node in chain (prepare for parallel)\n\tsingleNodeChain := NewChain[map[string]any, map[string]any]()\n\tsingleNodeChain.\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in sequence item 01\")\n\t\t\treturn kvs, nil\n\t\t}))\n\n\t// add parallel\n\tparallel := NewParallel()\n\tparallel.\n\t\tAddLambda(\"test_parallel1_lambda\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\tt.Log(\"in parallel item 01\")\n\t\t\treturn kvs, nil\n\t\t}))\n\n\tparallel.AddGraph(\"test_parallel2_chain\", singleNodeChain)\n\n\tctx := context.Background()\n\n\tchain.AppendParallel(parallel)\n\tr, err := chain.Compile(ctx)\n\tassert.Nil(t, err)\n\n\tout, err := r.Invoke(ctx, map[string]any{\"test\": \"test\"})\n\tassert.Nil(t, err)\n\tt.Log(\"out is : \", out)\n}\n\nfunc TestParallelModels(t *testing.T) {\n\tcm := &mockIntentChatModel{}\n\tchain := NewChain[map[string]any, map[string]any]()\n\tchatSuite := NewChain[map[string]any, string]()\n\tchatSuite.\n\t\tAppendChatTemplate(prompt.FromMessages(schema.FString, schema.SystemMessage(`You are a {role}.`), schema.UserMessage(`{input}`))).\n\t\tAppendChatModel(cm).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, msg *schema.Message) (string, error) {\n\t\t\tt.Log(\"in parallel item 01\")\n\t\t\treturn msg.Content, nil\n\t\t}))\n\n\tparallel := NewParallel()\n\tparallel.\n\t\tAddGraph(\"time001\", chatSuite).\n\t\tAddGraph(\"time002\", chatSuite).\n\t\tAddGraph(\"time003\", chatSuite)\n\n\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\tt.Log(\"in view lambda: \", kvs)\n\t\treturn kvs, nil\n\t}))\n\n\tchain.AppendParallel(parallel)\n\n\tctx := context.Background()\n\n\tr, err := chain.Compile(ctx)\n\tassert.Nil(t, err)\n\n\tout, err := r.Invoke(ctx, map[string]any{\"role\": \"cat\", \"input\": \"你怎么叫的？\"})\n\tassert.Nil(t, err)\n\n\tt.Log(\"out is : \", out)\n}\n\nfunc TestChainMultiNodes(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"test embedding Node\", func(t *testing.T) {\n\t\tchain := NewChain[[]string, [][]float64]()\n\n\t\tmockCtrl := gomock.NewController(t)\n\t\teb := embedding.NewMockEmbedder(mockCtrl)\n\t\tchain.AppendEmbedding(eb)\n\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, r)\n\t})\n\n\tt.Run(\"test retriever Node\", func(t *testing.T) {\n\t\tchain := NewChain[string, []*schema.Document]()\n\n\t\tchain.AppendRetriever(retriever.NewMockRetriever(gomock.NewController(t)))\n\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, r)\n\t})\n\n\tt.Run(\"test chat model\", func(t *testing.T) {\n\t\tchain := NewChain[[]*schema.Message, *schema.Message]()\n\n\t\tcm := &mockIntentChatModel{}\n\t\tchain.AppendChatModel(cm)\n\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, r)\n\t})\n\n\tt.Run(\"test chat template\", func(t *testing.T) {\n\t\tchain := NewChain[map[string]any, []*schema.Message]()\n\n\t\tchatTemplate := prompt.FromMessages(schema.FString)\n\t\tchain.AppendChatTemplate(chatTemplate)\n\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, r)\n\t})\n\n\tt.Run(\"test lambda\", func(t *testing.T) {\n\t\tchain := NewChain[map[string]any, map[string]any]()\n\n\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, r)\n\t})\n\n\tt.Run(\"test indexer\", func(t *testing.T) {\n\t\tchain := NewChain[[]*schema.Document, []string]()\n\n\t\tchain.AppendIndexer(indexer.NewMockIndexer(gomock.NewController(t)))\n\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, r)\n\t})\n\n\tt.Run(\"test parallel\", func(t *testing.T) {\n\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\tparallel := NewParallel()\n\t\tparallel.AddLambda(\"test_parallel\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tchain.AppendParallel(parallel)\n\t\t_, err := chain.Compile(ctx)\n\t\tassert.Error(t, err)\n\n\t\tchain = NewChain[map[string]any, map[string]any]()\n\t\tparallel = NewParallel()\n\t\tparallel.AddLambda(\"test_parallel\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tparallel.AddLambda(\"test_parallel\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tchain.AppendParallel(parallel)\n\t\t_, err = chain.Compile(ctx)\n\t\tassert.Error(t, err)\n\n\t\tchain = NewChain[map[string]any, map[string]any]()\n\t\tparallel = NewParallel()\n\t\tparallel.AddLambda(\"test_parallel\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tparallel.AddLambda(\"test_parallel1\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tchain.AppendParallel(parallel)\n\t\t_, err = chain.Compile(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tchain = NewChain[map[string]any, map[string]any]()\n\t\tparallel = NewParallel()\n\t\tparallel.AddLambda(\"test_parallel\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tparallel.AddLambda(\"test_parallel1\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tchain.AppendParallel(parallel)\n\n\t\tparallel1 := NewParallel()\n\t\tparallel1.AddLambda(\"test_parallel\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tparallel1.AddLambda(\"test_parallel1\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tchain.AppendParallel(parallel1)\n\n\t\t_, err = chain.Compile(ctx)\n\t\tassert.Error(t, err)\n\t})\n\n\tt.Run(\"test tools Node\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\ttoolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{})\n\t\tassert.NoError(t, err)\n\t\tchain.AppendToolsNode(toolsNode)\n\t})\n\n\tt.Run(\"test chain with compile option\", func(t *testing.T) {\n\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}))\n\t\tr, err := chain.Compile(ctx, WithMaxRunSteps(10))\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, r)\n\t})\n\n\tt.Run(\"test chain return type\", func(t *testing.T) {\n\t\tt.Run(\"test chain any output type\", func(t *testing.T) {\n\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (any, error) {\n\t\t\t\treturn 1, nil\n\t\t\t}))\n\t\t\t_, err := chain.Compile(ctx)\n\t\t\tassert.Nil(t, err)\n\t\t})\n\n\t\tt.Run(\"test chain error output type\", func(t *testing.T) {\n\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\t\treturn \"123\", nil\n\t\t\t}))\n\t\t\t_, err := chain.Compile(ctx)\n\t\t\tassert.Error(t, err)\n\t\t})\n\n\t\tt.Run(\"test chain error input type\", func(t *testing.T) {\n\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, input string) (map[string]any, error) {\n\t\t\t\treturn nil, nil\n\t\t\t}))\n\t\t\t_, err := chain.Compile(ctx)\n\t\t\tassert.Error(t, err)\n\t\t})\n\t})\n\n}\n\nfunc TestParallelMultiNodes(t *testing.T) {\n\tctx := context.Background()\n\tp := NewParallel()\n\tp.AddLambda(\"lambda\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\treturn kvs, nil\n\t}))\n\tp.AddGraph(\"graph\", NewChain[map[string]any, map[string]any]())\n\tp.AddIndexer(\"indexer\", indexer.NewMockIndexer(gomock.NewController(t)))\n\tp.AddLoader(\"loader\", document.NewMockLoader(gomock.NewController(t)))\n\tp.AddDocumentTransformer(\"document transformer\", document.NewMockTransformer(gomock.NewController(t)))\n\tp.AddRetriever(\"retriever\", retriever.NewMockRetriever(gomock.NewController(t)))\n\tp.AddChatModel(\"chatmodel\", model.NewMockChatModel(gomock.NewController(t)))\n\tp.AddChatTemplate(\"chatTemplate\", prompt.FromMessages(schema.FString, schema.SystemMessage(\"hello\")))\n\tp.AddEmbedding(\"embedding\", embedding.NewMockEmbedder(gomock.NewController(t)))\n\tp.AddPassthrough(\"passthrough\")\n\ttoolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{})\n\tassert.NoError(t, err)\n\tp.AddToolsNode(\"tools\", toolsNode)\n\n\tassert.Greater(t, len(p.nodes), 6)\n\n\tctrl := gomock.NewController(t)\n\tp = NewParallel()\n\tp.AddIndexer(\"key\", indexer.NewMockIndexer(ctrl))\n\tp.AddLoader(\"key\", document.NewMockLoader(ctrl))\n\tp.AddRetriever(\"r\", retriever.NewMockRetriever(ctrl))\n\tassert.NotNil(t, p.err)\n\n\tp = NewParallel()\n\tp.addNode(\"k\", nil, nil)\n\tassert.NotNil(t, p.err)\n\n\tp = &Parallel{\n\t\toutputKeys: nil,\n\t}\n\tp.addNode(\"k\", &graphNode{}, nil)\n\tassert.NotNil(t, p.err)\n}\n\ntype FakeLambdaOptions struct {\n\tInfo string\n}\n\ntype FakeLambdaOption func(opt *FakeLambdaOptions)\n\nfunc FakeWithLambdaInfo(info string) FakeLambdaOption {\n\treturn func(opt *FakeLambdaOptions) {\n\t\topt.Info = info\n\t}\n}\n\nfunc TestChainWithNodeKey(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"test normal chain with node key option\", func(t *testing.T) {\n\n\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\treturn kvs, nil\n\t\t}), WithNodeKey(\"lambda_01\"))\n\n\t\tb := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) {\n\t\t\treturn \"lambda_02\", nil\n\t\t})\n\n\t\tb.AddLambda(\"lambda_02\", InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (map[string]any, error) {\n\t\t\topt := &FakeLambdaOptions{}\n\t\t\tfor _, optFn := range opts {\n\t\t\t\toptFn(opt)\n\t\t\t}\n\t\t\tkvs[\"lambda_02\"] = opt.Info\n\t\t\treturn kvs, nil\n\t\t}), WithNodeKey(\"lambda_02\"))\n\n\t\tb.AddLambda(\"lambda_03\", InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (map[string]any, error) {\n\t\t\topt := &FakeLambdaOptions{}\n\t\t\tfor _, optFn := range opts {\n\t\t\t\toptFn(opt)\n\t\t\t}\n\t\t\tkvs[\"lambda_03\"] = opt.Info\n\t\t\treturn kvs, nil\n\t\t}), WithNodeKey(\"lambda_03\"))\n\n\t\tchain.AppendBranch(b)\n\n\t\tchain.AppendPassthrough()\n\n\t\tp := NewParallel()\n\t\tp.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\treturn kvs[\"lambda_02\"].(string), nil\n\t\t}))\n\t\tp.AddLambda(\"lambda_04\", InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (string, error) {\n\t\t\topt := &FakeLambdaOptions{}\n\t\t\tfor _, optFn := range opts {\n\t\t\t\toptFn(opt)\n\t\t\t}\n\t\t\treturn opt.Info, nil\n\t\t}), WithNodeKey(\"lambda_04\"))\n\n\t\tp.AddLambda(\"lambda_05\", InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (string, error) {\n\t\t\topt := &FakeLambdaOptions{}\n\t\t\tfor _, optFn := range opts {\n\t\t\t\toptFn(opt)\n\t\t\t}\n\t\t\treturn opt.Info, nil\n\t\t}), WithNodeKey(\"lambda_05\"))\n\t\tchain.AppendParallel(p)\n\n\t\tchain.AppendLambda(InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (map[string]any, error) {\n\t\t\topt := &FakeLambdaOptions{}\n\t\t\tfor _, optFn := range opts {\n\t\t\t\toptFn(opt)\n\t\t\t}\n\t\t\tkvs[\"lambda_06\"] = opt.Info\n\t\t\treturn kvs, nil\n\t\t}), WithNodeKey(\"lambda_06\"))\n\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.Nil(t, err)\n\n\t\tres, err := r.Invoke(ctx, map[string]any{},\n\t\t\tWithLambdaOption(FakeWithLambdaInfo(\"normal\")),\n\t\t\tWithLambdaOption(FakeWithLambdaInfo(\"info_lambda_02\")).DesignateNode(\"lambda_02\"), // branch\n\t\t\tWithLambdaOption(FakeWithLambdaInfo(\"info_lambda_03\")).DesignateNode(\"lambda_03\"), // branch (wont run)\n\t\t\tWithLambdaOption(FakeWithLambdaInfo(\"info_lambda_05\")).DesignateNode(\"lambda_05\"), // parallel\n\t\t)\n\t\tassert.Nil(t, err)\n\n\t\tassert.Equal(t, \"info_lambda_02\", res[\"lambda_02\"]) // transmit option with DesigateNode\n\t\tassert.Equal(t, \"info_lambda_05\", res[\"lambda_05\"]) // transmit option with DesigateNode\n\t\tassert.Equal(t, \"normal\", res[\"lambda_06\"])         // without DesigateNode, using default option\n\t})\n\n\tt.Run(\"test chain with node key option and error with correct error info\", func(t *testing.T) {\n\n\t\tt.Run(\"compile error of chain\", func(t *testing.T) {\n\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\t\treturn \"123\", nil\n\t\t\t}), WithNodeKey(\"lambda_01\"))\n\n\t\t\tc, err := chain.Compile(ctx)\n\t\t\tassert.Nil(t, c)\n\t\t\tfmt.Printf(\"%+v\\n\", err)\n\n\t\t\tassert.Contains(t, err.Error(), \"edge[lambda_01]\")\n\t\t})\n\n\t\tt.Run(\"compile error of branch\", func(t *testing.T) {\n\t\t\tt.Run(\"without node key\", func(t *testing.T) {\n\t\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn kvs, nil\n\t\t\t\t}), WithNodeKey(\"lambda_01\"))\n\t\t\t\tb := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) {\n\t\t\t\t\treturn \"lambda_02\", nil\n\t\t\t\t})\n\t\t\t\tb.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn kvs, nil\n\t\t\t\t}))\n\t\t\t\tb.AddLambda(\"lambda_03\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\t\t\treturn \"\", nil\n\t\t\t\t}))\n\t\t\t\tchain.AppendBranch(b)\n\t\t\t\tc, err := chain.Compile(ctx)\n\t\t\t\tassert.Nil(t, c)\n\t\t\t\tfmt.Printf(\"%+v\\n\", err)\n\t\t\t\tassert.Contains(t, err.Error(), \"edge[node_1_branch_lambda_03]\") // with no node key option, will use default node key\n\t\t\t})\n\n\t\t\tt.Run(\"with node key\", func(t *testing.T) {\n\t\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn kvs, nil\n\t\t\t\t}), WithNodeKey(\"lambda_01\"))\n\n\t\t\t\tb := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) {\n\t\t\t\t\treturn \"lambda_02\", nil\n\t\t\t\t})\n\t\t\t\tb.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn kvs, nil\n\t\t\t\t}), WithNodeKey(\"lambda_02\"))\n\n\t\t\t\tb.AddLambda(\"lambda_03\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) {\n\t\t\t\t\treturn \"123\", nil\n\t\t\t\t}), WithNodeKey(\"key_of_lambda_03\"))\n\n\t\t\t\tchain.AppendBranch(b)\n\t\t\t\tc, err := chain.Compile(ctx)\n\t\t\t\tassert.Nil(t, c)\n\t\t\t\tfmt.Printf(\"%+v\\n\", err)\n\t\t\t\tassert.Contains(t, err.Error(), \"edge[key_of_lambda_03]\")\n\t\t\t})\n\t\t})\n\n\t\tt.Run(\"compile error of parallel\", func(t *testing.T) {\n\t\t\tt.Run(\"without node key\", func(t *testing.T) {\n\t\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn kvs, nil\n\t\t\t\t}), WithNodeKey(\"lambda_01\"))\n\t\t\t\tp := NewParallel()\n\t\t\t\tp.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn kvs, nil\n\t\t\t\t}))\n\t\t\t\tp.AddLambda(\"lambda_03\", InvokableLambda(func(ctx context.Context, v string) (string, error) {\n\t\t\t\t\treturn \"\", nil\n\t\t\t\t}))\n\n\t\t\t\tchain.AppendParallel(p)\n\n\t\t\t\tc, err := chain.Compile(ctx)\n\t\t\t\tassert.Nil(t, c)\n\t\t\t\tfmt.Printf(\"%+v\\n\", err)\n\t\t\t\tassert.Contains(t, err.Error(), \"to=node_1_parallel_1\") // with no node key option, will use default node key\n\t\t\t})\n\n\t\t\tt.Run(\"with node key\", func(t *testing.T) {\n\t\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\t\tchain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn kvs, nil\n\t\t\t\t}), WithNodeKey(\"lambda_01\"))\n\t\t\t\tp := NewParallel()\n\t\t\t\tp.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn kvs, nil\n\t\t\t\t}), WithNodeKey(\"lambda_02\"))\n\t\t\t\tp.AddLambda(\"lambda_03\", InvokableLambda(func(ctx context.Context, v string) (string, error) {\n\t\t\t\t\treturn \"\", nil\n\t\t\t\t}), WithNodeKey(\"key_of_lambda_03\"))\n\t\t\t\tchain.AppendParallel(p)\n\t\t\t\tc, err := chain.Compile(ctx)\n\t\t\t\tassert.Nil(t, c)\n\t\t\t\tfmt.Printf(\"%+v\\n\", err)\n\t\t\t\tassert.Contains(t, err.Error(), \"to=key_of_lambda_03\")\n\t\t\t})\n\t\t})\n\n\t\tt.Run(\"invoke error\", func(t *testing.T) {\n\t\t\tt.Run(\"branch with out node key\", func(t *testing.T) {\n\t\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\n\t\t\t\tb := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) {\n\t\t\t\t\treturn \"lambda_01\", nil\n\t\t\t\t})\n\n\t\t\t\tb.AddLambda(\"lambda_01\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn nil, fmt.Errorf(\"fake error\")\n\t\t\t\t}))\n\t\t\t\tb.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn nil, nil\n\t\t\t\t}))\n\n\t\t\t\tchain.AppendBranch(b)\n\t\t\t\tc, err := chain.Compile(ctx)\n\t\t\t\tassert.Nil(t, err)\n\n\t\t\t\t_, err = c.Invoke(ctx, map[string]any{})\n\t\t\t\tfmt.Printf(\"%+v\\n\", err)\n\t\t\t\tassert.Error(t, err)\n\t\t\t\tassert.Contains(t, err.Error(), \"node_0_branch_lambda_01\") // with no node key option, will use default node key\n\t\t\t})\n\n\t\t\tt.Run(\"branch with node key\", func(t *testing.T) {\n\t\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\t\tb := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) {\n\t\t\t\t\treturn \"lambda_01\", nil\n\t\t\t\t})\n\t\t\t\tb.AddLambda(\"lambda_01\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn nil, fmt.Errorf(\"fake error\")\n\t\t\t\t}), WithNodeKey(\"key_of_lambda_01\"))\n\t\t\t\tb.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn nil, nil\n\t\t\t\t}))\n\n\t\t\t\tchain.AppendBranch(b)\n\t\t\t\tc, err := chain.Compile(ctx)\n\t\t\t\tassert.Nil(t, err)\n\t\t\t\t_, err = c.Invoke(ctx, map[string]any{})\n\t\t\t\tfmt.Printf(\"%+v\\n\", err)\n\t\t\t\tassert.Error(t, err)\n\t\t\t\tassert.Contains(t, err.Error(), \"key_of_lambda_01\")\n\t\t\t})\n\n\t\t\tt.Run(\"parallel with out node key\", func(t *testing.T) {\n\t\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\t\tp := NewParallel()\n\t\t\t\tp.AddLambda(\"lambda_01\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn nil, fmt.Errorf(\"fake error\")\n\t\t\t\t}))\n\t\t\t\tp.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn nil, nil\n\t\t\t\t}))\n\t\t\t\tchain.AppendParallel(p)\n\t\t\t\tc, err := chain.Compile(ctx)\n\t\t\t\tassert.Nil(t, err)\n\t\t\t\t_, err = c.Invoke(ctx, map[string]any{})\n\t\t\t\tfmt.Printf(\"%+v\\n\", err)\n\t\t\t\tassert.Error(t, err)\n\t\t\t\tassert.Contains(t, err.Error(), \"node_0_parallel_0\") // with no node key option, will use default node key\n\t\t\t})\n\n\t\t\tt.Run(\"parallel with node key\", func(t *testing.T) {\n\t\t\t\tchain := NewChain[map[string]any, map[string]any]()\n\t\t\t\tp := NewParallel()\n\t\t\t\tp.AddLambda(\"lambda_01\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn nil, fmt.Errorf(\"fake error\")\n\t\t\t\t}), WithNodeKey(\"key_of_lambda_01\"))\n\t\t\t\tp.AddLambda(\"lambda_02\", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) {\n\t\t\t\t\treturn nil, nil\n\t\t\t\t}))\n\t\t\t\tchain.AppendParallel(p)\n\t\t\t\tc, err := chain.Compile(ctx)\n\t\t\t\tassert.Nil(t, err)\n\t\t\t\t_, err = c.Invoke(ctx, map[string]any{})\n\t\t\t\tfmt.Printf(\"%+v\\n\", err)\n\t\t\t\tassert.Error(t, err)\n\t\t\t\tassert.Contains(t, err.Error(), \"key_of_lambda_01\")\n\t\t\t})\n\t\t})\n\n\t})\n\n}\n"
  },
  {
    "path": "compose/checkpoint.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/internal/core\"\n\t\"github.com/cloudwego/eino/internal/serialization\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc init() {\n\tschema.RegisterName[*checkpoint](\"_eino_checkpoint\")\n\tschema.RegisterName[*dagChannel](\"_eino_dag_channel\")\n\tschema.RegisterName[*pregelChannel](\"_eino_pregel_channel\")\n\tschema.RegisterName[dependencyState](\"_eino_dependency_state\")\n\t_ = serialization.GenericRegister[channel](\"_eino_channel\")\n}\n\n// RegisterSerializableType registers a custom type for eino serialization.\n// This allows eino to properly serialize and deserialize custom types.\n// Both custom interfaces and structs need to be registered using this function.\n// Types only need to be registered once - pointers and other references will be handled automatically.\n// All built-in eino types are already registered.\n// Parameters:\n// - name: A unique identifier for the type being registered (should not start with \"_eino\")\n// - T: The generic type parameter representing the type to register\n// Returns:\n// - error: An error if registration fails (e.g., if the type is already registered)\n// Deprecated: RegisterSerializableType is deprecated. Use schema.RegisterName[T](name) instead.\nfunc RegisterSerializableType[T any](name string) (err error) {\n\treturn serialization.GenericRegister[T](name)\n}\n\ntype CheckPointStore = core.CheckPointStore\n\ntype Serializer interface {\n\tMarshal(v any) ([]byte, error)\n\tUnmarshal(data []byte, v any) error\n}\n\n// WithCheckPointStore sets the checkpoint store implementation for a graph.\nfunc WithCheckPointStore(store CheckPointStore) GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\to.checkPointStore = store\n\t}\n}\n\n// WithSerializer sets the serializer used to persist checkpoint state.\nfunc WithSerializer(serializer Serializer) GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\to.serializer = serializer\n\t}\n}\n\n// WithCheckPointID sets the checkpoint ID to load from and write to by default.\nfunc WithCheckPointID(checkPointID string) Option {\n\treturn Option{\n\t\tcheckPointID: &checkPointID,\n\t}\n}\n\n// WithWriteToCheckPointID specifies a different checkpoint ID to write to.\n// If not provided, the checkpoint ID from WithCheckPointID will be used for writing.\n// This is useful for scenarios where you want to load from an existed checkpoint\n// but save the progress to a new, separate checkpoint.\nfunc WithWriteToCheckPointID(checkPointID string) Option {\n\treturn Option{\n\t\twriteToCheckPointID: &checkPointID,\n\t}\n}\n\n// WithForceNewRun forces the graph to run from the beginning, ignoring any checkpoints.\nfunc WithForceNewRun() Option {\n\treturn Option{\n\t\tforceNewRun: true,\n\t}\n}\n\n// StateModifier modifies state during checkpoint operations for a given node path.\ntype StateModifier func(ctx context.Context, path NodePath, state any) error\n\n// WithStateModifier installs a state modifier invoked during checkpoint read/write.\nfunc WithStateModifier(sm StateModifier) Option {\n\treturn Option{\n\t\tstateModifier: sm,\n\t}\n}\n\ntype checkpoint struct {\n\tChannels       map[string]channel\n\tInputs         map[string] /*node key*/ any /*input*/\n\tState          any\n\tSkipPreHandler map[string]bool\n\tRerunNodes     []string\n\n\tSubGraphs map[string]*checkpoint\n\n\tInterruptID2Addr  map[string]Address\n\tInterruptID2State map[string]core.InterruptState\n}\n\ntype stateModifierKey struct{}\ntype checkPointKey struct{} // *checkpoint\n\nfunc getStateModifier(ctx context.Context) StateModifier {\n\tif sm, ok := ctx.Value(stateModifierKey{}).(StateModifier); ok {\n\t\treturn sm\n\t}\n\treturn nil\n}\n\nfunc setStateModifier(ctx context.Context, modifier StateModifier) context.Context {\n\treturn context.WithValue(ctx, stateModifierKey{}, modifier)\n}\n\nfunc getCheckPointFromStore(ctx context.Context, id string, cpr *checkPointer) (cp *checkpoint, err error) {\n\tcp, existed, err := cpr.get(ctx, id)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif !existed {\n\t\treturn nil, nil\n\t}\n\n\treturn cp, nil\n}\n\nfunc setCheckPointToCtx(ctx context.Context, cp *checkpoint) context.Context {\n\tctx = core.PopulateInterruptState(ctx, cp.InterruptID2Addr, cp.InterruptID2State)\n\treturn context.WithValue(ctx, checkPointKey{}, cp)\n}\n\nfunc getCheckPointFromCtx(ctx context.Context) *checkpoint {\n\tif cp, ok := ctx.Value(checkPointKey{}).(*checkpoint); ok {\n\t\treturn cp\n\t}\n\treturn nil\n}\n\nfunc forwardCheckPoint(ctx context.Context, nodeKey string) context.Context {\n\tcp := getCheckPointFromCtx(ctx)\n\tif cp == nil {\n\t\treturn ctx\n\t}\n\n\tif subCP, ok := cp.SubGraphs[nodeKey]; ok {\n\t\tdelete(cp.SubGraphs, nodeKey) // only forward once\n\t\treturn context.WithValue(ctx, checkPointKey{}, subCP)\n\t}\n\treturn context.WithValue(ctx, checkPointKey{}, (*checkpoint)(nil))\n}\n\nfunc newCheckPointer(\n\tinputPairs, outputPairs map[string]streamConvertPair,\n\tstore CheckPointStore,\n\tserializer Serializer,\n) *checkPointer {\n\tif serializer == nil {\n\t\tserializer = &serialization.InternalSerializer{}\n\t}\n\treturn &checkPointer{\n\t\tsc:         newStreamConverter(inputPairs, outputPairs),\n\t\tstore:      store,\n\t\tserializer: serializer,\n\t}\n}\n\ntype checkPointer struct {\n\tsc         *streamConverter\n\tstore      CheckPointStore\n\tserializer Serializer\n}\n\nfunc (c *checkPointer) get(ctx context.Context, id string) (*checkpoint, bool, error) {\n\tdata, existed, err := c.store.Get(ctx, id)\n\tif err != nil || existed == false {\n\t\treturn nil, existed, err\n\t}\n\n\tcp := &checkpoint{}\n\terr = c.serializer.Unmarshal(data, cp)\n\tif err != nil {\n\t\treturn nil, false, err\n\t}\n\n\treturn cp, true, nil\n}\n\nfunc (c *checkPointer) set(ctx context.Context, id string, cp *checkpoint) error {\n\tdata, err := c.serializer.Marshal(cp)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treturn c.store.Set(ctx, id, data)\n}\n\n// MigrateCheckpointState is an advanced compatibility utility for checkpoint upgrades.\n//\n// It decodes checkpoint bytes using the given serializer, applies migrate to checkpoint.State and\n// all nested SubGraphs' states, then re-encodes the checkpoint.\n//\n// Typical use cases:\n//   - Resume-time migration when you changed your graph state type/schema and need to load old\n//     checkpoints without discarding them.\n//   - Framework-level backward compatibility (e.g. ADK upgrading checkpoints across versions).\n//\n// Migrate callback contract:\n//   - Returns (newState, changed, error).\n//   - If changed is false, the state is left as-is.\n//   - If error is non-nil, migration stops and the error is returned to the caller.\n//\n// The original bytes are returned only if no state was changed anywhere in the checkpoint tree.\nfunc MigrateCheckpointState(data []byte, serializer Serializer, migrate func(state any) (any, bool, error)) ([]byte, error) {\n\tcp := &checkpoint{}\n\tif err := serializer.Unmarshal(data, cp); err != nil {\n\t\treturn nil, err\n\t}\n\tchanged, err := migrateCheckpoint(cp, migrate)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif !changed {\n\t\treturn data, nil\n\t}\n\treturn serializer.Marshal(cp)\n}\n\n// migrateCheckpoint recursively applies migrate to cp.State and all SubGraphs.\nfunc migrateCheckpoint(cp *checkpoint, migrate func(state any) (any, bool, error)) (bool, error) {\n\tanyChanged := false\n\tif cp.State != nil {\n\t\tnewState, changed, err := migrate(cp.State)\n\t\tif err != nil {\n\t\t\treturn false, err\n\t\t}\n\t\tif changed {\n\t\t\tcp.State = newState\n\t\t\tanyChanged = true\n\t\t}\n\t}\n\tfor _, sub := range cp.SubGraphs {\n\t\tchanged, err := migrateCheckpoint(sub, migrate)\n\t\tif err != nil {\n\t\t\treturn false, err\n\t\t}\n\t\tif changed {\n\t\t\tanyChanged = true\n\t\t}\n\t}\n\treturn anyChanged, nil\n}\n\n// convertCheckPoint if value in checkpoint is streamReader, convert it to non-stream\nfunc (c *checkPointer) convertCheckPoint(cp *checkpoint, isStream bool) (err error) {\n\tfor _, ch := range cp.Channels {\n\t\terr = ch.convertValues(func(m map[string]any) error {\n\t\t\treturn c.sc.convertOutputs(isStream, m)\n\t\t})\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\terr = c.sc.convertInputs(isStream, cp.Inputs)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treturn nil\n}\n\n// convertCheckPoint convert values in checkpoint to streamReader if needed\nfunc (c *checkPointer) restoreCheckPoint(cp *checkpoint, isStream bool) (err error) {\n\tfor _, ch := range cp.Channels {\n\t\terr = ch.convertValues(func(m map[string]any) error {\n\t\t\treturn c.sc.restoreOutputs(isStream, m)\n\t\t})\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\terr = c.sc.restoreInputs(isStream, cp.Inputs)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treturn nil\n}\n\nfunc newStreamConverter(inputPairs, outputPairs map[string]streamConvertPair) *streamConverter {\n\treturn &streamConverter{\n\t\tinputPairs:  inputPairs,\n\t\toutputPairs: outputPairs,\n\t}\n}\n\ntype streamConverter struct {\n\tinputPairs, outputPairs map[string]streamConvertPair\n}\n\nfunc (s *streamConverter) convertInputs(isStream bool, values map[string]any) error {\n\treturn convert(values, s.inputPairs, isStream)\n}\n\nfunc (s *streamConverter) restoreInputs(isStream bool, values map[string]any) error {\n\treturn restore(values, s.inputPairs, isStream)\n}\n\nfunc (s *streamConverter) convertOutputs(isStream bool, values map[string]any) error {\n\treturn convert(values, s.outputPairs, isStream)\n}\n\nfunc (s *streamConverter) restoreOutputs(isStream bool, values map[string]any) error {\n\treturn restore(values, s.outputPairs, isStream)\n}\n\nfunc convert(values map[string]any, convPairs map[string]streamConvertPair, isStream bool) error {\n\tif !isStream {\n\t\treturn nil\n\t}\n\tfor key, v := range values {\n\t\tconvPair, ok := convPairs[key]\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"checkpoint conv stream fail, node[%s] have not been registered\", key)\n\t\t}\n\t\tsr, ok := v.(streamReader)\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"checkpoint conv stream fail, value of [%s] isn't stream\", key)\n\t\t}\n\t\tnValue, err := convPair.concatStream(sr)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tvalues[key] = nValue\n\t}\n\treturn nil\n}\n\nfunc restore(values map[string]any, convPairs map[string]streamConvertPair, isStream bool) error {\n\tif !isStream {\n\t\treturn nil\n\t}\n\tfor key, v := range values {\n\t\tconvPair, ok := convPairs[key]\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"checkpoint restore stream fail, node[%s] have not been registered\", key)\n\t\t}\n\t\tsr, err := convPair.restoreStream(v)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tvalues[key] = sr\n\t}\n\treturn nil\n}\n"
  },
  {
    "path": "compose/checkpoint_migrate_test.go",
    "content": "/*\n * Copyright 2026 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"errors\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\ntype stubSerializer struct {\n\tunmarshal func(data []byte, v any) error\n\tmarshal   func(v any) ([]byte, error)\n}\n\nfunc (s stubSerializer) Marshal(v any) ([]byte, error) {\n\treturn s.marshal(v)\n}\n\nfunc (s stubSerializer) Unmarshal(data []byte, v any) error {\n\treturn s.unmarshal(data, v)\n}\n\nfunc TestMigrateCheckpointState_UnmarshalError(t *testing.T) {\n\tin := []byte(\"in\")\n\tcodec := stubSerializer{\n\t\tunmarshal: func(_ []byte, _ any) error { return errors.New(\"bad\") },\n\t\tmarshal:   func(_ any) ([]byte, error) { return []byte(\"unused\"), nil },\n\t}\n\t_, err := MigrateCheckpointState(in, codec, func(state any) (any, bool, error) {\n\t\treturn state, false, nil\n\t})\n\tassert.Error(t, err)\n}\n\nfunc TestMigrateCheckpointState_NoChangeReturnsOriginalBytes(t *testing.T) {\n\tin := []byte(\"in\")\n\tcp := &checkpoint{State: \"s\"}\n\tcodec := stubSerializer{\n\t\tunmarshal: func(_ []byte, v any) error {\n\t\t\t*(v.(*checkpoint)) = *cp\n\t\t\treturn nil\n\t\t},\n\t\tmarshal: func(_ any) ([]byte, error) {\n\t\t\treturn []byte(\"marshaled\"), nil\n\t\t},\n\t}\n\tout, err := MigrateCheckpointState(in, codec, func(state any) (any, bool, error) {\n\t\treturn state, false, nil\n\t})\n\tassert.NoError(t, err)\n\tassert.Equal(t, in, out)\n}\n\nfunc TestMigrateCheckpointState_ChangeTriggersMarshal(t *testing.T) {\n\tin := []byte(\"in\")\n\tcp := &checkpoint{State: \"s\"}\n\tvar sawState any\n\tcodec := stubSerializer{\n\t\tunmarshal: func(_ []byte, v any) error {\n\t\t\t*(v.(*checkpoint)) = *cp\n\t\t\treturn nil\n\t\t},\n\t\tmarshal: func(v any) ([]byte, error) {\n\t\t\tsawState = v.(*checkpoint).State\n\t\t\treturn []byte(\"marshaled\"), nil\n\t\t},\n\t}\n\tout, err := MigrateCheckpointState(in, codec, func(state any) (any, bool, error) {\n\t\treturn \"s2\", true, nil\n\t})\n\tassert.NoError(t, err)\n\tassert.Equal(t, []byte(\"marshaled\"), out)\n\tassert.Equal(t, \"s2\", sawState)\n}\n\nfunc TestMigrateCheckpointState_MigrateErrorStops(t *testing.T) {\n\tin := []byte(\"in\")\n\tcp := &checkpoint{\n\t\tState: \"root\",\n\t\tSubGraphs: map[string]*checkpoint{\n\t\t\t\"sub\": {State: \"sub\"},\n\t\t},\n\t}\n\tcodec := stubSerializer{\n\t\tunmarshal: func(_ []byte, v any) error {\n\t\t\t*(v.(*checkpoint)) = *cp\n\t\t\treturn nil\n\t\t},\n\t\tmarshal: func(_ any) ([]byte, error) {\n\t\t\treturn []byte(\"marshaled\"), nil\n\t\t},\n\t}\n\t_, err := MigrateCheckpointState(in, codec, func(state any) (any, bool, error) {\n\t\tif state == \"sub\" {\n\t\t\treturn nil, false, errors.New(\"boom\")\n\t\t}\n\t\treturn state, false, nil\n\t})\n\tassert.Error(t, err)\n}\n"
  },
  {
    "path": "compose/checkpoint_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"io\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/internal/callbacks\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/internal/serialization\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype inMemoryStore struct {\n\tm map[string][]byte\n}\n\nfunc (i *inMemoryStore) Get(_ context.Context, checkPointID string) ([]byte, bool, error) {\n\tv, ok := i.m[checkPointID]\n\treturn v, ok, nil\n}\n\nfunc (i *inMemoryStore) Set(_ context.Context, checkPointID string, checkPoint []byte) error {\n\ti.m[checkPointID] = checkPoint\n\treturn nil\n}\n\nfunc newInMemoryStore() *inMemoryStore {\n\treturn &inMemoryStore{\n\t\tm: make(map[string][]byte),\n\t}\n}\n\ntype testStruct struct {\n\tA string\n}\n\nfunc init() {\n\tschema.Register[testStruct]()\n}\n\nfunc TestSimpleCheckPoint(t *testing.T) {\n\tstore := newInMemoryStore()\n\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) {\n\t\treturn &testStruct{A: \"\"}\n\t}))\n\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"2\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) {\n\t\treturn in + state.A, nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor), WithCheckPointStore(store), WithInterruptAfterNodes([]string{\"1\"}), WithInterruptBeforeNodes([]string{\"2\"}), WithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, &testStruct{A: \"\"}, info.State)\n\tassert.Equal(t, []string{\"2\"}, info.BeforeNodes)\n\tassert.Equal(t, []string{\"1\"}, info.AfterNodes)\n\tassert.Empty(t, info.RerunNodesExtra)\n\tassert.Empty(t, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t}))\n\n\trCtx := ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\tresult, err := r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"start1state2\", result)\n\n\t/*\t_, err = r.Stream(ctx, \"start\", WithCheckPointID(\"2\"))\n\t\tassert.NotNil(t, err)\n\t\tinfo, ok = ExtractInterruptInfo(err)\n\t\tassert.True(t, ok)\n\t\tassert.Equal(t, &testStruct{A: \"\"}, info.State)\n\t\tassert.Equal(t, []string{\"2\"}, info.BeforeNodes)\n\t\tassert.Equal(t, []string{\"1\"}, info.AfterNodes)\n\t\tassert.Empty(t, info.RerunNodesExtra)\n\t\tassert.Empty(t, info.SubGraphs)\n\t\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tInfo: &testStruct{\n\t\t\t\tA: \"\",\n\t\t\t},\n\t\t\tIsRootCause: true,\n\t\t}))\n\n\t\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t\tstreamResult, err := r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"))\n\t\tassert.NoError(t, err)\n\t\tresult = \"\"\n\t\tfor {\n\t\t\tchunk, err := streamResult.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\tresult += chunk\n\t\t}\n\n\t\tassert.Equal(t, \"start1state2\", result)*/\n}\n\nfunc TestCustomStructInAn2y(t *testing.T) {\n\tstore := newInMemoryStore()\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) {\n\t\treturn &testStruct{A: \"\"}\n\t}))\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output *testStruct, err error) {\n\t\treturn &testStruct{A: input + \"1\"}, nil\n\t}), WithOutputKey(\"1\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) {\n\t\treturn input[\"1\"].(*testStruct).A + \"2\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in map[string]any, state *testStruct) (map[string]any, error) {\n\t\tin[\"1\"].(*testStruct).A += state.A\n\t\treturn in, nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(store), WithInterruptAfterNodes([]string{\"1\"}),\n\t\tWithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, &testStruct{A: \"\"}, info.State)\n\tassert.Equal(t, []string{\"1\"}, info.AfterNodes)\n\tassert.Empty(t, info.RerunNodesExtra)\n\tassert.Empty(t, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t}))\n\trCtx := ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\tresult, err := r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"start1state2\", result)\n\n\t_, err = r.Stream(ctx, \"start\", WithCheckPointID(\"2\"))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, &testStruct{A: \"\"}, info.State)\n\tassert.Equal(t, []string{\"1\"}, info.AfterNodes)\n\tassert.Empty(t, info.RerunNodesExtra)\n\tassert.Empty(t, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t}))\n\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\tstreamResult, err := r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"))\n\tassert.NoError(t, err)\n\tresult = \"\"\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tresult += chunk\n\t}\n\n\tassert.Equal(t, \"start1state2\", result)\n}\n\nfunc TestSubGraph(t *testing.T) {\n\tsubG := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) {\n\t\treturn &testStruct{A: \"\"}\n\t}))\n\terr := subG.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = subG.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"2\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) {\n\t\treturn in + state.A, nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = subG.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\n\tg := NewGraph[string, string]()\n\terr = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddGraphNode(\"2\", subG, WithGraphCompileOptions(WithInterruptAfterNodes([]string{\"1\"})))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"3\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"3\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", \"3\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()), WithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"\"},\n\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\n\trCtx := ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\tresult, err := r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"start11state23\", result)\n\n\t_, err = r.Stream(ctx, \"start\", WithCheckPointID(\"2\"))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"\"},\n\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\tRerunNodesExtra: make(map[string]any),\n\t\t\tSubGraphs:       map[string]*InterruptInfo{},\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\tstreamResult, err := r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"))\n\tassert.NoError(t, err)\n\tresult = \"\"\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tresult += chunk\n\t}\n\n\tassert.Equal(t, \"start11state23\", result)\n}\n\ntype testGraphCallback struct {\n\tonStartTimes       int\n\tonEndTimes         int\n\tonStreamStartTimes int\n\tonStreamEndTimes   int\n\tonErrorTimes       int\n}\n\nfunc (t *testGraphCallback) OnStart(ctx context.Context, info *callbacks.RunInfo, _ callbacks.CallbackInput) context.Context {\n\tif info.Component == ComponentOfGraph {\n\t\tt.onStartTimes++\n\t}\n\treturn ctx\n}\n\nfunc (t *testGraphCallback) OnEnd(ctx context.Context, info *callbacks.RunInfo, _ callbacks.CallbackOutput) context.Context {\n\tif info.Component == ComponentOfGraph {\n\t\tt.onEndTimes++\n\t}\n\treturn ctx\n}\n\nfunc (t *testGraphCallback) OnError(ctx context.Context, info *callbacks.RunInfo, _ error) context.Context {\n\tif info.Component == ComponentOfGraph {\n\t\tt.onErrorTimes++\n\t}\n\treturn ctx\n}\n\nfunc (t *testGraphCallback) OnStartWithStreamInput(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {\n\tinput.Close()\n\tif info.Component == ComponentOfGraph {\n\t\tt.onStreamStartTimes++\n\t}\n\treturn ctx\n}\n\nfunc (t *testGraphCallback) OnEndWithStreamOutput(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context {\n\toutput.Close()\n\tif info.Component == ComponentOfGraph {\n\t\tt.onStreamEndTimes++\n\t}\n\treturn ctx\n}\n\nfunc TestNestedSubGraph(t *testing.T) {\n\tsSubG := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) {\n\t\treturn &testStruct{A: \"\"}\n\t}))\n\terr := sSubG.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = sSubG.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"2\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) {\n\t\treturn in + state.A, nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = sSubG.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = sSubG.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = sSubG.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\n\tsubG := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) {\n\t\treturn &testStruct{A: \"\"}\n\t}))\n\terr = subG.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = subG.AddGraphNode(\"2\", sSubG, WithGraphCompileOptions(WithInterruptAfterNodes([]string{\"1\"})), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) {\n\t\treturn in + state.A, nil\n\t}), WithOutputKey(\"2\"))\n\tassert.NoError(t, err)\n\terr = subG.AddLambdaNode(\"3\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"3\", nil\n\t}), WithOutputKey(\"3\"))\n\tassert.NoError(t, err)\n\terr = subG.AddLambdaNode(\"4\", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) {\n\t\treturn input[\"2\"].(string) + \"4\\n\" + input[\"3\"].(string) + \"4\\n\" + input[\"state\"].(string) + \"4\\n\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in map[string]any, state *testStruct) (map[string]any, error) {\n\t\tin[\"state\"] = state.A\n\t\treturn in, nil\n\t}))\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(\"1\", \"3\")\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(\"3\", \"4\")\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(\"2\", \"4\")\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(\"4\", END)\n\tassert.NoError(t, err)\n\n\tg := NewGraph[string, string]()\n\terr = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddGraphNode(\"2\", subG, WithGraphCompileOptions(WithInterruptAfterNodes([]string{\"1\", \"3\"}), WithInterruptBeforeNodes([]string{\"4\"})))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"3\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"3\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", \"3\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()), WithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\n\ttGCB := &testGraphCallback{}\n\t_, err = r.Invoke(ctx, \"start\", WithCheckPointID(\"1\"), WithCallbacks(tGCB))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"\"},\n\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\n\trCtx := ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t_, err = r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"), WithCallbacks(tGCB))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"state\"},\n\t\t\tAfterNodes:      []string{\"3\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs: map[string]*InterruptInfo{\n\t\t\t\t\"2\": {\n\t\t\t\t\tState:           &testStruct{A: \"\"},\n\t\t\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentNode,\n\t\t\t\t\tID:   \"2\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tInfo: &testStruct{\n\t\t\t\tA: \"state\",\n\t\t\t},\n\t\t\tParent: &InterruptCtx{\n\t\t\t\tID: \"runnable:root\",\n\t\t\t\tAddress: Address{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\t\tID:   \"root\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t_, err = r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"), WithCallbacks(tGCB))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"state\"},\n\t\t\tBeforeNodes:     []string{\"4\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"state\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state2\"})\n\tresult, err := r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"), WithCallbacks(tGCB))\n\tassert.NoError(t, err)\n\tassert.Equal(t, `start11state1state24\nstart1134\nstate24\n3`, result)\n\n\t_, err = r.Stream(ctx, \"start\", WithCheckPointID(\"2\"), WithCallbacks(tGCB))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"\"},\n\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t_, err = r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"), WithCallbacks(tGCB))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"state\"},\n\t\t\tAfterNodes:      []string{\"3\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs: map[string]*InterruptInfo{\n\t\t\t\t\"2\": {\n\t\t\t\t\tState:           &testStruct{A: \"\"},\n\t\t\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentNode,\n\t\t\t\t\tID:   \"2\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tInfo: &testStruct{\n\t\t\t\tA: \"state\",\n\t\t\t},\n\t\t\tParent: &InterruptCtx{\n\t\t\t\tAddress: Address{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\t\tID:   \"root\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t_, err = r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"), WithCallbacks(tGCB))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"state\"},\n\t\t\tBeforeNodes:     []string{\"4\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"state\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state2\"})\n\tstreamResult, err := r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"), WithCallbacks(tGCB))\n\tassert.NoError(t, err)\n\tresult = \"\"\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tresult += chunk\n\t}\n\tassert.Equal(t, `start11state1state24\nstart1134\nstate24\n3`, result)\n\n\tassert.Equal(t, 10, tGCB.onStartTimes)       // 3+sSubG*1*3+subG*2*2+g*0\n\tassert.Equal(t, 3, tGCB.onEndTimes)          // success*3\n\tassert.Equal(t, 10, tGCB.onStreamStartTimes) // 3+sSubG*1*3+subG*2*2+g*0\n\tassert.Equal(t, 3, tGCB.onStreamEndTimes)    // success*3\n\tassert.Equal(t, 14, tGCB.onErrorTimes)       // 2*(sSubG*1*3+subG*2*2+g*0)\n\n\t// dag\n\tr, err = g.Compile(ctx, WithCheckPointStore(newInMemoryStore()), WithNodeTriggerMode(AllPredecessor),\n\t\tWithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"\"},\n\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t_, err = r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"state\"},\n\t\t\tAfterNodes:      []string{\"3\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs: map[string]*InterruptInfo{\n\t\t\t\t\"2\": {\n\t\t\t\t\tState:           &testStruct{A: \"\"},\n\t\t\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tID: \"runnable:root;node:2;node:2\",\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentNode,\n\t\t\t\t\tID:   \"2\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tInfo: &testStruct{\n\t\t\t\tA: \"state\",\n\t\t\t},\n\t\t\tParent: &InterruptCtx{\n\t\t\t\tAddress: Address{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\t\tID:   \"root\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t_, err = r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"state\"},\n\t\t\tBeforeNodes:     []string{\"4\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"state\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state2\"})\n\tresult, err = r.Invoke(rCtx, \"start\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, `start11state1state24\nstart1134\nstate24\n3`, result)\n\n\t_, err = r.Stream(ctx, \"start\", WithCheckPointID(\"2\"))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"\"},\n\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t_, err = r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"state\"},\n\t\t\tAfterNodes:      []string{\"3\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs: map[string]*InterruptInfo{\n\t\t\t\t\"2\": {\n\t\t\t\t\tState:           &testStruct{A: \"\"},\n\t\t\t\t\tAfterNodes:      []string{\"1\"},\n\t\t\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentNode,\n\t\t\t\t\tID:   \"2\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tInfo: &testStruct{\n\t\t\t\tA: \"state\",\n\t\t\t},\n\t\t\tParent: &InterruptCtx{\n\t\t\t\tAddress: Address{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\t\tID:   \"root\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state\"})\n\t_, err = r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"))\n\tassert.NotNil(t, err)\n\tinfo, ok = ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, map[string]*InterruptInfo{\n\t\t\"2\": {\n\t\t\tState:           &testStruct{A: \"state\"},\n\t\t\tBeforeNodes:     []string{\"4\"},\n\t\t\tRerunNodesExtra: make(map[string]interface{}),\n\t\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t\t},\n\t}, info.SubGraphs)\n\tassert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{\n\t\tAddress: Address{\n\t\t\t{\n\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\tID:   \"root\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tType: AddressSegmentNode,\n\t\t\t\tID:   \"2\",\n\t\t\t},\n\t\t},\n\t\tInfo: &testStruct{\n\t\t\tA: \"state\",\n\t\t},\n\t\tIsRootCause: true,\n\t\tParent: &InterruptCtx{\n\t\t\tAddress: Address{\n\t\t\t\t{\n\t\t\t\t\tType: AddressSegmentRunnable,\n\t\t\t\t\tID:   \"root\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}))\n\trCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: \"state2\"})\n\tstreamResult, err = r.Stream(rCtx, \"start\", WithCheckPointID(\"2\"))\n\tassert.NoError(t, err)\n\tresult = \"\"\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tresult += chunk\n\t}\n\tassert.Equal(t, `start11state1state24\nstart1134\nstate24\n3`, result)\n}\n\nfunc TestDAGInterrupt(t *testing.T) {\n\tg := NewGraph[string, map[string]any]()\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\ttime.Sleep(time.Millisecond * 100)\n\t\treturn input, nil\n\t}), WithOutputKey(\"1\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\ttime.Sleep(time.Millisecond * 200)\n\t\treturn input, nil\n\t}), WithOutputKey(\"2\"))\n\tassert.NoError(t, err)\n\terr = g.AddPassthroughNode(\"3\")\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(START, \"2\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", \"3\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", \"3\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()), WithInterruptAfterNodes([]string{\"1\", \"2\"}))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"1\"))\n\tinfo, existed := ExtractInterruptInfo(err)\n\tassert.True(t, existed)\n\tassert.Equal(t, []string{\"1\", \"2\"}, info.AfterNodes)\n\n\tresult, err := r.Invoke(ctx, \"\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\"1\": \"input\", \"2\": \"input\"}, result)\n}\n\nfunc TestRerunNodeInterrupt(t *testing.T) {\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) {\n\t\treturn &testStruct{}\n\t}))\n\n\ttimes := 0\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\tdefer func() { times++ }()\n\t\tif times%2 == 0 {\n\t\t\treturn \"\", NewInterruptAndRerunErr(\"test extra\")\n\t\t}\n\t\treturn input, nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) {\n\t\treturn state.A, nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"1\"))\n\tinfo, existed := ExtractInterruptInfo(err)\n\tassert.True(t, existed)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\n\tresult, err := r.Invoke(ctx, \"\", WithCheckPointID(\"1\"), WithStateModifier(func(ctx context.Context, path NodePath, state any) error {\n\t\tstate.(*testStruct).A = \"state\"\n\t\treturn nil\n\t}))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"state\", result)\n\n\t_, err = r.Stream(ctx, \"input\", WithCheckPointID(\"2\"))\n\tinfo, existed = ExtractInterruptInfo(err)\n\tassert.True(t, existed)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\tassert.Equal(t, \"test extra\", info.RerunNodesExtra[\"1\"].(string))\n\n\tstreamResult, err := r.Stream(ctx, \"\", WithCheckPointID(\"2\"), WithStateModifier(func(ctx context.Context, path NodePath, state any) error {\n\t\tstate.(*testStruct).A = \"state\"\n\t\treturn nil\n\t}))\n\tassert.NoError(t, err)\n\tchunk, err := streamResult.Recv()\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"state\", chunk)\n\t_, err = streamResult.Recv()\n\tassert.Equal(t, io.EOF, err)\n}\n\ntype myInterface interface {\n\tA()\n}\n\nfunc TestInterfaceResume(t *testing.T) {\n\tg := NewGraph[myInterface, string]()\n\ttimes := 0\n\tassert.NoError(t, g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input myInterface) (output string, err error) {\n\t\tif times == 0 {\n\t\t\ttimes++\n\t\t\treturn \"\", NewInterruptAndRerunErr(\"test extra\")\n\t\t}\n\t\treturn \"success\", nil\n\t})))\n\tassert.NoError(t, g.AddEdge(START, \"1\"))\n\tassert.NoError(t, g.AddEdge(\"1\", END))\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, nil, WithCheckPointID(\"1\"))\n\tinfo, existed := ExtractInterruptInfo(err)\n\tassert.True(t, existed)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\tresult, err := r.Invoke(ctx, nil, WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"success\", result)\n}\n\nfunc TestEarlyFailCallback(t *testing.T) {\n\tg := NewGraph[string, string]()\n\tassert.NoError(t, g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t})))\n\tassert.NoError(t, g.AddEdge(START, \"1\"))\n\tassert.NoError(t, g.AddEdge(\"1\", END))\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor))\n\tassert.NoError(t, err)\n\ttGCB := &testGraphCallback{}\n\t_, _ = r.Invoke(ctx, \"\", WithCallbacks(tGCB), WithRuntimeMaxSteps(1))\n\tassert.Equal(t, 1, tGCB.onStartTimes)\n\tassert.Equal(t, 1, tGCB.onErrorTimes)\n\tassert.Equal(t, 0, tGCB.onEndTimes)\n}\n\nfunc TestGraphStartInterrupt(t *testing.T) {\n\tsubG := NewGraph[string, string]()\n\t_ = subG.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"sub1\", nil\n\t}))\n\t_ = subG.AddEdge(START, \"1\")\n\t_ = subG.AddEdge(\"1\", END)\n\n\tg := NewGraph[string, string]()\n\t_ = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\t_ = g.AddGraphNode(\"2\", subG, WithGraphCompileOptions(WithInterruptBeforeNodes([]string{\"1\"})))\n\t_ = g.AddEdge(START, \"1\")\n\t_ = g.AddEdge(\"1\", \"2\")\n\t_ = g.AddEdge(\"2\", END)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"1\"))\n\tinfo, existed := ExtractInterruptInfo(err)\n\tassert.True(t, existed)\n\tassert.Equal(t, []string{\"1\"}, info.SubGraphs[\"2\"].BeforeNodes)\n\tresult, err := r.Invoke(ctx, \"\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input1sub1\", result)\n}\n\nfunc TestWithForceNewRun(t *testing.T) {\n\tg := NewGraph[string, string]()\n\t_ = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\t_ = g.AddEdge(START, \"1\")\n\t_ = g.AddEdge(\"1\", END)\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(&failStore{t: t}))\n\tassert.NoError(t, err)\n\tresult, err := r.Invoke(ctx, \"input\", WithCheckPointID(\"1\"), WithForceNewRun())\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input1\", result)\n}\n\ntype failStore struct {\n\tt *testing.T\n}\n\nfunc (f *failStore) Get(_ context.Context, _ string) ([]byte, bool, error) {\n\tf.t.Fatalf(\"cannot call store\")\n\treturn nil, false, errors.New(\"fail\")\n}\n\nfunc (f *failStore) Set(_ context.Context, _ string, _ []byte) error {\n\tf.t.Fatalf(\"cannot call store\")\n\treturn errors.New(\"fail\")\n}\n\nfunc TestPreHandlerInterrupt(t *testing.T) {\n\ttype state struct{}\n\tassert.NoError(t, serialization.GenericRegister[state](\"_eino_TestPreHandlerInterrupt_state\"))\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) state {\n\t\treturn state{}\n\t}))\n\ttimes := 0\n\t_ = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, state state) (string, error) {\n\t\tif times == 0 {\n\t\t\ttimes++\n\t\t\treturn \"\", NewInterruptAndRerunErr(\"\")\n\t\t}\n\t\treturn in, nil\n\t}))\n\t_ = g.AddEdge(START, \"1\")\n\t_ = g.AddEdge(\"1\", END)\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\t_, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"1\"))\n\tinfo, existed := ExtractInterruptInfo(err)\n\tassert.True(t, existed)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\tresult, err := r.Invoke(ctx, \"\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"1\", result)\n}\n\nfunc TestCancelInterrupt(t *testing.T) {\n\tg := NewGraph[string, string]()\n\t_ = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\ttime.Sleep(3 * time.Second)\n\t\treturn input + \"1\", nil\n\t}))\n\t_ = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"2\", nil\n\t}))\n\t_ = g.AddEdge(START, \"1\")\n\t_ = g.AddEdge(\"1\", \"2\")\n\t_ = g.AddEdge(\"2\", END)\n\tctx := context.Background()\n\n\t// pregel\n\tr, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\t// interrupt after nodes\n\tcanceledCtx, cancel := WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(500 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(time.Hour))\n\t}()\n\t_, err = r.Invoke(canceledCtx, \"input\", WithCheckPointID(\"1\"))\n\tassert.Error(t, err)\n\tinfo, success := ExtractInterruptInfo(err)\n\tassert.True(t, success)\n\tassert.Equal(t, []string{\"1\"}, info.AfterNodes)\n\tresult, err := r.Invoke(ctx, \"input\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input12\", result)\n\t// infinite timeout\n\tcanceledCtx, cancel = WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(500 * time.Millisecond)\n\t\tcancel()\n\t}()\n\t_, err = r.Invoke(canceledCtx, \"input\", WithCheckPointID(\"2\"))\n\tassert.Error(t, err)\n\tinfo, success = ExtractInterruptInfo(err)\n\tassert.True(t, success)\n\tassert.Equal(t, []string{\"1\"}, info.AfterNodes)\n\tresult, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"2\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input12\", result)\n\n\t// interrupt rerun nodes - with auto-enabled PersistRerunInput, input is preserved\n\tcanceledCtx, cancel = WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(500 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(0))\n\t}()\n\t_, err = r.Invoke(canceledCtx, \"input\", WithCheckPointID(\"3\"))\n\tassert.Error(t, err)\n\tinfo, success = ExtractInterruptInfo(err)\n\tassert.True(t, success)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\tresult, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"3\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input12\", result)\n\n\t// dag\n\tg = NewGraph[string, string]()\n\t_ = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\ttime.Sleep(3 * time.Second)\n\t\treturn input + \"1\", nil\n\t}))\n\t_ = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"2\", nil\n\t}))\n\t_ = g.AddEdge(START, \"1\")\n\t_ = g.AddEdge(\"1\", \"2\")\n\t_ = g.AddEdge(\"2\", END)\n\tr, err = g.Compile(ctx, WithNodeTriggerMode(AllPredecessor), WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\t// interrupt after nodes\n\tcanceledCtx, cancel = WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(500 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(time.Hour))\n\t}()\n\t_, err = r.Invoke(canceledCtx, \"input\", WithCheckPointID(\"1\"))\n\tassert.Error(t, err)\n\tinfo, success = ExtractInterruptInfo(err)\n\tassert.True(t, success)\n\tassert.Equal(t, []string{\"1\"}, info.AfterNodes)\n\tresult, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input12\", result)\n\t// infinite timeout\n\tcanceledCtx, cancel = WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(500 * time.Millisecond)\n\t\tcancel()\n\t}()\n\t_, err = r.Invoke(canceledCtx, \"input\", WithCheckPointID(\"2\"))\n\tassert.Error(t, err)\n\tinfo, success = ExtractInterruptInfo(err)\n\tassert.True(t, success)\n\tassert.Equal(t, []string{\"1\"}, info.AfterNodes)\n\tresult, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"2\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input12\", result)\n\n\t// interrupt rerun nodes - with auto-enabled PersistRerunInput, input is preserved\n\tcanceledCtx, cancel = WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(300 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(0))\n\t}()\n\t_, err = r.Invoke(canceledCtx, \"input\", WithCheckPointID(\"3\"))\n\tassert.Error(t, err)\n\tinfo, success = ExtractInterruptInfo(err)\n\tassert.True(t, success)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\tresult, err = r.Invoke(ctx, \"input\", WithCheckPointID(\"3\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input12\", result)\n\n\t// dag multi canceled nodes\n\tgg := NewGraph[string, map[string]any]()\n\t_ = gg.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"1\", nil\n\t}))\n\t_ = gg.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\ttime.Sleep(3 * time.Second)\n\t\treturn input + \"2\", nil\n\t}), WithOutputKey(\"2\"))\n\t_ = gg.AddLambdaNode(\"3\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\ttime.Sleep(3 * time.Second)\n\t\treturn input + \"3\", nil\n\t}), WithOutputKey(\"3\"))\n\t_ = gg.AddLambdaNode(\"4\", InvokableLambda(func(ctx context.Context, input map[string]any) (output map[string]any, err error) {\n\t\treturn input, nil\n\t}))\n\t_ = gg.AddEdge(START, \"1\")\n\t_ = gg.AddEdge(\"1\", \"2\")\n\t_ = gg.AddEdge(\"1\", \"3\")\n\t_ = gg.AddEdge(\"2\", \"4\")\n\t_ = gg.AddEdge(\"3\", \"4\")\n\t_ = gg.AddEdge(\"4\", END)\n\tctx = context.Background()\n\trr, err := gg.Compile(ctx, WithNodeTriggerMode(AllPredecessor), WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\t// interrupt after nodes\n\tcanceledCtx, cancel = WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(500 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(time.Hour))\n\t}()\n\t_, err = rr.Invoke(canceledCtx, \"input\", WithCheckPointID(\"1\"))\n\tassert.Error(t, err)\n\tinfo, success = ExtractInterruptInfo(err)\n\tassert.True(t, success)\n\tassert.Equal(t, 2, len(info.AfterNodes))\n\tresult2, err := rr.Invoke(ctx, \"input\", WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\n\t\t\"2\": \"input12\",\n\t\t\"3\": \"input13\",\n\t}, result2)\n\n\t// interrupt rerun nodes - with auto-enabled PersistRerunInput, input is preserved\n\tcanceledCtx, cancel = WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(500 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(0))\n\t}()\n\t_, err = rr.Invoke(canceledCtx, \"input\", WithCheckPointID(\"2\"))\n\tassert.Error(t, err)\n\tinfo, success = ExtractInterruptInfo(err)\n\tassert.True(t, success)\n\tassert.Equal(t, 2, len(info.RerunNodes))\n\tresult2, err = rr.Invoke(ctx, \"input\", WithCheckPointID(\"2\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\n\t\t\"2\": \"input12\",\n\t\t\"3\": \"input13\",\n\t}, result2)\n}\n\nfunc TestPersistRerunInputNonStream(t *testing.T) {\n\tstore := newInMemoryStore()\n\n\tvar mu sync.Mutex\n\tvar receivedInput string\n\tvar callCount int\n\n\tg := NewGraph[string, string]()\n\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\tmu.Lock()\n\t\tcallCount++\n\t\tcurrentCount := callCount\n\t\treceivedInput = input\n\t\tmu.Unlock()\n\n\t\tif currentCount == 1 {\n\t\t\ttime.Sleep(2 * time.Second)\n\t\t}\n\t\treturn input + \"_processed\", nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx,\n\t\tWithNodeTriggerMode(AllPredecessor),\n\t\tWithCheckPointStore(store),\n\t)\n\tassert.NoError(t, err)\n\n\tcanceledCtx, cancel := WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(0))\n\t}()\n\n\t_, err = r.Invoke(canceledCtx, \"test_input\", WithCheckPointID(\"cp1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\n\tmu.Lock()\n\tassert.Equal(t, \"test_input\", receivedInput)\n\tmu.Unlock()\n\n\tresult, err := r.Invoke(ctx, \"\", WithCheckPointID(\"cp1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"test_input_processed\", result)\n\n\tmu.Lock()\n\tassert.Equal(t, \"test_input\", receivedInput)\n\tassert.Equal(t, 2, callCount)\n\tmu.Unlock()\n}\n\nfunc TestPersistRerunInputStream(t *testing.T) {\n\tstore := newInMemoryStore()\n\n\tvar mu sync.Mutex\n\tvar receivedInput string\n\tvar callCount int\n\n\tg := NewGraph[string, string]()\n\n\terr := g.AddLambdaNode(\"1\", TransformableLambda(func(ctx context.Context, input *schema.StreamReader[string]) (output *schema.StreamReader[string], err error) {\n\t\tmu.Lock()\n\t\tcallCount++\n\t\tcurrentCount := callCount\n\t\tmu.Unlock()\n\n\t\tvar sb string\n\t\tfor {\n\t\t\tchunk, err := input.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tsb += chunk\n\t\t}\n\n\t\tmu.Lock()\n\t\treceivedInput = sb\n\t\tmu.Unlock()\n\n\t\tif currentCount == 1 {\n\t\t\ttime.Sleep(2 * time.Second)\n\t\t}\n\n\t\treturn schema.StreamReaderFromArray([]string{sb + \"_processed\"}), nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx,\n\t\tWithNodeTriggerMode(AllPredecessor),\n\t\tWithCheckPointStore(store),\n\t)\n\tassert.NoError(t, err)\n\n\tinputStream := schema.StreamReaderFromArray([]string{\"chunk1\", \"chunk2\", \"chunk3\"})\n\n\tcanceledCtx, cancel := WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(0))\n\t}()\n\n\t_, err = r.Transform(canceledCtx, inputStream, WithCheckPointID(\"cp1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\n\tmu.Lock()\n\tassert.Equal(t, \"chunk1chunk2chunk3\", receivedInput)\n\tmu.Unlock()\n\n\temptyInputStream := schema.StreamReaderFromArray([]string{})\n\n\tresultStream, err := r.Transform(ctx, emptyInputStream, WithCheckPointID(\"cp1\"))\n\tassert.NoError(t, err)\n\n\tvar result string\n\tfor {\n\t\tchunk, err := resultStream.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tresult += chunk\n\t}\n\n\tassert.Equal(t, \"chunk1chunk2chunk3_processed\", result)\n\n\tmu.Lock()\n\tassert.Equal(t, \"chunk1chunk2chunk3\", receivedInput)\n\tassert.Equal(t, 2, callCount)\n\tmu.Unlock()\n}\n\ntype testPersistRerunInputState struct {\n\tPrefix string\n}\n\nfunc TestPersistRerunInputWithPreHandler(t *testing.T) {\n\tstore := newInMemoryStore()\n\n\tvar mu sync.Mutex\n\tvar receivedInput string\n\tvar callCount int\n\n\tschema.Register[testPersistRerunInputState]()\n\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) *testPersistRerunInputState {\n\t\treturn &testPersistRerunInputState{Prefix: \"prefix_\"}\n\t}))\n\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\tmu.Lock()\n\t\tcallCount++\n\t\tcurrentCount := callCount\n\t\treceivedInput = input\n\t\tmu.Unlock()\n\n\t\tif currentCount == 1 {\n\t\t\ttime.Sleep(2 * time.Second)\n\t\t}\n\t\treturn input + \"_processed\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, s *testPersistRerunInputState) (string, error) {\n\t\treturn s.Prefix + in, nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx,\n\t\tWithNodeTriggerMode(AllPredecessor),\n\t\tWithCheckPointStore(store),\n\t)\n\tassert.NoError(t, err)\n\n\tcanceledCtx, cancel := WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(0))\n\t}()\n\n\t_, err = r.Invoke(canceledCtx, \"test_input\", WithCheckPointID(\"cp1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tif ok {\n\t\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\t}\n\n\tmu.Lock()\n\tassert.Equal(t, \"prefix_test_input\", receivedInput)\n\tmu.Unlock()\n\n\tresult, err := r.Invoke(ctx, \"\", WithCheckPointID(\"cp1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"prefix_test_input_processed\", result)\n\n\tmu.Lock()\n\tassert.Equal(t, \"prefix_test_input\", receivedInput)\n\tassert.Equal(t, 2, callCount)\n\tmu.Unlock()\n}\n\nfunc TestPersistRerunInputBackwardCompatibility(t *testing.T) {\n\tstore := newInMemoryStore()\n\n\tvar receivedInput string\n\tvar callCount int\n\n\tg := NewGraph[string, string]()\n\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\tcallCount++\n\t\treceivedInput = input\n\t\tif len(input) > 0 {\n\t\t\treturn \"\", StatefulInterrupt(ctx, \"interrupt\", input)\n\t\t}\n\n\t\t_, _, restoredInput := GetInterruptState[string](ctx)\n\t\treturn restoredInput + \"_processed\", nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx,\n\t\tWithNodeTriggerMode(AllPredecessor),\n\t\tWithCheckPointStore(store),\n\t)\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"test_input\", WithCheckPointID(\"cp1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, []string{\"1\"}, info.RerunNodes)\n\n\tassert.Equal(t, \"test_input\", receivedInput)\n\n\tresult, err := r.Invoke(ctx, \"\", WithCheckPointID(\"cp1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"test_input_processed\", result)\n\tassert.Equal(t, \"\", receivedInput)\n\tassert.Equal(t, 2, callCount)\n}\n\nfunc TestPersistRerunInputSubGraph(t *testing.T) {\n\tstore := newInMemoryStore()\n\n\tvar mu sync.Mutex\n\tvar receivedInput string\n\tvar callCount int\n\n\tsubG := NewGraph[string, string]()\n\terr := subG.AddLambdaNode(\"sub1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\tmu.Lock()\n\t\tcallCount++\n\t\tcurrentCount := callCount\n\t\treceivedInput = input\n\t\tmu.Unlock()\n\n\t\tif currentCount == 1 {\n\t\t\ttime.Sleep(2 * time.Second)\n\t\t}\n\t\treturn input + \"_sub_processed\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(START, \"sub1\")\n\tassert.NoError(t, err)\n\terr = subG.AddEdge(\"sub1\", END)\n\tassert.NoError(t, err)\n\n\tg := NewGraph[string, string]()\n\terr = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input + \"_main\", nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddGraphNode(\"2\", subG)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx,\n\t\tWithNodeTriggerMode(AllPredecessor),\n\t\tWithCheckPointStore(store),\n\t)\n\tassert.NoError(t, err)\n\n\tcanceledCtx, cancel := WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(0))\n\t}()\n\n\t_, err = r.Invoke(canceledCtx, \"test\", WithCheckPointID(\"cp1\"))\n\tassert.NotNil(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok, \"Expected interrupt error, got: %v\", err)\n\tif len(info.SubGraphs) > 0 {\n\t\tassert.Contains(t, info.SubGraphs, \"2\")\n\t\tsubInfo := info.SubGraphs[\"2\"]\n\t\tassert.Equal(t, []string{\"sub1\"}, subInfo.RerunNodes)\n\t} else {\n\t\tassert.Equal(t, []string{\"2\"}, info.RerunNodes)\n\t}\n\n\tmu.Lock()\n\tassert.Equal(t, \"test_main\", receivedInput)\n\tmu.Unlock()\n\n\tresult, err := r.Invoke(ctx, \"\", WithCheckPointID(\"cp1\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"test_main_sub_processed\", result)\n\n\tmu.Lock()\n\tassert.Equal(t, \"test_main\", receivedInput)\n\tassert.Equal(t, 2, callCount)\n\tmu.Unlock()\n}\n\ntype longRunningToolInput struct {\n\tInput string `json:\"input\"`\n}\n\nfunc TestToolsNodeWithExternalGraphInterrupt(t *testing.T) {\n\tstore := newInMemoryStore()\n\tctx := context.Background()\n\n\tvar mu sync.Mutex\n\tvar callCount int\n\n\tlongRunningToolInfo := &schema.ToolInfo{\n\t\tName: \"long_running_tool\",\n\t\tDesc: \"A tool that takes a long time to run\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"input\": {Type: \"string\", Desc: \"input\"},\n\t\t}),\n\t}\n\n\tlongRunningTool := newCheckpointTestTool(longRunningToolInfo, func(ctx context.Context, in *longRunningToolInput) (string, error) {\n\t\tmu.Lock()\n\t\tcallCount++\n\t\tcurrentCount := callCount\n\t\tmu.Unlock()\n\n\t\tif currentCount == 1 {\n\t\t\ttime.Sleep(2 * time.Second)\n\t\t}\n\t\treturn \"result_\" + in.Input, nil\n\t})\n\n\ttoolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\tTools: []tool.BaseTool{longRunningTool},\n\t})\n\tassert.NoError(t, err)\n\n\tg := NewGraph[*schema.Message, []*schema.Message]()\n\terr = g.AddToolsNode(\"tools\", toolsNode)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(START, \"tools\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"tools\", END)\n\tassert.NoError(t, err)\n\n\tr, err := g.Compile(ctx,\n\t\tWithNodeTriggerMode(AllPredecessor),\n\t\tWithCheckPointStore(store),\n\t)\n\tassert.NoError(t, err)\n\n\tinputMsg := &schema.Message{\n\t\tRole: schema.Assistant,\n\t\tToolCalls: []schema.ToolCall{{\n\t\t\tID:   \"call_1\",\n\t\t\tType: \"function\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"long_running_tool\",\n\t\t\t\tArguments: `{\"input\": \"test\"}`,\n\t\t\t},\n\t\t}},\n\t}\n\n\tcanceledCtx, cancel := WithGraphInterrupt(ctx)\n\tgo func() {\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tcancel(WithGraphInterruptTimeout(0))\n\t}()\n\n\t_, err = r.Invoke(canceledCtx, inputMsg, WithCheckPointID(\"cp1\"))\n\tassert.Error(t, err)\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok, \"Expected interrupt error, got: %v\", err)\n\tif ok {\n\t\tassert.Equal(t, []string{\"tools\"}, info.RerunNodes)\n\t}\n\n\tresult, err := r.Invoke(ctx, &schema.Message{}, WithCheckPointID(\"cp1\"))\n\tassert.NoError(t, err)\n\tassert.Len(t, result, 1)\n\tassert.Equal(t, `\"result_test\"`, result[0].Content)\n\n\tmu.Lock()\n\tassert.Equal(t, 2, callCount)\n\tmu.Unlock()\n}\n\ntype checkpointTestTool[I, O any] struct {\n\tinfo *schema.ToolInfo\n\tfn   func(ctx context.Context, in I) (O, error)\n}\n\nfunc newCheckpointTestTool[I, O any](info *schema.ToolInfo, f func(ctx context.Context, in I) (O, error)) tool.InvokableTool {\n\treturn &checkpointTestTool[I, O]{\n\t\tinfo: info,\n\t\tfn:   f,\n\t}\n}\n\nfunc (f *checkpointTestTool[I, O]) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn f.info, nil\n}\n\nfunc (f *checkpointTestTool[I, O]) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\tt := generic.NewInstance[I]()\n\terr := sonic.UnmarshalString(argumentsInJSON, t)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\to, err := f.fn(ctx, t)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn sonic.MarshalString(o)\n}\n"
  },
  {
    "path": "compose/component_to_graph_node.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n)\n\nfunc toComponentNode[I, O, TOption any](\n\tnode any,\n\tcomponentType component,\n\tinvoke Invoke[I, O, TOption],\n\tstream Stream[I, O, TOption],\n\tcollect Collect[I, O, TOption],\n\ttransform Transform[I, O, TOption],\n\topts ...GraphAddNodeOpt,\n) (*graphNode, *graphAddNodeOpts) {\n\tmeta := parseExecutorInfoFromComponent(componentType, node)\n\tinfo, options := getNodeInfo(opts...)\n\trun := runnableLambda(invoke, stream, collect, transform,\n\t\t!meta.isComponentCallbackEnabled,\n\t)\n\n\tgn := toNode(info, run, nil, meta, node, opts...)\n\n\treturn gn, options\n}\n\nfunc toEmbeddingNode(node embedding.Embedder, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\treturn toComponentNode(\n\t\tnode,\n\t\tcomponents.ComponentOfEmbedding,\n\t\tnode.EmbedStrings,\n\t\tnil,\n\t\tnil,\n\t\tnil,\n\t\topts...)\n}\n\nfunc toRetrieverNode(node retriever.Retriever, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\treturn toComponentNode(\n\t\tnode,\n\t\tcomponents.ComponentOfRetriever,\n\t\tnode.Retrieve,\n\t\tnil,\n\t\tnil,\n\t\tnil,\n\t\topts...)\n}\n\nfunc toLoaderNode(node document.Loader, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\treturn toComponentNode(\n\t\tnode,\n\t\tcomponents.ComponentOfLoader,\n\t\tnode.Load,\n\t\tnil,\n\t\tnil,\n\t\tnil,\n\t\topts...)\n}\n\nfunc toIndexerNode(node indexer.Indexer, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\treturn toComponentNode(\n\t\tnode,\n\t\tcomponents.ComponentOfIndexer,\n\t\tnode.Store,\n\t\tnil,\n\t\tnil,\n\t\tnil,\n\t\topts...)\n}\n\nfunc toChatModelNode(node model.BaseChatModel, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\treturn toComponentNode(\n\t\tnode,\n\t\tcomponents.ComponentOfChatModel,\n\t\tnode.Generate,\n\t\tnode.Stream,\n\t\tnil,\n\t\tnil,\n\t\topts...)\n}\n\nfunc toChatTemplateNode(node prompt.ChatTemplate, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\treturn toComponentNode(\n\t\tnode,\n\t\tcomponents.ComponentOfPrompt,\n\t\tnode.Format,\n\t\tnil,\n\t\tnil,\n\t\tnil,\n\t\topts...)\n}\n\nfunc toDocumentTransformerNode(node document.Transformer, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\treturn toComponentNode(\n\t\tnode,\n\t\tcomponents.ComponentOfTransformer,\n\t\tnode.Transform,\n\t\tnil,\n\t\tnil,\n\t\tnil,\n\t\topts...)\n}\n\nfunc toToolsNode(node *ToolsNode, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\treturn toComponentNode(\n\t\tnode,\n\t\tComponentOfToolsNode,\n\t\tnode.Invoke,\n\t\tnode.Stream,\n\t\tnil,\n\t\tnil,\n\t\topts...)\n}\n\nfunc toLambdaNode(node *Lambda, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\tinfo, options := getNodeInfo(opts...)\n\n\tgn := toNode(info, node.executor, nil, node.executor.meta, node, opts...)\n\n\treturn gn, options\n}\n\nfunc toAnyGraphNode(node AnyGraph, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\tmeta := parseExecutorInfoFromComponent(node.component(), node)\n\tinfo, options := getNodeInfo(opts...)\n\n\tgn := toNode(info, nil, node, meta, node, opts...)\n\n\treturn gn, options\n}\n\nfunc toPassthroughNode(opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {\n\tnode := composablePassthrough()\n\tinfo, options := getNodeInfo(opts...)\n\tgn := toNode(info, node, nil, node.meta, node, opts...)\n\treturn gn, options\n}\n\nfunc toNode(nodeInfo *nodeInfo, executor *composableRunnable, graph AnyGraph,\n\tmeta *executorMeta, instance any, opts ...GraphAddNodeOpt) *graphNode {\n\n\tif meta == nil {\n\t\tmeta = &executorMeta{}\n\t}\n\n\tgn := &graphNode{\n\t\tnodeInfo: nodeInfo,\n\n\t\tcr:           executor,\n\t\tg:            graph,\n\t\texecutorMeta: meta,\n\n\t\tinstance: instance,\n\t\topts:     opts,\n\t}\n\n\treturn gn\n}\n"
  },
  {
    "path": "compose/dag.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"fmt\"\n)\n\nfunc dagChannelBuilder(controlDependencies []string, dataDependencies []string, zeroValue func() any, emptyStream func() streamReader) channel {\n\tdeps := make(map[string]dependencyState, len(controlDependencies))\n\tfor _, dep := range controlDependencies {\n\t\tdeps[dep] = dependencyStateWaiting\n\t}\n\tindirect := make(map[string]bool, len(dataDependencies))\n\tfor _, dep := range dataDependencies {\n\t\tindirect[dep] = false\n\t}\n\n\treturn &dagChannel{\n\t\tValues:              make(map[string]any),\n\t\tControlPredecessors: deps,\n\t\tDataPredecessors:    indirect,\n\t\tzeroValue:           zeroValue,\n\t\temptyStream:         emptyStream,\n\t}\n}\n\ntype dependencyState uint8\n\nconst (\n\tdependencyStateWaiting dependencyState = iota\n\tdependencyStateReady\n\tdependencyStateSkipped\n)\n\ntype dagChannel struct {\n\tzeroValue   func() any\n\temptyStream func() streamReader\n\n\tControlPredecessors map[string]dependencyState\n\tValues              map[string]any\n\tDataPredecessors    map[string]bool // if all dependencies have been skipped, indirect dependencies won't effect.\n\tSkipped             bool\n\n\tmergeConfig FanInMergeConfig\n}\n\nfunc (ch *dagChannel) setMergeConfig(cfg FanInMergeConfig) {\n\tch.mergeConfig.StreamMergeWithSourceEOF = cfg.StreamMergeWithSourceEOF\n}\n\nfunc (ch *dagChannel) load(c channel) error {\n\tdc, ok := c.(*dagChannel)\n\tif !ok {\n\t\treturn fmt.Errorf(\"load dag channel fail, got %T, want *dagChannel\", c)\n\t}\n\tch.ControlPredecessors = dc.ControlPredecessors\n\tch.DataPredecessors = dc.DataPredecessors\n\tch.Skipped = dc.Skipped\n\tch.Values = dc.Values\n\treturn nil\n}\n\nfunc (ch *dagChannel) reportValues(ins map[string]any) error {\n\tif ch.Skipped {\n\t\treturn nil\n\t}\n\n\tfor k, v := range ins {\n\t\tif _, ok := ch.DataPredecessors[k]; !ok {\n\t\t\tcontinue\n\t\t}\n\t\tch.DataPredecessors[k] = true\n\t\tch.Values[k] = v\n\t}\n\treturn nil\n}\n\nfunc (ch *dagChannel) reportDependencies(dependencies []string) {\n\tif ch.Skipped {\n\t\treturn\n\t}\n\n\tfor _, dep := range dependencies {\n\t\tif _, ok := ch.ControlPredecessors[dep]; ok {\n\t\t\tch.ControlPredecessors[dep] = dependencyStateReady\n\t\t}\n\t}\n\treturn\n}\n\nfunc (ch *dagChannel) reportSkip(keys []string) bool {\n\tfor _, k := range keys {\n\t\tif _, ok := ch.ControlPredecessors[k]; ok {\n\t\t\tch.ControlPredecessors[k] = dependencyStateSkipped\n\t\t}\n\t\tif _, ok := ch.DataPredecessors[k]; ok {\n\t\t\tch.DataPredecessors[k] = true\n\t\t}\n\t}\n\n\tallSkipped := true\n\tfor _, state := range ch.ControlPredecessors {\n\t\tif state != dependencyStateSkipped {\n\t\t\tallSkipped = false\n\t\t\tbreak\n\t\t}\n\t}\n\tch.Skipped = allSkipped\n\n\treturn allSkipped\n}\n\nfunc (ch *dagChannel) get(isStream bool, name string, edgeHandler *edgeHandlerManager) (\n\tany, bool, error) {\n\tif ch.Skipped {\n\t\treturn nil, false, nil\n\t}\n\n\tif len(ch.ControlPredecessors)+len(ch.DataPredecessors) == 0 {\n\t\treturn nil, false, nil\n\t}\n\n\tfor _, state := range ch.ControlPredecessors {\n\t\tif state == dependencyStateWaiting {\n\t\t\treturn nil, false, nil\n\t\t}\n\t}\n\tfor _, ready := range ch.DataPredecessors {\n\t\tif !ready {\n\t\t\treturn nil, false, nil\n\t\t}\n\t}\n\n\tdefer func() {\n\t\tch.Values = make(map[string]any)\n\t\tfor k := range ch.ControlPredecessors {\n\t\t\tch.ControlPredecessors[k] = dependencyStateWaiting\n\t\t}\n\t\tfor k := range ch.DataPredecessors {\n\t\t\tch.DataPredecessors[k] = false\n\t\t}\n\t}()\n\n\tvalueList := make([]any, len(ch.Values))\n\tnames := make([]string, len(ch.Values))\n\ti := 0\n\tfor k, value := range ch.Values {\n\t\tresolvedV, err := edgeHandler.handle(k, name, value, isStream)\n\t\tif err != nil {\n\t\t\treturn nil, false, err\n\t\t}\n\t\tvalueList[i] = resolvedV\n\t\tnames[i] = k\n\t\ti++\n\t}\n\n\tif len(valueList) == 0 {\n\t\tif isStream {\n\t\t\treturn ch.emptyStream(), true, nil\n\t\t}\n\t\treturn ch.zeroValue(), true, nil\n\t}\n\tif len(valueList) == 1 {\n\t\treturn valueList[0], true, nil\n\t}\n\n\tmergeOpts := &mergeOptions{\n\t\tstreamMergeWithSourceEOF: ch.mergeConfig.StreamMergeWithSourceEOF,\n\t\tnames:                    names,\n\t}\n\tv, err := mergeValues(valueList, mergeOpts)\n\tif err != nil {\n\t\treturn nil, false, err\n\t}\n\treturn v, true, nil\n}\n\nfunc (ch *dagChannel) convertValues(fn func(map[string]any) error) error {\n\treturn fn(ch.Values)\n}\n"
  },
  {
    "path": "compose/dag_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"io\"\n\t\"testing\"\n)\n\nfunc TestDAG(t *testing.T) {\n\tvar err error\n\n\tg := NewGraph[string, string]()\n\terr = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}), WithOutputKey(\"1\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}), WithOutputKey(\"2\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = g.AddLambdaNode(\"3\", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) {\n\t\tif _, ok := input[\"1\"]; !ok {\n\t\t\treturn \"\", fmt.Errorf(\"node 1 output fail: %+v\", input)\n\t\t}\n\t\tif _, ok := input[\"2\"]; !ok {\n\t\t\treturn \"\", fmt.Errorf(\"node 2 output fail: %+v\", input)\n\t\t}\n\t\treturn input[\"1\"].(string) + input[\"2\"].(string), nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = g.AddLambdaNode(\"4\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = g.AddLambdaNode(\"5\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = g.AddLambdaNode(\"6\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}), WithOutputKey(\"6\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = g.AddLambdaNode(\"7\", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) {\n\t\tif _, ok := input[\"1\"]; !ok {\n\t\t\treturn \"\", fmt.Errorf(\"7:node 1 output fail: %+v\", input)\n\t\t}\n\t\tif _, ok := input[\"6\"]; !ok {\n\t\t\treturn \"\", fmt.Errorf(\"7:node 6 output fail: %+v\", input)\n\t\t}\n\t\treturn input[\"1\"].(string) + input[\"6\"].(string), nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = g.AddEdge(\"1\", \"3\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"2\", \"3\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"3\", \"4\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"4\", \"5\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"4\", \"6\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"6\", \"7\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"1\", \"7\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = g.AddEdge(START, \"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"7\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\trunner, err := g.Compile(context.Background(), WithNodeTriggerMode(AllPredecessor))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\t// success\n\tctx := context.Background()\n\tout, err := runner.Invoke(ctx, \"hello\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif out != \"hellohellohello\" {\n\t\tt.Fatalf(\"node7 fail\")\n\t}\n\n\tresult, err := runner.Invoke(ctx, \"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif result != \"111\" {\n\t\tt.Fatalf(\"runner invoke fail, output: %s\", result)\n\t}\n\tstreamResult, err := runner.Stream(ctx, \"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tdefer streamResult.Close()\n\tret := \"\"\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tret += chunk\n\t}\n\tif ret != \"111\" {\n\t\tt.Fatalf(\"runner stream fail, output: %s\", ret)\n\t}\n\n\t// loop\n\tgg := NewGraph[string, map[string]any]()\n\terr = gg.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}), WithOutputKey(\"1\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = gg.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) {\n\t\treturn input[\"1\"].(string), nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = gg.AddLambdaNode(\"3\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}), WithOutputKey(\"3\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\terr = gg.AddEdge(\"1\", \"2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = gg.AddEdge(\"2\", \"3\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = gg.AddEdge(\"3\", \"2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = gg.AddEdge(START, \"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = gg.AddEdge(\"3\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\t_, err = gg.compile(ctx, &graphCompileOptions{nodeTriggerMode: AllPredecessor})\n\tif err == nil {\n\t\tt.Fatal(\"cannot validate loop\")\n\t}\n}\n"
  },
  {
    "path": "compose/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package compose provides graph and workflow primitives to build\n// composable, interruptible execution pipelines with callback support.\npackage compose\n"
  },
  {
    "path": "compose/error.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"reflect\"\n\t\"strings\"\n)\n\n// ErrExceedMaxSteps graph will throw this error when the number of steps exceeds the maximum number of steps.\nvar ErrExceedMaxSteps = errors.New(\"exceeds max steps\")\n\nfunc newUnexpectedInputTypeErr(expected reflect.Type, got reflect.Type) error {\n\treturn fmt.Errorf(\"unexpected input type. expected: %v, got: %v\", expected, got)\n}\n\ntype defaultImplAction string\n\nconst (\n\tactionInvokeByStream     defaultImplAction = \"InvokeByStream\"\n\tactionInvokeByCollect    defaultImplAction = \"InvokeByCollect\"\n\tactionInvokeByTransform  defaultImplAction = \"InvokeByTransform\"\n\tactionStreamByInvoke     defaultImplAction = \"StreamByInvoke\"\n\tactionStreamByTransform  defaultImplAction = \"StreamByTransform\"\n\tactionStreamByCollect    defaultImplAction = \"StreamByCollect\"\n\tactionCollectByTransform defaultImplAction = \"CollectByTransform\"\n\tactionCollectByInvoke    defaultImplAction = \"CollectByInvoke\"\n\tactionCollectByStream    defaultImplAction = \"CollectByStream\"\n\tactionTransformByStream  defaultImplAction = \"TransformByStream\"\n\tactionTransformByCollect defaultImplAction = \"TransformByCollect\"\n\tactionTransformByInvoke  defaultImplAction = \"TransformByInvoke\"\n)\n\nfunc newStreamReadError(err error) error {\n\treturn fmt.Errorf(\"failed to read from stream. error: %w\", err)\n}\n\nfunc newGraphRunError(err error) error {\n\treturn &internalError{\n\t\ttyp:       internalErrorTypeGraphRun,\n\t\tnodePath:  NodePath{},\n\t\torigError: err,\n\t}\n}\n\nfunc wrapGraphNodeError(nodeKey string, err error) error {\n\tif ok := isInterruptError(err); ok {\n\t\treturn err\n\t}\n\tvar ie *internalError\n\tok := errors.As(err, &ie)\n\tif !ok {\n\t\treturn &internalError{\n\t\t\ttyp:       internalErrorTypeNodeRun,\n\t\t\tnodePath:  NodePath{path: []string{nodeKey}},\n\t\t\torigError: err,\n\t\t}\n\t}\n\tie.nodePath.path = append([]string{nodeKey}, ie.nodePath.path...)\n\treturn ie\n}\n\ntype internalErrorType string\n\nconst (\n\tinternalErrorTypeNodeRun  = \"NodeRunError\"\n\tinternalErrorTypeGraphRun = \"GraphRunError\"\n)\n\ntype internalError struct {\n\ttyp       internalErrorType\n\tnodePath  NodePath\n\torigError error\n}\n\nfunc (i *internalError) Error() string {\n\tsb := strings.Builder{}\n\tsb.WriteString(string(\"[\" + i.typ + \"] \"))\n\tsb.WriteString(i.origError.Error())\n\tif len(i.nodePath.path) > 0 {\n\t\tsb.WriteString(\"\\n------------------------\\n\")\n\t\tsb.WriteString(\"node path: [\")\n\t\tfor j := 0; j < len(i.nodePath.path)-1; j++ {\n\t\t\tsb.WriteString(i.nodePath.path[j] + \", \")\n\t\t}\n\t\tsb.WriteString(i.nodePath.path[len(i.nodePath.path)-1])\n\t\tsb.WriteString(\"]\")\n\t}\n\tsb.WriteString(\"\")\n\treturn sb.String()\n}\n\nfunc (i *internalError) Unwrap() error {\n\treturn i.origError\n}\n"
  },
  {
    "path": "compose/error_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestCommonError(t *testing.T) {\n\tg := NewGraph[string, string]()\n\tassert.NoError(t, g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"\", errors.New(\"my error\")\n\t})))\n\tassert.NoError(t, g.AddEdge(START, \"1\"))\n\tassert.NoError(t, g.AddEdge(\"1\", END))\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx)\n\tassert.NoError(t, err)\n\n\t// node error\n\t_, err = r.Invoke(ctx, \"input\")\n\tvar ie *internalError\n\n\tassert.True(t, errors.As(err, &ie))\n\tassert.Equal(t, \"my error\", ie.origError.Error())\n\n\t// wrapper error\n\tsr, sw := schema.Pipe[string](0)\n\tsw.Close()\n\t_, err = r.Transform(ctx, sr)\n\tassert.True(t, errors.As(err, &ie))\n\tassert.ErrorContains(t, ie.origError, \"stream reader is empty, concat fail\")\n\tassert.Equal(t, []string{\"1\"}, ie.nodePath.path)\n\tprintln(err.Error())\n}\n\nfunc TestSubGraphNodeError(t *testing.T) {\n\tsubG := NewGraph[string, string]()\n\tassert.NoError(t, subG.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"\", errors.New(\"my error\")\n\t})))\n\tassert.NoError(t, subG.AddEdge(START, \"1\"))\n\tassert.NoError(t, subG.AddEdge(\"1\", END))\n\n\tg := NewGraph[string, string]()\n\tassert.NoError(t, g.AddGraphNode(\"a\", subG))\n\tassert.NoError(t, g.AddEdge(START, \"a\"))\n\tassert.NoError(t, g.AddEdge(\"a\", END))\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx)\n\tassert.NoError(t, err)\n\t_, err = r.Invoke(ctx, \"input\")\n\tvar ie *internalError\n\tassert.True(t, errors.As(err, &ie))\n\tassert.Equal(t, \"my error\", ie.origError.Error())\n\tassert.Equal(t, []string{\"a\", \"1\"}, ie.nodePath.path)\n}\n\nfunc TestContextCancelDuringRun(t *testing.T) {\n\t// Create a graph with a long-running node to test context cancellation\n\tg := NewGraph[string, string]()\n\n\t// Add a node that waits for some time (long enough to be cancelled)\n\tassert.NoError(t, g.AddLambdaNode(\"slow_node\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\tselect {\n\t\tcase <-ctx.Done():\n\t\t\t// Return context's error when cancelled\n\t\t\treturn \"\", ctx.Err()\n\t\tcase <-time.After(200 * time.Millisecond):\n\t\t\treturn input + \"_processed\", nil\n\t\t}\n\t})))\n\n\tassert.NoError(t, g.AddEdge(START, \"slow_node\"))\n\tassert.NoError(t, g.AddEdge(\"slow_node\", END))\n\n\t// Create a context that we can cancel\n\tctx, cancel := context.WithCancel(context.Background())\n\n\t// Compile the graph\n\tr, err := g.Compile(ctx)\n\tassert.NoError(t, err)\n\n\t// Run the invoke in a goroutine\n\tresultCh := make(chan error)\n\tgo func() {\n\t\t_, err := r.Invoke(ctx, \"input\")\n\t\tresultCh <- err\n\t}()\n\n\t// Cancel the context after a short delay\n\ttime.Sleep(50 * time.Millisecond)\n\tcancel()\n\n\t// Get the result\n\terr = <-resultCh\n\n\t// Verify the error is related to context cancellation\n\tassert.Error(t, err)\n\n\t// Check error type and content\n\tvar ie *internalError\n\tassert.True(t, errors.As(err, &ie))\n\n\t// Error path should contain the node\n\tassert.Equal(t, []string{\"slow_node\"}, ie.nodePath.path)\n\n\t// Original error should be context.Canceled\n\tassert.ErrorIs(t, ie.origError, context.Canceled)\n\n\t// Test unwrap capability\n\tunwrappedErr := ie.Unwrap()\n\tassert.ErrorIs(t, unwrappedErr, context.Canceled)\n}\n"
  },
  {
    "path": "compose/field_mapping.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"reflect\"\n\t\"runtime/debug\"\n\t\"strings\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype FieldMapping struct {\n\tfromNodeKey string\n\tfrom        string\n\tto          string\n\n\tcustomExtractor func(input any) (any, error)\n}\n\n// String returns the string representation of the FieldMapping.\nfunc (m *FieldMapping) String() string {\n\tvar sb strings.Builder\n\tsb.WriteString(\"[from \")\n\n\tif m.from != \"\" {\n\t\tsb.WriteString(m.from)\n\t\tsb.WriteString(\"(field) of \")\n\t}\n\n\tsb.WriteString(m.fromNodeKey)\n\n\tif m.to != \"\" {\n\t\tsb.WriteString(\" to \")\n\t\tsb.WriteString(m.to)\n\t\tsb.WriteString(\"(field)\")\n\t}\n\n\tsb.WriteString(\"]\")\n\treturn sb.String()\n}\n\n// FromField creates a FieldMapping that maps a single predecessor field to the entire successor input.\n// This is an exclusive mapping - once set, no other field mappings can be added since the successor input\n// has already been fully mapped.\n// Field: either the field of a struct, or the key of a map.\nfunc FromField(from string) *FieldMapping {\n\treturn &FieldMapping{\n\t\tfrom: from,\n\t}\n}\n\n// ToField creates a FieldMapping that maps the entire predecessor output to a single successor field.\n// Field: either the field of a struct, or the key of a map.\nfunc ToField(to string, opts ...FieldMappingOption) *FieldMapping {\n\tfm := &FieldMapping{\n\t\tto: to,\n\t}\n\tfor _, opt := range opts {\n\t\topt(fm)\n\t}\n\treturn fm\n}\n\n// MapFields creates a FieldMapping that maps a single predecessor field to a single successor field.\n// Field: either the field of a struct, or the key of a map.\nfunc MapFields(from, to string) *FieldMapping {\n\treturn &FieldMapping{\n\t\tfrom: from,\n\t\tto:   to,\n\t}\n}\n\nfunc (m *FieldMapping) FromNodeKey() string {\n\treturn m.fromNodeKey\n}\n\nfunc (m *FieldMapping) FromPath() FieldPath {\n\treturn splitFieldPath(m.from)\n}\n\nfunc (m *FieldMapping) ToPath() FieldPath {\n\treturn splitFieldPath(m.to)\n}\n\nfunc (m *FieldMapping) Equals(o *FieldMapping) bool {\n\tif m == nil {\n\t\treturn o == nil\n\t}\n\n\tif o == nil || m.customExtractor != nil || o.customExtractor != nil {\n\t\treturn false\n\t}\n\n\treturn m.from == o.from && m.to == o.to && m.fromNodeKey == o.fromNodeKey\n}\n\n// FieldPath represents a path to a nested field in a struct or map.\n// Each element in the path is either:\n// - a struct field name\n// - a map key\n//\n// Example paths:\n//   - []string{\"user\"}            // top-level field\n//   - []string{\"user\", \"name\"}    // nested struct field\n//   - []string{\"users\", \"admin\"}  // map key access\ntype FieldPath []string\n\nfunc (fp *FieldPath) join() string {\n\treturn strings.Join(*fp, pathSeparator)\n}\n\nfunc splitFieldPath(path string) FieldPath {\n\tp := strings.Split(path, pathSeparator)\n\tif len(p) == 1 && p[0] == \"\" {\n\t\treturn FieldPath{}\n\t}\n\n\treturn p\n}\n\n// pathSeparator is a special character (Unit Separator) used internally to join path elements.\n// This character is chosen because it's extremely unlikely to appear in user-defined field names or map keys.\nconst pathSeparator = \"\\x1F\"\n\n// FromFieldPath creates a FieldMapping that maps a single predecessor field path to the entire successor input.\n// This is an exclusive mapping - once set, no other field mappings can be added since the successor input\n// has already been fully mapped.\n//\n// Example:\n//\n//\t// Maps the 'name' field from nested 'user.profile' to the entire successor input\n//\tFromFieldPath(FieldPath{\"user\", \"profile\", \"name\"})\n//\n// Note: The field path elements must not contain the internal path separator character ('\\x1F').\nfunc FromFieldPath(fromFieldPath FieldPath) *FieldMapping {\n\treturn &FieldMapping{\n\t\tfrom: fromFieldPath.join(),\n\t}\n}\n\n// ToFieldPath creates a FieldMapping that maps the entire predecessor output to a single successor field path.\n//\n// Example:\n//\n//\t// Maps the entire predecessor output to response.data.userName\n//\tToFieldPath(FieldPath{\"response\", \"data\", \"userName\"})\n//\n// Note: The field path elements must not contain the internal path separator character ('\\x1F').\nfunc ToFieldPath(toFieldPath FieldPath, opts ...FieldMappingOption) *FieldMapping {\n\tfm := &FieldMapping{\n\t\tto: toFieldPath.join(),\n\t}\n\tfor _, opt := range opts {\n\t\topt(fm)\n\t}\n\treturn fm\n}\n\n// MapFieldPaths creates a FieldMapping that maps a single predecessor field path to a single successor field path.\n//\n// Example:\n//\n//\t// Maps user.profile.name to response.userName\n//\tMapFieldPaths(\n//\t    FieldPath{\"user\", \"profile\", \"name\"},\n//\t    FieldPath{\"response\", \"userName\"},\n//\t)\n//\n// Note: The field path elements must not contain the internal path separator character ('\\x1F').\nfunc MapFieldPaths(fromFieldPath, toFieldPath FieldPath) *FieldMapping {\n\treturn &FieldMapping{\n\t\tfrom: fromFieldPath.join(),\n\t\tto:   toFieldPath.join(),\n\t}\n}\n\n// FieldMappingOption is a functional option for configuring a FieldMapping.\ntype FieldMappingOption func(*FieldMapping)\n\n// WithCustomExtractor sets a custom extractor function for the FieldMapping.\n// The extractor function is used to extract a value from the 'source' of the FieldMapping.\n// NOTE: if specified in this way, Eino can only check the validity of the field mapping at request time..\nfunc WithCustomExtractor(extractor func(input any) (any, error)) FieldMappingOption {\n\treturn func(m *FieldMapping) {\n\t\tm.customExtractor = extractor\n\t}\n}\n\nfunc (m *FieldMapping) targetPath() FieldPath {\n\treturn splitFieldPath(m.to)\n}\n\nfunc buildFieldMappingConverter[I any]() func(input any) (any, error) {\n\treturn func(input any) (any, error) {\n\t\tin, ok := input.(map[string]any)\n\t\tif !ok {\n\t\t\tpanic(newUnexpectedInputTypeErr(reflect.TypeOf(map[string]any{}), reflect.TypeOf(input)))\n\t\t}\n\n\t\treturn convertTo(in, generic.TypeOf[I]()), nil\n\t}\n}\n\nfunc buildStreamFieldMappingConverter[I any]() func(input streamReader) streamReader {\n\treturn func(input streamReader) streamReader {\n\t\ts, ok := unpackStreamReader[map[string]any](input)\n\t\tif !ok {\n\t\t\tpanic(\"mappingStreamAssign incoming streamReader chunk type not map[string]any\")\n\t\t}\n\n\t\treturn packStreamReader(schema.StreamReaderWithConvert(s, func(v map[string]any) (I, error) {\n\t\t\tt := convertTo(v, generic.TypeOf[I]())\n\t\t\treturn t.(I), nil\n\t\t}))\n\t}\n}\n\nfunc convertTo(mappings map[string]any, typ reflect.Type) any {\n\ttValue := newInstanceByType(typ)\n\tif !tValue.CanAddr() {\n\t\ttValue = newInstanceByType(reflect.PointerTo(typ)).Elem()\n\t}\n\n\tfor mapping, taken := range mappings {\n\t\ttValue = assignOne(tValue, taken, mapping)\n\t}\n\n\treturn tValue.Interface()\n}\n\nfunc assignOne(destValue reflect.Value, taken any, to string) reflect.Value {\n\tif len(to) == 0 { // assign to output directly\n\t\tdestValue.Set(reflect.ValueOf(taken))\n\t\treturn destValue\n\t}\n\n\tvar (\n\t\ttoPaths           = splitFieldPath(to)\n\t\toriginalDestValue = destValue\n\t\tparentMap         reflect.Value\n\t\tparentKey         string\n\t)\n\n\tfor {\n\t\tpath := toPaths[0]\n\t\ttoPaths = toPaths[1:]\n\t\tif len(toPaths) == 0 {\n\t\t\ttoSet := reflect.ValueOf(taken)\n\n\t\t\tif destValue.Type() == reflect.TypeOf((*any)(nil)).Elem() {\n\t\t\t\texistingMap, ok := destValue.Interface().(map[string]any)\n\t\t\t\tif ok {\n\t\t\t\t\tdestValue = reflect.ValueOf(existingMap)\n\t\t\t\t} else {\n\t\t\t\t\tmapValue := reflect.MakeMap(reflect.TypeOf(map[string]any{}))\n\t\t\t\t\tdestValue.Set(mapValue)\n\t\t\t\t\tdestValue = mapValue\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tif destValue.Kind() == reflect.Map {\n\t\t\t\tkey := reflect.ValueOf(path)\n\t\t\t\tkeyType := destValue.Type().Key()\n\t\t\t\tif keyType != strType {\n\t\t\t\t\tkey = key.Convert(keyType)\n\t\t\t\t}\n\n\t\t\t\tif !toSet.IsValid() {\n\t\t\t\t\ttoSet = reflect.Zero(destValue.Type().Elem())\n\t\t\t\t}\n\t\t\t\tdestValue.SetMapIndex(key, toSet)\n\n\t\t\t\tif parentMap.IsValid() {\n\t\t\t\t\tparentMap.SetMapIndex(reflect.ValueOf(parentKey), destValue)\n\t\t\t\t}\n\n\t\t\t\treturn originalDestValue\n\t\t\t}\n\n\t\t\tptrValue := destValue\n\t\t\tfor destValue.Kind() == reflect.Ptr {\n\t\t\t\tdestValue = destValue.Elem()\n\t\t\t}\n\n\t\t\tif !toSet.IsValid() {\n\t\t\t\t// just skip it, because this 'nil' is the zero value of the corresponding struct field\n\t\t\t} else {\n\t\t\t\tfield := destValue.FieldByName(path)\n\t\t\t\tfield.Set(toSet)\n\t\t\t}\n\n\t\t\tif parentMap.IsValid() {\n\t\t\t\tparentMap.SetMapIndex(reflect.ValueOf(parentKey), ptrValue)\n\t\t\t}\n\n\t\t\treturn originalDestValue\n\t\t}\n\n\t\tif destValue.Type() == reflect.TypeOf((*any)(nil)).Elem() {\n\t\t\texistingMap, ok := destValue.Interface().(map[string]any)\n\t\t\tif ok {\n\t\t\t\tdestValue = reflect.ValueOf(existingMap)\n\t\t\t} else {\n\t\t\t\tmapValue := reflect.MakeMap(reflect.TypeOf(map[string]any{}))\n\t\t\t\tdestValue.Set(mapValue)\n\t\t\t\tdestValue = mapValue\n\t\t\t}\n\t\t}\n\n\t\tif destValue.Kind() == reflect.Map {\n\t\t\tkeyValue := reflect.ValueOf(path)\n\t\t\tvalueValue := destValue.MapIndex(keyValue)\n\t\t\tif !valueValue.IsValid() {\n\t\t\t\tvalueValue = newInstanceByType(destValue.Type().Elem())\n\t\t\t\tdestValue.SetMapIndex(keyValue, valueValue)\n\t\t\t}\n\n\t\t\tif parentMap.IsValid() {\n\t\t\t\tparentMap.SetMapIndex(reflect.ValueOf(parentKey), destValue)\n\t\t\t}\n\n\t\t\tparentMap = destValue\n\t\t\tparentKey = path\n\t\t\tdestValue = valueValue\n\n\t\t\tcontinue\n\t\t}\n\n\t\tptrValue := destValue\n\t\tfor destValue.Kind() == reflect.Ptr {\n\t\t\tdestValue = destValue.Elem()\n\t\t}\n\n\t\tfield := destValue.FieldByName(path)\n\t\tinstantiateIfNeeded(field)\n\n\t\tif parentMap.IsValid() {\n\t\t\tparentMap.SetMapIndex(reflect.ValueOf(parentKey), ptrValue)\n\t\t\tparentMap = reflect.Value{}\n\t\t\tparentKey = \"\"\n\t\t}\n\n\t\tdestValue = field\n\t}\n}\n\nfunc instantiateIfNeeded(field reflect.Value) {\n\tif field.Kind() == reflect.Ptr {\n\t\tif field.IsNil() {\n\t\t\tfield.Set(reflect.New(field.Type().Elem()))\n\t\t}\n\t} else if field.Kind() == reflect.Map {\n\t\tif field.IsNil() {\n\t\t\tfield.Set(reflect.MakeMap(field.Type()))\n\t\t}\n\t}\n}\n\nfunc newInstanceByType(typ reflect.Type) reflect.Value {\n\tswitch typ.Kind() {\n\tcase reflect.Map:\n\t\treturn reflect.MakeMap(typ)\n\tcase reflect.Slice, reflect.Array:\n\t\tslice := reflect.New(typ).Elem()\n\t\tslice.Set(reflect.MakeSlice(typ, 0, 0))\n\t\treturn slice\n\tcase reflect.Ptr:\n\t\ttyp = typ.Elem()\n\t\torigin := reflect.New(typ)\n\t\tnested := newInstanceByType(typ)\n\t\torigin.Elem().Set(nested)\n\n\t\treturn origin\n\tdefault:\n\t\treturn reflect.New(typ).Elem()\n\t}\n}\n\nfunc checkAndExtractFromField(fromField string, input reflect.Value) (reflect.Value, error) {\n\tf := input.FieldByName(fromField)\n\tif !f.IsValid() {\n\t\treturn reflect.Value{}, fmt.Errorf(\"field mapping from a struct field, but field not found. field=%v, inputType=%v\", fromField, input.Type())\n\t}\n\n\tif !f.CanInterface() {\n\t\treturn reflect.Value{}, fmt.Errorf(\"field mapping from a struct field, but field not exported. field= %v, inputType=%v\", fromField, input.Type())\n\t}\n\n\treturn f, nil\n}\n\ntype errMapKeyNotFound struct {\n\tmapKey string\n}\n\nfunc (e *errMapKeyNotFound) Error() string {\n\treturn fmt.Sprintf(\"key=%s\", e.mapKey)\n}\n\ntype errInterfaceNotValidForFieldMapping struct {\n\tinterfaceType reflect.Type\n\tactualType    reflect.Type\n}\n\nfunc (e *errInterfaceNotValidForFieldMapping) Error() string {\n\treturn fmt.Sprintf(\"field mapping from an interface type, but actual type is not struct, struct ptr or map. InterfaceType= %v, ActualType= %v\", e.interfaceType, e.actualType)\n}\n\nfunc checkAndExtractFromMapKey(fromMapKey string, input reflect.Value) (reflect.Value, error) {\n\tkey := reflect.ValueOf(fromMapKey)\n\tif input.Type().Key() != strType {\n\t\tkey = key.Convert(input.Type().Key())\n\t}\n\n\tv := input.MapIndex(key)\n\tif !v.IsValid() {\n\t\treturn reflect.Value{}, fmt.Errorf(\"field mapping from a map key, but key not found in input. %w\", &errMapKeyNotFound{mapKey: fromMapKey})\n\t}\n\n\treturn v, nil\n}\n\nfunc checkAndExtractFieldType(paths []string, typ reflect.Type) (extracted reflect.Type, remainingPaths FieldPath, err error) {\n\textracted = typ\n\tfor i, field := range paths {\n\t\tfor extracted.Kind() == reflect.Ptr {\n\t\t\textracted = extracted.Elem()\n\t\t}\n\n\t\tif extracted.Kind() == reflect.Map {\n\t\t\tif !strType.ConvertibleTo(extracted.Key()) {\n\t\t\t\treturn nil, nil, fmt.Errorf(\"type[%v] is not a map with string or string alias key\", extracted)\n\t\t\t}\n\n\t\t\textracted = extracted.Elem()\n\t\t\tcontinue\n\t\t}\n\n\t\tif extracted.Kind() == reflect.Struct {\n\t\t\tf, ok := extracted.FieldByName(field)\n\t\t\tif !ok {\n\t\t\t\treturn nil, nil, fmt.Errorf(\"type[%v] has no field[%s]\", extracted, field)\n\t\t\t}\n\n\t\t\tif !f.IsExported() {\n\t\t\t\treturn nil, nil, fmt.Errorf(\"type[%v] has an unexported field[%s]\", extracted.String(), field)\n\t\t\t}\n\n\t\t\textracted = f.Type\n\t\t\tcontinue\n\t\t}\n\n\t\tif extracted.Kind() == reflect.Interface {\n\t\t\treturn extracted, paths[i:], nil\n\t\t}\n\n\t\treturn nil, nil, fmt.Errorf(\"intermediate type[%v] is not valid\", extracted)\n\t}\n\n\treturn extracted, nil, nil\n}\n\nvar strType = reflect.TypeOf(\"\")\n\nfunc fieldMap(mappings []*FieldMapping, allowMapKeyNotFound bool, uncheckedSourcePaths map[string]FieldPath) func(any) (map[string]any, error) {\n\treturn func(input any) (result map[string]any, err error) {\n\t\tresult = make(map[string]any, len(mappings))\n\t\tvar inputValue reflect.Value\n\tloop:\n\t\tfor _, mapping := range mappings {\n\t\t\tif mapping.customExtractor != nil {\n\t\t\t\tresult[mapping.to], err = mapping.customExtractor(input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tif len(mapping.from) == 0 {\n\t\t\t\tresult[mapping.to] = input\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tfromPath := splitFieldPath(mapping.from)\n\n\t\t\tif !inputValue.IsValid() {\n\t\t\t\tinputValue = reflect.ValueOf(input)\n\t\t\t}\n\n\t\t\tvar (\n\t\t\t\tpathInputValue = inputValue\n\t\t\t\tpathInputType  = inputValue.Type()\n\t\t\t\ttaken          = input\n\t\t\t)\n\n\t\t\tfor i, path := range fromPath {\n\t\t\t\tfor pathInputValue.Kind() == reflect.Ptr {\n\t\t\t\t\tpathInputValue = pathInputValue.Elem()\n\t\t\t\t}\n\n\t\t\t\tif !pathInputValue.IsValid() {\n\t\t\t\t\treturn nil, fmt.Errorf(\"intermediate source value on path=%v is nil for type [%v]\", fromPath[:i+1], pathInputType)\n\t\t\t\t}\n\n\t\t\t\tif pathInputValue.Kind() == reflect.Map && pathInputValue.IsNil() {\n\t\t\t\t\treturn nil, fmt.Errorf(\"intermediate source value on path=%v is nil for map type [%v]\", fromPath[:i+1], pathInputType)\n\t\t\t\t}\n\n\t\t\t\ttaken, pathInputType, err = takeOne(pathInputValue, pathInputType, path)\n\t\t\t\tif err != nil {\n\t\t\t\t\t// we deferred check from Compile time to request time for interface types, so we won't panic here\n\t\t\t\t\tvar interfaceNotValidErr *errInterfaceNotValidForFieldMapping\n\t\t\t\t\tif errors.As(err, &interfaceNotValidErr) {\n\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t}\n\n\t\t\t\t\t// map key not found can only be a request time error, so we won't panic here\n\t\t\t\t\tvar mapKeyNotFoundErr *errMapKeyNotFound\n\t\t\t\t\tif errors.As(err, &mapKeyNotFoundErr) {\n\t\t\t\t\t\tif allowMapKeyNotFound {\n\t\t\t\t\t\t\tcontinue loop\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t}\n\n\t\t\t\t\tif uncheckedSourcePaths != nil {\n\t\t\t\t\t\tuncheckedPath, ok := uncheckedSourcePaths[mapping.from]\n\t\t\t\t\t\tif ok && len(uncheckedPath) >= len(fromPath)-i {\n\t\t\t\t\t\t\t// the err happens on the mapping source path which is unchecked at request time, so we won't panic here\n\t\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\n\t\t\t\t\tpanic(safe.NewPanicErr(err, debug.Stack()))\n\t\t\t\t}\n\n\t\t\t\tif i < len(fromPath)-1 {\n\t\t\t\t\tpathInputValue = reflect.ValueOf(taken)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tresult[mapping.to] = taken\n\t\t}\n\n\t\treturn result, nil\n\t}\n}\n\nfunc streamFieldMap(mappings []*FieldMapping, uncheckedSourcePaths map[string]FieldPath) func(streamReader) streamReader {\n\treturn func(input streamReader) streamReader {\n\t\treturn packStreamReader(schema.StreamReaderWithConvert(input.toAnyStreamReader(), fieldMap(mappings, true, uncheckedSourcePaths)))\n\t}\n}\n\nfunc takeOne(inputValue reflect.Value, inputType reflect.Type, from string) (taken any, takenType reflect.Type, err error) {\n\tvar f reflect.Value\n\tswitch k := inputValue.Kind(); k {\n\tcase reflect.Map:\n\t\tf, err = checkAndExtractFromMapKey(from, inputValue)\n\t\tif err != nil {\n\t\t\treturn nil, nil, err\n\t\t}\n\n\t\treturn f.Interface(), f.Type(), nil\n\tcase reflect.Struct:\n\t\tf, err = checkAndExtractFromField(from, inputValue)\n\t\tif err != nil {\n\t\t\treturn nil, nil, err\n\t\t}\n\n\t\treturn f.Interface(), f.Type(), nil\n\tdefault:\n\t\tif inputType.Kind() == reflect.Interface {\n\t\t\treturn nil, nil, &errInterfaceNotValidForFieldMapping{\n\t\t\t\tinterfaceType: inputType,\n\t\t\t\tactualType:    inputValue.Type(),\n\t\t\t}\n\t\t}\n\n\t\tpanic(\"when take one value from source, value not map or struct, and type not interface\")\n\t}\n}\n\nfunc isFromAll(mappings []*FieldMapping) bool {\n\tfor _, mapping := range mappings {\n\t\tif len(mapping.from) == 0 && mapping.customExtractor == nil {\n\t\t\treturn true\n\t\t}\n\t}\n\treturn false\n}\n\nfunc fromFields(mappings []*FieldMapping) bool {\n\tfor _, mapping := range mappings {\n\t\tif len(mapping.from) == 0 || mapping.customExtractor != nil {\n\t\t\treturn false\n\t\t}\n\t}\n\n\treturn true\n}\n\nfunc isToAll(mappings []*FieldMapping) bool {\n\tfor _, mapping := range mappings {\n\t\tif len(mapping.to) == 0 {\n\t\t\treturn true\n\t\t}\n\t}\n\treturn false\n}\n\nfunc validateStructOrMap(t reflect.Type) bool {\n\tswitch t.Kind() {\n\tcase reflect.Map:\n\t\treturn true\n\tcase reflect.Ptr:\n\t\tt = t.Elem()\n\t\tfallthrough\n\tcase reflect.Struct:\n\t\treturn true\n\tdefault:\n\t\treturn false\n\t}\n}\n\nfunc validateFieldMapping(predecessorType reflect.Type, successorType reflect.Type, mappings []*FieldMapping) (\n\t// type checkers that are deferred to request-time\n\ttypeHandler *handlerPair,\n\t// the remaining predecessor field paths that are not checked at compile time because of interface type found\n\tuncheckedSourcePath map[string]FieldPath,\n\terr error) {\n\t// check if mapping is legal\n\tif isFromAll(mappings) && isToAll(mappings) {\n\t\t// unreachable\n\t\tpanic(fmt.Errorf(\"invalid field mappings: from all fields to all, use common edge instead\"))\n\t} else if !isToAll(mappings) && (!validateStructOrMap(successorType) && successorType != reflect.TypeOf((*any)(nil)).Elem()) {\n\t\t// if user has not provided a specific struct type, graph cannot construct any struct in the runtime\n\t\treturn nil, nil, fmt.Errorf(\"static check fail: successor input type should be struct or map, actual: %v\", successorType)\n\t} else if fromFields(mappings) && !validateStructOrMap(predecessorType) {\n\t\treturn nil, nil, fmt.Errorf(\"static check fail: predecessor output type should be struct or map, actual: %v\", predecessorType)\n\t}\n\n\tvar fieldCheckers map[string]handlerPair\n\n\tfor i := range mappings {\n\t\tmapping := mappings[i]\n\n\t\tsuccessorFieldType, successorRemaining, err := checkAndExtractFieldType(splitFieldPath(mapping.to), successorType)\n\t\tif err != nil {\n\t\t\treturn nil, nil, fmt.Errorf(\"static check failed for mapping %s: %w\", mapping, err)\n\t\t}\n\n\t\tif len(successorRemaining) > 0 {\n\t\t\tif successorFieldType == reflect.TypeOf((*any)(nil)).Elem() {\n\t\t\t\tcontinue // at request time expand this 'any' to 'map[string]any'\n\t\t\t}\n\t\t\treturn nil, nil, fmt.Errorf(\"static check failed for mapping %s, the successor has intermediate interface type %v\", mapping, successorFieldType)\n\t\t}\n\n\t\tif mapping.customExtractor != nil { // custom extractor applies to request-time data, so skip compile-time check\n\t\t\tcontinue\n\t\t}\n\n\t\tpredecessorFieldType, predecessorRemaining, err := checkAndExtractFieldType(splitFieldPath(mapping.from), predecessorType)\n\t\tif err != nil {\n\t\t\treturn nil, nil, fmt.Errorf(\"static check failed for mapping %s: %w\", mapping, err)\n\t\t}\n\n\t\tif len(predecessorRemaining) > 0 {\n\t\t\tif uncheckedSourcePath == nil {\n\t\t\t\tuncheckedSourcePath = make(map[string]FieldPath)\n\t\t\t}\n\t\t\tuncheckedSourcePath[mapping.from] = predecessorRemaining\n\t\t}\n\n\t\tchecker := func(a any) (any, error) {\n\t\t\ttrueInType := reflect.TypeOf(a)\n\t\t\tif trueInType == nil {\n\t\t\t\tswitch successorFieldType.Kind() {\n\t\t\t\tcase reflect.Map, reflect.Slice, reflect.Ptr, reflect.Interface:\n\t\t\t\tdefault:\n\t\t\t\t\treturn nil, fmt.Errorf(\"runtime check failed for mapping %s, field[%v]-[%v] is absolutely not assignable\", mapping, trueInType, successorFieldType)\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tif !trueInType.AssignableTo(successorFieldType) {\n\t\t\t\t\treturn nil, fmt.Errorf(\"runtime check failed for mapping %s, field[%v]-[%v] is absolutely not assignable\", mapping, trueInType, successorFieldType)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\treturn a, nil\n\t\t}\n\n\t\tif len(predecessorRemaining) > 0 {\n\t\t\t// can't check if types match at compile time, because there is interface type at some point along the source path. Defer to request time\n\t\t\tif fieldCheckers == nil {\n\t\t\t\tfieldCheckers = make(map[string]handlerPair)\n\t\t\t}\n\t\t\tfieldCheckers[mapping.to] = handlerPair{\n\t\t\t\tinvoke: checker,\n\t\t\t\ttransform: func(input streamReader) streamReader {\n\t\t\t\t\treturn packStreamReader(schema.StreamReaderWithConvert(input.toAnyStreamReader(), checker))\n\t\t\t\t},\n\t\t\t}\n\t\t} else {\n\t\t\tat := checkAssignable(predecessorFieldType, successorFieldType)\n\t\t\tif at == assignableTypeMustNot {\n\t\t\t\treturn nil, nil, fmt.Errorf(\"static check failed for mapping %s, field[%v]-[%v] is absolutely not assignable\", mapping, predecessorFieldType, successorFieldType)\n\t\t\t} else if at == assignableTypeMay {\n\t\t\t\t// can't decide if types match, because the successorFieldType implements predecessorFieldType, which is an interface type\n\t\t\t\tif fieldCheckers == nil {\n\t\t\t\t\tfieldCheckers = make(map[string]handlerPair)\n\t\t\t\t}\n\t\t\t\tfieldCheckers[mapping.to] = handlerPair{\n\t\t\t\t\tinvoke: checker,\n\t\t\t\t\ttransform: func(input streamReader) streamReader {\n\t\t\t\t\t\treturn packStreamReader(schema.StreamReaderWithConvert(input.toAnyStreamReader(), checker))\n\t\t\t\t\t},\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t}\n\n\tif len(fieldCheckers) == 0 {\n\t\treturn nil, uncheckedSourcePath, nil\n\t}\n\n\tchecker := func(value map[string]any) (map[string]any, error) {\n\t\tvar err error\n\t\tfor k, v := range fieldCheckers {\n\t\t\tfor mapping := range value {\n\t\t\t\tif mapping == k {\n\t\t\t\t\tvalue[mapping], err = v.invoke(value[mapping])\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t\treturn value, nil\n\t}\n\treturn &handlerPair{\n\t\tinvoke: func(value any) (any, error) {\n\t\t\treturn checker(value.(map[string]any))\n\t\t},\n\t\ttransform: func(input streamReader) streamReader {\n\t\t\ts, ok := unpackStreamReader[map[string]any](input)\n\t\t\tif !ok {\n\t\t\t\t// impossible\n\t\t\t\tpanic(\"field mapping edge stream value isn't map[string]any\")\n\t\t\t}\n\t\t\treturn packStreamReader(schema.StreamReaderWithConvert(s, checker))\n\t\t},\n\t}, uncheckedSourcePath, nil\n}\n"
  },
  {
    "path": "compose/generic_graph.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\ntype newGraphOptions struct {\n\twithState func(ctx context.Context) any\n\tstateType reflect.Type\n}\n\n// NewGraphOption configures behavior when creating a new graph, such as\n// providing local state generation.\ntype NewGraphOption func(ngo *newGraphOptions)\n\n// WithGenLocalState registers a function to generate per-run local state\n// that can be shared across nodes in the graph.\nfunc WithGenLocalState[S any](gls GenLocalState[S]) NewGraphOption {\n\treturn func(ngo *newGraphOptions) {\n\t\tngo.withState = func(ctx context.Context) any {\n\t\t\treturn gls(ctx)\n\t\t}\n\t\tngo.stateType = generic.TypeOf[S]()\n\t}\n}\n\n// NewGraph create a directed graph that can compose components, lambda, chain, parallel etc.\n// simultaneously provide flexible and multi-granular aspect governance capabilities.\n// I: the input type of graph compiled product\n// O: the output type of graph compiled product\n//\n// To share state between nodes, use WithGenLocalState option:\n//\n//\ttype testState struct {\n//\t\tUserInfo *UserInfo\n//\t\tKVs     map[string]any\n//\t}\n//\n//\tgenStateFunc := func(ctx context.Context) *testState {\n//\t\treturn &testState{}\n//\t}\n//\n//\tgraph := compose.NewGraph[string, string](WithGenLocalState(genStateFunc))\n//\n//\t// you can use WithStatePreHandler and WithStatePostHandler to do something with state\n//\tgraph.AddNode(\"node1\", someNode, compose.WithPreHandler(func(ctx context.Context, in string, state *testState) (string, error) {\n//\t\t// do something with state\n//\t\treturn in, nil\n//\t}), compose.WithPostHandler(func(ctx context.Context, out string, state *testState) (string, error) {\n//\t\t// do something with state\n//\t\treturn out, nil\n//\t}))\nfunc NewGraph[I, O any](opts ...NewGraphOption) *Graph[I, O] {\n\toptions := &newGraphOptions{}\n\tfor _, opt := range opts {\n\t\topt(options)\n\t}\n\n\tg := &Graph[I, O]{\n\t\tnewGraphFromGeneric[I, O](\n\t\t\tComponentOfGraph,\n\t\t\toptions.withState,\n\t\t\toptions.stateType,\n\t\t\topts,\n\t\t),\n\t}\n\n\treturn g\n}\n\n// Graph is a generic graph that can be used to compose components.\n// I: the input type of graph compiled product\n// O: the output type of graph compiled product\ntype Graph[I, O any] struct {\n\t*graph\n}\n\n// AddEdge adds an edge to the graph, edge means a data flow from startNode to endNode.\n// the previous node's output type must be set to the next node's input type.\n// NOTE: startNode and endNode must have been added to the graph before adding edge.\n// e.g.\n//\n//\tgraph.AddNode(\"start_node_key\", compose.NewPassthroughNode())\n//\tgraph.AddNode(\"end_node_key\", compose.NewPassthroughNode())\n//\n//\terr := graph.AddEdge(\"start_node_key\", \"end_node_key\")\nfunc (g *Graph[I, O]) AddEdge(startNode, endNode string) (err error) {\n\treturn g.graph.addEdgeWithMappings(startNode, endNode, false, false)\n}\n\n// Compile take the raw graph and compile it into a form ready to be run.\n// e.g.\n//\n//\tgraph, err := compose.NewGraph[string, string]()\n//\tif err != nil {...}\n//\n//\trunnable, err := graph.Compile(ctx, compose.WithGraphName(\"my_graph\"))\n//\tif err != nil {...}\n//\n//\trunnable.Invoke(ctx, \"input\") // invoke\n//\trunnable.Stream(ctx, \"input\") // stream\n//\trunnable.Collect(ctx, inputReader) // collect\n//\trunnable.Transform(ctx, inputReader) // transform\nfunc (g *Graph[I, O]) Compile(ctx context.Context, opts ...GraphCompileOption) (Runnable[I, O], error) {\n\treturn compileAnyGraph[I, O](ctx, g, opts...)\n}\n\nfunc compileAnyGraph[I, O any](ctx context.Context, g AnyGraph, opts ...GraphCompileOption) (Runnable[I, O], error) {\n\tif len(globalGraphCompileCallbacks) > 0 {\n\t\topts = append([]GraphCompileOption{WithGraphCompileCallbacks(globalGraphCompileCallbacks...)}, opts...)\n\t}\n\toption := newGraphCompileOptions(opts...)\n\n\tcr, err := g.compile(ctx, option)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tcr.meta = &executorMeta{\n\t\tcomponent:                  g.component(),\n\t\tisComponentCallbackEnabled: true,\n\t\tcomponentImplType:          \"\",\n\t}\n\n\tcr.nodeInfo = &nodeInfo{\n\t\tname: option.graphName,\n\t}\n\n\tctxWrapper := func(ctx context.Context, opts ...Option) context.Context {\n\t\treturn initGraphCallbacks(AppendAddressSegment(ctx, AddressSegmentRunnable, option.graphName), cr.nodeInfo, cr.meta, opts...)\n\t}\n\n\trp, err := toGenericRunnable[I, O](cr, ctxWrapper)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn rp, nil\n}\n"
  },
  {
    "path": "compose/generic_helper.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc newGenericHelper[I, O any]() *genericHelper {\n\treturn &genericHelper{\n\t\tinputStreamFilter:  defaultStreamMapFilter[I],\n\t\toutputStreamFilter: defaultStreamMapFilter[O],\n\t\tinputConverter: handlerPair{\n\t\t\tinvoke:    defaultValueChecker[I],\n\t\t\ttransform: defaultStreamConverter[I],\n\t\t},\n\t\toutputConverter: handlerPair{\n\t\t\tinvoke:    defaultValueChecker[O],\n\t\t\ttransform: defaultStreamConverter[O],\n\t\t},\n\t\tinputFieldMappingConverter: handlerPair{\n\t\t\tinvoke:    buildFieldMappingConverter[I](),\n\t\t\ttransform: buildStreamFieldMappingConverter[I](),\n\t\t},\n\t\toutputFieldMappingConverter: handlerPair{\n\t\t\tinvoke:    buildFieldMappingConverter[O](),\n\t\t\ttransform: buildStreamFieldMappingConverter[O](),\n\t\t},\n\t\tinputStreamConvertPair:  defaultStreamConvertPair[I](),\n\t\toutputStreamConvertPair: defaultStreamConvertPair[O](),\n\t\tinputZeroValue:          zeroValueFromGeneric[I],\n\t\toutputZeroValue:         zeroValueFromGeneric[O],\n\t\tinputEmptyStream:        emptyStreamFromGeneric[I],\n\t\toutputEmptyStream:       emptyStreamFromGeneric[O],\n\t}\n}\n\ntype genericHelper struct {\n\t// when set input key, use this method to convert input from map[string]any to T\n\tinputStreamFilter, outputStreamFilter streamMapFilter\n\t// when predecessor's output is assignableTypeMay to current node's input, validate and convert(if needed) types using the following two methods\n\tinputConverter, outputConverter handlerPair\n\t// when current node enable field mapping, convert map input to expected struct using the following two methods\n\tinputFieldMappingConverter, outputFieldMappingConverter handlerPair\n\t// can convert input/output from stream to non-stream or non-stream to stream, used for checkpoint\n\tinputStreamConvertPair, outputStreamConvertPair streamConvertPair\n\n\tinputZeroValue, outputZeroValue     func() any\n\tinputEmptyStream, outputEmptyStream func() streamReader\n}\n\nfunc (g *genericHelper) forMapInput() *genericHelper {\n\treturn &genericHelper{\n\t\toutputStreamFilter:          g.outputStreamFilter,\n\t\toutputConverter:             g.outputConverter,\n\t\toutputFieldMappingConverter: g.outputFieldMappingConverter,\n\t\toutputStreamConvertPair:     g.outputStreamConvertPair,\n\t\toutputZeroValue:             g.outputZeroValue,\n\t\toutputEmptyStream:           g.outputEmptyStream,\n\n\t\tinputStreamFilter: defaultStreamMapFilter[map[string]any],\n\t\tinputConverter: handlerPair{\n\t\t\tinvoke:    defaultValueChecker[map[string]any],\n\t\t\ttransform: defaultStreamConverter[map[string]any],\n\t\t},\n\t\tinputFieldMappingConverter: handlerPair{\n\t\t\tinvoke:    buildFieldMappingConverter[map[string]any](),\n\t\t\ttransform: buildStreamFieldMappingConverter[map[string]any](),\n\t\t},\n\t\tinputStreamConvertPair: defaultStreamConvertPair[map[string]any](),\n\t\tinputZeroValue:         zeroValueFromGeneric[map[string]any],\n\t\tinputEmptyStream:       emptyStreamFromGeneric[map[string]any],\n\t}\n}\n\nfunc (g *genericHelper) forMapOutput() *genericHelper {\n\treturn &genericHelper{\n\t\tinputStreamFilter:          g.inputStreamFilter,\n\t\tinputConverter:             g.inputConverter,\n\t\tinputFieldMappingConverter: g.inputFieldMappingConverter,\n\t\tinputStreamConvertPair:     g.inputStreamConvertPair,\n\t\tinputZeroValue:             g.inputZeroValue,\n\t\tinputEmptyStream:           g.inputEmptyStream,\n\n\t\toutputStreamFilter: defaultStreamMapFilter[map[string]any],\n\t\toutputConverter: handlerPair{\n\t\t\tinvoke:    defaultValueChecker[map[string]any],\n\t\t\ttransform: defaultStreamConverter[map[string]any],\n\t\t},\n\t\toutputFieldMappingConverter: handlerPair{\n\t\t\tinvoke:    buildFieldMappingConverter[map[string]any](),\n\t\t\ttransform: buildStreamFieldMappingConverter[map[string]any](),\n\t\t},\n\t\toutputStreamConvertPair: defaultStreamConvertPair[map[string]any](),\n\t\toutputZeroValue:         zeroValueFromGeneric[map[string]any],\n\t\toutputEmptyStream:       emptyStreamFromGeneric[map[string]any],\n\t}\n}\n\nfunc (g *genericHelper) forPredecessorPassthrough() *genericHelper {\n\treturn &genericHelper{\n\t\tinputStreamFilter:           g.inputStreamFilter,\n\t\toutputStreamFilter:          g.inputStreamFilter,\n\t\tinputConverter:              g.inputConverter,\n\t\toutputConverter:             g.inputConverter,\n\t\tinputFieldMappingConverter:  g.inputFieldMappingConverter,\n\t\toutputFieldMappingConverter: g.inputFieldMappingConverter,\n\t\tinputStreamConvertPair:      g.inputStreamConvertPair,\n\t\toutputStreamConvertPair:     g.inputStreamConvertPair,\n\t\tinputZeroValue:              g.inputZeroValue,\n\t\toutputZeroValue:             g.inputZeroValue,\n\t\tinputEmptyStream:            g.inputEmptyStream,\n\t\toutputEmptyStream:           g.inputEmptyStream,\n\t}\n}\n\nfunc (g *genericHelper) forSuccessorPassthrough() *genericHelper {\n\treturn &genericHelper{\n\t\tinputStreamFilter:           g.outputStreamFilter,\n\t\toutputStreamFilter:          g.outputStreamFilter,\n\t\tinputConverter:              g.outputConverter,\n\t\toutputConverter:             g.outputConverter,\n\t\tinputFieldMappingConverter:  g.outputFieldMappingConverter,\n\t\toutputFieldMappingConverter: g.outputFieldMappingConverter,\n\t\tinputStreamConvertPair:      g.outputStreamConvertPair,\n\t\toutputStreamConvertPair:     g.outputStreamConvertPair,\n\t\tinputZeroValue:              g.outputZeroValue,\n\t\toutputZeroValue:             g.outputZeroValue,\n\t\tinputEmptyStream:            g.outputEmptyStream,\n\t\toutputEmptyStream:           g.outputEmptyStream,\n\t}\n}\n\ntype streamMapFilter func(key string, isr streamReader) (streamReader, bool)\n\ntype valueHandler func(value any) (any, error)\ntype streamHandler func(streamReader) streamReader\n\ntype handlerPair struct {\n\tinvoke    valueHandler\n\ttransform streamHandler\n}\n\ntype streamConvertPair struct {\n\tconcatStream  func(sr streamReader) (any, error)\n\trestoreStream func(any) (streamReader, error)\n}\n\nfunc defaultStreamConvertPair[T any]() streamConvertPair {\n\tvar t T\n\treturn streamConvertPair{\n\t\tconcatStream: func(sr streamReader) (any, error) {\n\t\t\ttsr, ok := unpackStreamReader[T](sr)\n\t\t\tif !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"cannot convert sr to streamReader[%T]\", t)\n\t\t\t}\n\t\t\tvalue, err := concatStreamReader(tsr)\n\t\t\tif err != nil {\n\t\t\t\tif errors.Is(err, emptyStreamConcatErr) {\n\t\t\t\t\treturn nil, nil\n\t\t\t\t}\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\treturn value, nil\n\t\t},\n\t\trestoreStream: func(a any) (streamReader, error) {\n\t\t\tif a == nil {\n\t\t\t\treturn packStreamReader(schema.StreamReaderFromArray([]T{})), nil\n\t\t\t}\n\t\t\tvalue, ok := a.(T)\n\t\t\tif !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"cannot convert value[%T] to streamReader[%T]\", a, t)\n\t\t\t}\n\t\t\treturn packStreamReader(schema.StreamReaderFromArray([]T{value})), nil\n\t\t},\n\t}\n}\n\nfunc defaultStreamMapFilter[T any](key string, isr streamReader) (streamReader, bool) {\n\tsr, ok := unpackStreamReader[map[string]any](isr)\n\tif !ok {\n\t\treturn nil, false\n\t}\n\n\tcvt := func(m map[string]any) (T, error) {\n\t\tvar t T\n\t\tv, ok_ := m[key]\n\t\tif !ok_ {\n\t\t\treturn t, schema.ErrNoValue\n\t\t}\n\t\tvv, ok_ := v.(T)\n\t\tif !ok_ {\n\t\t\treturn t, fmt.Errorf(\n\t\t\t\t\"[defaultStreamMapFilter]fail, key[%s]'s value type[%s] isn't expected type[%s]\",\n\t\t\t\tkey, reflect.TypeOf(v).String(),\n\t\t\t\tgeneric.TypeOf[T]().String())\n\t\t}\n\t\treturn vv, nil\n\t}\n\n\tret := schema.StreamReaderWithConvert[map[string]any, T](sr, cvt)\n\n\treturn packStreamReader(ret), true\n}\n\nfunc defaultStreamConverter[T any](reader streamReader) streamReader {\n\treturn packStreamReader(schema.StreamReaderWithConvert(reader.toAnyStreamReader(), func(v any) (T, error) {\n\t\tvv, ok := v.(T)\n\t\tif !ok {\n\t\t\tvar t T\n\t\t\treturn t, fmt.Errorf(\"runtime type check fail, expected type: %T, actual type: %T\", t, v)\n\t\t}\n\t\treturn vv, nil\n\t}))\n}\n\nfunc defaultValueChecker[T any](v any) (any, error) {\n\tnValue, ok := v.(T)\n\tif !ok {\n\t\tvar t T\n\t\treturn nil, fmt.Errorf(\"runtime type check fail, expected type: %T, actual type: %T\", t, v)\n\t}\n\treturn nValue, nil\n}\n\nfunc zeroValueFromGeneric[T any]() any {\n\tvar t T\n\treturn t\n}\n\nfunc emptyStreamFromGeneric[T any]() streamReader {\n\tvar t T\n\tsr, sw := schema.Pipe[T](1)\n\tsw.Send(t, nil)\n\tsw.Close()\n\treturn packStreamReader(sr)\n}\n"
  },
  {
    "path": "compose/graph.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"reflect\"\n\t\"strings\"\n\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/internal/gmap\"\n)\n\n// START is the start node of the graph. You can add your first edge with START.\nconst START = \"start\"\n\n// END is the end node of the graph. You can add your last edge with END.\nconst END = \"end\"\n\n// graphRunType is a custom type used to control the running mode of the graph.\ntype graphRunType string\n\nconst (\n\t// runTypePregel is a running mode of the graph that is suitable for large-scale graph processing tasks. Can have cycles in graph. Compatible with NodeTriggerType.AnyPredecessor.\n\trunTypePregel graphRunType = \"Pregel\"\n\t// runTypeDAG is a running mode of the graph that represents the graph as a directed acyclic graph, suitable for tasks that can be represented as a directed acyclic graph. Compatible with NodeTriggerType.AllPredecessor.\n\trunTypeDAG graphRunType = \"DAG\"\n)\n\n// String returns the string representation of the graph run type.\nfunc (g graphRunType) String() string {\n\treturn string(g)\n}\n\ntype graph struct {\n\tnodes        map[string]*graphNode\n\tcontrolEdges map[string][]string\n\tdataEdges    map[string][]string\n\tbranches     map[string][]*GraphBranch\n\tstartNodes   []string\n\tendNodes     []string\n\n\ttoValidateMap map[string][]struct {\n\t\tendNode  string\n\t\tmappings []*FieldMapping\n\t}\n\n\tstateType      reflect.Type\n\tstateGenerator func(ctx context.Context) any\n\tnewOpts        []NewGraphOption\n\n\texpectedInputType, expectedOutputType reflect.Type\n\n\t*genericHelper\n\n\tfieldMappingRecords map[string][]*FieldMapping\n\n\tbuildError error\n\n\tcmp component\n\n\tcompiled bool\n\n\thandlerOnEdges   map[string]map[string][]handlerPair\n\thandlerPreNode   map[string][]handlerPair\n\thandlerPreBranch map[string][][]handlerPair\n}\n\ntype newGraphConfig struct {\n\tinputType, outputType reflect.Type\n\tgh                    *genericHelper\n\tcmp                   component\n\tstateType             reflect.Type\n\tstateGenerator        func(ctx context.Context) any\n\tnewOpts               []NewGraphOption\n}\n\nfunc newGraphFromGeneric[I, O any](\n\tcmp component,\n\tstateGenerator func(ctx context.Context) any,\n\tstateType reflect.Type,\n\topts []NewGraphOption,\n) *graph {\n\treturn newGraph(&newGraphConfig{\n\t\tinputType:      generic.TypeOf[I](),\n\t\toutputType:     generic.TypeOf[O](),\n\t\tgh:             newGenericHelper[I, O](),\n\t\tcmp:            cmp,\n\t\tstateType:      stateType,\n\t\tstateGenerator: stateGenerator,\n\t\tnewOpts:        opts,\n\t})\n}\n\nfunc newGraph(cfg *newGraphConfig) *graph {\n\treturn &graph{\n\t\tnodes:        make(map[string]*graphNode),\n\t\tdataEdges:    make(map[string][]string),\n\t\tcontrolEdges: make(map[string][]string),\n\t\tbranches:     make(map[string][]*GraphBranch),\n\n\t\ttoValidateMap: make(map[string][]struct {\n\t\t\tendNode  string\n\t\t\tmappings []*FieldMapping\n\t\t}),\n\n\t\texpectedInputType:  cfg.inputType,\n\t\texpectedOutputType: cfg.outputType,\n\t\tgenericHelper:      cfg.gh,\n\n\t\tfieldMappingRecords: make(map[string][]*FieldMapping),\n\n\t\tcmp: cfg.cmp,\n\n\t\tstateType:      cfg.stateType,\n\t\tstateGenerator: cfg.stateGenerator,\n\t\tnewOpts:        cfg.newOpts,\n\n\t\thandlerOnEdges:   make(map[string]map[string][]handlerPair),\n\t\thandlerPreNode:   make(map[string][]handlerPair),\n\t\thandlerPreBranch: make(map[string][][]handlerPair),\n\t}\n}\n\nfunc (g *graph) component() component {\n\treturn g.cmp\n}\n\nfunc isChain(cmp component) bool {\n\treturn cmp == ComponentOfChain\n}\n\nfunc isWorkflow(cmp component) bool {\n\treturn cmp == ComponentOfWorkflow\n}\n\n// ErrGraphCompiled is returned when attempting to modify a graph after it has been compiled\nvar ErrGraphCompiled = errors.New(\"graph has been compiled, cannot be modified\")\n\nfunc (g *graph) addNode(key string, node *graphNode, options *graphAddNodeOpts) (err error) {\n\tif g.buildError != nil {\n\t\treturn g.buildError\n\t}\n\n\tif g.compiled {\n\t\treturn ErrGraphCompiled\n\t}\n\n\tdefer func() {\n\t\tif err != nil {\n\t\t\tg.buildError = err\n\t\t}\n\t}()\n\n\tif key == END || key == START {\n\t\treturn fmt.Errorf(\"node '%s' is reserved, cannot add manually\", key)\n\t}\n\n\tif _, ok := g.nodes[key]; ok {\n\t\treturn fmt.Errorf(\"node '%s' already present\", key)\n\t}\n\n\t// check options\n\tif options.needState {\n\t\tif g.stateGenerator == nil {\n\t\t\treturn fmt.Errorf(\"node '%s' needs state but graph state is not enabled\", key)\n\t\t}\n\t}\n\n\tif options.nodeOptions.nodeKey != \"\" {\n\t\tif !isChain(g.cmp) {\n\t\t\treturn errors.New(\"only chain support node key option\")\n\t\t}\n\t}\n\t// end: check options\n\n\t// check pre- / post-handler type\n\tif options.processor != nil {\n\t\tif options.processor.statePreHandler != nil {\n\t\t\t// check state type\n\t\t\tif g.stateType != options.processor.preStateType {\n\t\t\t\treturn fmt.Errorf(\"node[%s]'s pre handler state type[%v] is different from graph[%v]\", key, options.processor.preStateType, g.stateType)\n\t\t\t}\n\t\t\t// check input type\n\t\t\tif node.inputType() == nil && options.processor.statePreHandler.outputType != reflect.TypeOf((*any)(nil)).Elem() {\n\t\t\t\treturn fmt.Errorf(\"passthrough node[%s]'s pre handler type isn't any\", key)\n\t\t\t} else if node.inputType() != nil && node.inputType() != options.processor.statePreHandler.outputType {\n\t\t\t\treturn fmt.Errorf(\"node[%s]'s pre handler type[%v] is different from its input type[%v]\", key, options.processor.statePreHandler.outputType, node.inputType())\n\t\t\t}\n\t\t}\n\t\tif options.processor.statePostHandler != nil {\n\t\t\t// check state type\n\t\t\tif g.stateType != options.processor.postStateType {\n\t\t\t\treturn fmt.Errorf(\"node[%s]'s post handler state type[%v] is different from graph[%v]\", key, options.processor.postStateType, g.stateType)\n\t\t\t}\n\t\t\t// check input type\n\t\t\tif node.outputType() == nil && options.processor.statePostHandler.inputType != reflect.TypeOf((*any)(nil)).Elem() {\n\t\t\t\treturn fmt.Errorf(\"passthrough node[%s]'s post handler type isn't any\", key)\n\t\t\t} else if node.outputType() != nil && node.outputType() != options.processor.statePostHandler.inputType {\n\t\t\t\treturn fmt.Errorf(\"node[%s]'s post handler type[%v] is different from its output type[%v]\", key, options.processor.statePostHandler.inputType, node.outputType())\n\t\t\t}\n\t\t}\n\t}\n\n\tg.nodes[key] = node\n\n\treturn nil\n}\n\nfunc (g *graph) addEdgeWithMappings(startNode, endNode string, noControl bool, noData bool, mappings ...*FieldMapping) (err error) {\n\tif g.buildError != nil {\n\t\treturn g.buildError\n\t}\n\tif g.compiled {\n\t\treturn ErrGraphCompiled\n\t}\n\n\tif noControl && noData {\n\t\treturn fmt.Errorf(\"edge[%s]-[%s] cannot be both noDirectDependency and noDataFlow\", startNode, endNode)\n\t}\n\n\tdefer func() {\n\t\tif err != nil {\n\t\t\tg.buildError = err\n\t\t}\n\t}()\n\tif startNode == END {\n\t\treturn errors.New(\"END cannot be a start node\")\n\t}\n\tif endNode == START {\n\t\treturn errors.New(\"START cannot be an end node\")\n\t}\n\n\tif _, ok := g.nodes[startNode]; !ok && startNode != START {\n\t\treturn fmt.Errorf(\"edge start node '%s' needs to be added to graph first\", startNode)\n\t}\n\tif _, ok := g.nodes[endNode]; !ok && endNode != END {\n\t\treturn fmt.Errorf(\"edge end node '%s' needs to be added to graph first\", endNode)\n\t}\n\n\tif !noControl {\n\t\tfor i := range g.controlEdges[startNode] {\n\t\t\tif g.controlEdges[startNode][i] == endNode {\n\t\t\t\treturn fmt.Errorf(\"control edge[%s]-[%s] have been added yet\", startNode, endNode)\n\t\t\t}\n\t\t}\n\n\t\tg.controlEdges[startNode] = append(g.controlEdges[startNode], endNode)\n\t\tif startNode == START {\n\t\t\tg.startNodes = append(g.startNodes, endNode)\n\t\t}\n\t\tif endNode == END {\n\t\t\tg.endNodes = append(g.endNodes, startNode)\n\t\t}\n\t}\n\tif !noData {\n\t\tfor i := range g.dataEdges[startNode] {\n\t\t\tif g.dataEdges[startNode][i] == endNode {\n\t\t\t\treturn fmt.Errorf(\"data edge[%s]-[%s] have been added yet\", startNode, endNode)\n\t\t\t}\n\t\t}\n\n\t\tg.addToValidateMap(startNode, endNode, mappings)\n\t\terr = g.updateToValidateMap()\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tg.dataEdges[startNode] = append(g.dataEdges[startNode], endNode)\n\t}\n\n\treturn nil\n}\n\n// AddEmbeddingNode adds a node that implements embedding.Embedder.\n// e.g.\n//\n//\tembeddingNode, err := openai.NewEmbedder(ctx, &openai.EmbeddingConfig{\n//\t\tModel: \"text-embedding-3-small\",\n//\t})\n//\n//\tgraph.AddEmbeddingNode(\"embedding_node_key\", embeddingNode)\nfunc (g *graph) AddEmbeddingNode(key string, node embedding.Embedder, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toEmbeddingNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddRetrieverNode adds a node that implements retriever.Retriever.\n// e.g.\n//\n//\tretriever, err := vikingdb.NewRetriever(ctx, &vikingdb.RetrieverConfig{})\n//\n//\tgraph.AddRetrieverNode(\"retriever_node_key\", retrieverNode)\nfunc (g *graph) AddRetrieverNode(key string, node retriever.Retriever, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toRetrieverNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddLoaderNode adds a node that implements document.Loader.\n// e.g.\n//\n//\tloader, err := file.NewLoader(ctx, &file.LoaderConfig{})\n//\n//\tgraph.AddLoaderNode(\"loader_node_key\", loader)\nfunc (g *graph) AddLoaderNode(key string, node document.Loader, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toLoaderNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddIndexerNode adds a node that implements indexer.Indexer.\n// e.g.\n//\n//\tindexer, err := vikingdb.NewIndexer(ctx, &vikingdb.IndexerConfig{})\n//\n//\tgraph.AddIndexerNode(\"indexer_node_key\", indexer)\nfunc (g *graph) AddIndexerNode(key string, node indexer.Indexer, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toIndexerNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddChatModelNode add node that implements model.BaseChatModel.\n// e.g.\n//\n//\tchatModel, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{\n//\t\tModel: \"gpt-4o\",\n//\t})\n//\n//\tgraph.AddChatModelNode(\"chat_model_node_key\", chatModel)\nfunc (g *graph) AddChatModelNode(key string, node model.BaseChatModel, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toChatModelNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddChatTemplateNode add node that implements prompt.ChatTemplate.\n// e.g.\n//\n//\tchatTemplate, err := prompt.FromMessages(schema.FString, &schema.Message{\n//\t\tRole:    schema.System,\n//\t\tContent: \"You are acting as a {role}.\",\n//\t})\n//\n//\tgraph.AddChatTemplateNode(\"chat_template_node_key\", chatTemplate)\nfunc (g *graph) AddChatTemplateNode(key string, node prompt.ChatTemplate, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toChatTemplateNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddToolsNode adds a node that implements tools.ToolsNode.\n// e.g.\n//\n//\ttoolsNode, err := tools.NewToolNode(ctx, &tools.ToolsNodeConfig{})\n//\n//\tgraph.AddToolsNode(\"tools_node_key\", toolsNode)\nfunc (g *graph) AddToolsNode(key string, node *ToolsNode, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toToolsNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddDocumentTransformerNode adds a node that implements document.Transformer.\n// e.g.\n//\n//\tmarkdownSplitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderSplitterConfig{})\n//\n//\tgraph.AddDocumentTransformerNode(\"document_transformer_node_key\", markdownSplitter)\nfunc (g *graph) AddDocumentTransformerNode(key string, node document.Transformer, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toDocumentTransformerNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddLambdaNode add node that implements at least one of Invoke[I, O], Stream[I, O], Collect[I, O], Transform[I, O].\n// due to the lack of supporting method generics, we need to use function generics to generate Lambda run as Runnable[I, O].\n// for Invoke[I, O], use compose.InvokableLambda()\n// for Stream[I, O], use compose.StreamableLambda()\n// for Collect[I, O], use compose.CollectableLambda()\n// for Transform[I, O], use compose.TransformableLambda()\n// for arbitrary combinations of 4 kinds of lambda, use compose.AnyLambda()\nfunc (g *graph) AddLambdaNode(key string, node *Lambda, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toLambdaNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddGraphNode add one kind of Graph[I, O]、Chain[I, O]、StateChain[I, O, S] as a node.\n// for Graph[I, O], comes from NewGraph[I, O]()\n// for Chain[I, O], comes from NewChain[I, O]()\nfunc (g *graph) AddGraphNode(key string, node AnyGraph, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toAnyGraphNode(node, opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddPassthroughNode adds a passthrough node to the graph.\n// mostly used in pregel mode of graph.\n// e.g.\n//\n//\tgraph.AddPassthroughNode(\"passthrough_node_key\")\nfunc (g *graph) AddPassthroughNode(key string, opts ...GraphAddNodeOpt) error {\n\tgNode, options := toPassthroughNode(opts...)\n\treturn g.addNode(key, gNode, options)\n}\n\n// AddBranch adds a branch to the graph.\n// e.g.\n//\n//\tcondition := func(ctx context.Context, in string) (string, error) {\n//\t\treturn \"next_node_key\", nil\n//\t}\n//\tendNodes := map[string]bool{\"path01\": true, \"path02\": true}\n//\tbranch := compose.NewGraphBranch(condition, endNodes)\n//\n//\tgraph.AddBranch(\"start_node_key\", branch)\nfunc (g *graph) AddBranch(startNode string, branch *GraphBranch) (err error) {\n\treturn g.addBranch(startNode, branch, false)\n}\n\nfunc (g *graph) addBranch(startNode string, branch *GraphBranch, skipData bool) (err error) {\n\tif g.buildError != nil {\n\t\treturn g.buildError\n\t}\n\n\tif g.compiled {\n\t\treturn ErrGraphCompiled\n\t}\n\n\tdefer func() {\n\t\tif err != nil {\n\t\t\tg.buildError = err\n\t\t}\n\t}()\n\n\tif startNode == END {\n\t\treturn errors.New(\"END cannot be a start node\")\n\t}\n\n\tif _, ok := g.nodes[startNode]; !ok && startNode != START {\n\t\treturn fmt.Errorf(\"branch start node '%s' needs to be added to graph first\", startNode)\n\t}\n\n\tif _, ok := g.handlerPreBranch[startNode]; !ok {\n\t\tg.handlerPreBranch[startNode] = [][]handlerPair{}\n\t}\n\tbranch.idx = len(g.handlerPreBranch[startNode])\n\n\tif startNode != START && g.nodes[startNode].executorMeta.component == ComponentOfPassthrough {\n\t\tg.nodes[startNode].cr.inputType = branch.inputType\n\t\tg.nodes[startNode].cr.outputType = branch.inputType\n\t\tg.nodes[startNode].cr.genericHelper = branch.genericHelper.forPredecessorPassthrough()\n\t}\n\n\t// check branch condition type\n\tresult := checkAssignable(g.getNodeOutputType(startNode), branch.inputType)\n\tif result == assignableTypeMustNot {\n\t\treturn fmt.Errorf(\"condition's input type[%s] and start node[%s]'s output type[%s] are mismatched\", branch.inputType.String(), startNode, g.getNodeOutputType(startNode).String())\n\t} else if result == assignableTypeMay {\n\t\tg.handlerPreBranch[startNode] = append(g.handlerPreBranch[startNode], []handlerPair{branch.inputConverter})\n\t} else {\n\t\tg.handlerPreBranch[startNode] = append(g.handlerPreBranch[startNode], []handlerPair{})\n\t}\n\n\tif !skipData {\n\t\tfor endNode := range branch.endNodes {\n\t\t\tif _, ok := g.nodes[endNode]; !ok {\n\t\t\t\tif endNode != END {\n\t\t\t\t\treturn fmt.Errorf(\"branch end node '%s' needs to be added to graph first\", endNode)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tg.addToValidateMap(startNode, endNode, nil)\n\t\t\te := g.updateToValidateMap()\n\t\t\tif e != nil {\n\t\t\t\treturn e\n\t\t\t}\n\n\t\t\tif startNode == START {\n\t\t\t\tg.startNodes = append(g.startNodes, endNode)\n\t\t\t}\n\t\t\tif endNode == END {\n\t\t\t\tg.endNodes = append(g.endNodes, startNode)\n\t\t\t}\n\t\t}\n\t} else {\n\t\tfor endNode := range branch.endNodes {\n\t\t\tif startNode == START {\n\t\t\t\tg.startNodes = append(g.startNodes, endNode)\n\t\t\t}\n\t\t\tif endNode == END {\n\t\t\t\tg.endNodes = append(g.endNodes, startNode)\n\t\t\t}\n\t\t}\n\t\tbranch.noDataFlow = true\n\t}\n\n\tg.branches[startNode] = append(g.branches[startNode], branch)\n\n\treturn nil\n}\n\nfunc (g *graph) addToValidateMap(startNode, endNode string, mapping []*FieldMapping) {\n\tg.toValidateMap[startNode] = append(g.toValidateMap[startNode], struct {\n\t\tendNode  string\n\t\tmappings []*FieldMapping\n\t}{endNode: endNode, mappings: mapping})\n}\n\n// updateToValidateMap after update node, check validate map\n// check again if nodes in toValidateMap have been updated. because when there are multiple linked passthrough nodes, in the worst scenario, only one node can be updated at a time.\nfunc (g *graph) updateToValidateMap() error {\n\tvar startNodeOutputType, endNodeInputType reflect.Type\n\tfor {\n\t\thasChanged := false\n\t\tfor startNode := range g.toValidateMap {\n\t\t\tstartNodeOutputType = g.getNodeOutputType(startNode)\n\n\t\t\tfor i := 0; i < len(g.toValidateMap[startNode]); i++ {\n\t\t\t\tendNode := g.toValidateMap[startNode][i]\n\n\t\t\t\tendNodeInputType = g.getNodeInputType(endNode.endNode)\n\t\t\t\tif startNodeOutputType == nil && endNodeInputType == nil {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\n\t\t\t\t// update toValidateMap\n\t\t\t\tg.toValidateMap[startNode] = append(g.toValidateMap[startNode][:i], g.toValidateMap[startNode][i+1:]...)\n\t\t\t\ti--\n\n\t\t\t\thasChanged = true\n\t\t\t\t// assume that START and END type isn't empty\n\t\t\t\tif startNodeOutputType != nil && endNodeInputType == nil {\n\t\t\t\t\tg.nodes[endNode.endNode].cr.inputType = startNodeOutputType\n\t\t\t\t\tg.nodes[endNode.endNode].cr.outputType = g.nodes[endNode.endNode].cr.inputType\n\t\t\t\t\tg.nodes[endNode.endNode].cr.genericHelper = g.getNodeGenericHelper(startNode).forSuccessorPassthrough()\n\t\t\t\t} else if startNodeOutputType == nil /* redundant condition && endNodeInputType != nil */ {\n\t\t\t\t\tg.nodes[startNode].cr.inputType = endNodeInputType\n\t\t\t\t\tg.nodes[startNode].cr.outputType = g.nodes[startNode].cr.inputType\n\t\t\t\t\tg.nodes[startNode].cr.genericHelper = g.getNodeGenericHelper(endNode.endNode).forPredecessorPassthrough()\n\t\t\t\t} else if len(endNode.mappings) == 0 {\n\t\t\t\t\t// common node check\n\t\t\t\t\tresult := checkAssignable(startNodeOutputType, endNodeInputType)\n\t\t\t\t\tif result == assignableTypeMustNot {\n\t\t\t\t\t\treturn fmt.Errorf(\"graph edge[%s]-[%s]: start node's output type[%s] and end node's input type[%s] mismatch\",\n\t\t\t\t\t\t\tstartNode, endNode.endNode, startNodeOutputType.String(), endNodeInputType.String())\n\t\t\t\t\t} else if result == assignableTypeMay {\n\t\t\t\t\t\t// add runtime check edges\n\t\t\t\t\t\tif _, ok := g.handlerOnEdges[startNode]; !ok {\n\t\t\t\t\t\t\tg.handlerOnEdges[startNode] = make(map[string][]handlerPair)\n\t\t\t\t\t\t}\n\t\t\t\t\t\tg.handlerOnEdges[startNode][endNode.endNode] = append(g.handlerOnEdges[startNode][endNode.endNode], g.getNodeGenericHelper(endNode.endNode).inputConverter)\n\t\t\t\t\t}\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\n\t\t\t\tif len(endNode.mappings) > 0 {\n\t\t\t\t\tif _, ok := g.handlerOnEdges[startNode]; !ok {\n\t\t\t\t\t\tg.handlerOnEdges[startNode] = make(map[string][]handlerPair)\n\t\t\t\t\t}\n\t\t\t\t\tg.fieldMappingRecords[endNode.endNode] = append(g.fieldMappingRecords[endNode.endNode], endNode.mappings...)\n\n\t\t\t\t\t// field mapping check\n\t\t\t\t\tchecker, uncheckedSourcePaths, err := validateFieldMapping(g.getNodeOutputType(startNode), g.getNodeInputType(endNode.endNode), endNode.mappings)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn err\n\t\t\t\t\t}\n\n\t\t\t\t\tg.handlerOnEdges[startNode][endNode.endNode] = append(g.handlerOnEdges[startNode][endNode.endNode], handlerPair{\n\t\t\t\t\t\tinvoke: func(value any) (any, error) {\n\t\t\t\t\t\t\treturn fieldMap(endNode.mappings, false, uncheckedSourcePaths)(value)\n\t\t\t\t\t\t},\n\t\t\t\t\t\ttransform: streamFieldMap(endNode.mappings, uncheckedSourcePaths),\n\t\t\t\t\t})\n\n\t\t\t\t\tif checker != nil {\n\t\t\t\t\t\tg.handlerOnEdges[startNode][endNode.endNode] = append(g.handlerOnEdges[startNode][endNode.endNode], *checker)\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t\tif !hasChanged {\n\t\t\tbreak\n\t\t}\n\t}\n\n\treturn nil\n}\n\nfunc (g *graph) getNodeGenericHelper(name string) *genericHelper {\n\tif name == START {\n\t\treturn g.genericHelper.forPredecessorPassthrough()\n\t} else if name == END {\n\t\treturn g.genericHelper.forSuccessorPassthrough()\n\t}\n\treturn g.nodes[name].getGenericHelper()\n}\n\nfunc (g *graph) getNodeInputType(name string) reflect.Type {\n\tif name == START {\n\t\treturn g.inputType()\n\t} else if name == END {\n\t\treturn g.outputType()\n\t}\n\treturn g.nodes[name].inputType()\n}\n\nfunc (g *graph) getNodeOutputType(name string) reflect.Type {\n\tif name == START {\n\t\treturn g.inputType()\n\t} else if name == END {\n\t\treturn g.outputType()\n\t}\n\treturn g.nodes[name].outputType()\n}\n\nfunc (g *graph) inputType() reflect.Type {\n\treturn g.expectedInputType\n}\n\nfunc (g *graph) outputType() reflect.Type {\n\treturn g.expectedOutputType\n}\n\nfunc (g *graph) compile(ctx context.Context, opt *graphCompileOptions) (*composableRunnable, error) {\n\tif g.buildError != nil {\n\t\treturn nil, g.buildError\n\t}\n\n\t// get run type\n\trunType := runTypePregel\n\tcb := pregelChannelBuilder\n\tif isChain(g.cmp) || isWorkflow(g.cmp) {\n\t\tif opt != nil && opt.nodeTriggerMode != \"\" {\n\t\t\treturn nil, errors.New(fmt.Sprintf(\"%s doesn't support node trigger mode option\", g.cmp))\n\t\t}\n\t}\n\tif (opt != nil && opt.nodeTriggerMode == AllPredecessor) || isWorkflow(g.cmp) {\n\t\trunType = runTypeDAG\n\t\tcb = dagChannelBuilder\n\t}\n\n\t// get eager type\n\teager := false\n\tif isWorkflow(g.cmp) || runType == runTypeDAG {\n\t\teager = true\n\t}\n\tif opt != nil && opt.eagerDisabled {\n\t\teager = false\n\t}\n\n\tif len(g.startNodes) == 0 {\n\t\treturn nil, errors.New(\"start node not set\")\n\t}\n\tif len(g.endNodes) == 0 {\n\t\treturn nil, errors.New(\"end node not set\")\n\t}\n\n\t// toValidateMap isn't empty means there are nodes that cannot infer type\n\tfor _, v := range g.toValidateMap {\n\t\tif len(v) > 0 {\n\t\t\treturn nil, fmt.Errorf(\"some node's input or output types cannot be inferred: %v\", g.toValidateMap)\n\t\t}\n\t}\n\n\tfor key := range g.fieldMappingRecords {\n\t\t// not allowed to map multiple fields to the same field\n\t\ttoMap := make(map[string]bool)\n\t\tfor _, mapping := range g.fieldMappingRecords[key] {\n\t\t\tif _, ok := toMap[mapping.to]; ok {\n\t\t\t\treturn nil, fmt.Errorf(\"duplicate mapping target field: %s of node[%s]\", mapping.to, key)\n\t\t\t}\n\t\t\ttoMap[mapping.to] = true\n\t\t}\n\n\t\t// add map to input converter\n\t\tg.handlerPreNode[key] = append(g.handlerPreNode[key], g.getNodeGenericHelper(key).inputFieldMappingConverter)\n\t}\n\n\tkey2SubGraphs := g.beforeChildGraphsCompile(opt)\n\tchanSubscribeTo := make(map[string]*chanCall)\n\tfor name, node := range g.nodes {\n\t\tnode.beforeChildGraphCompile(name, key2SubGraphs)\n\n\t\tr, err := node.compileIfNeeded(ctx)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tchCall := &chanCall{\n\t\t\taction:   r,\n\t\t\twriteTo:  g.dataEdges[name],\n\t\t\tcontrols: g.controlEdges[name],\n\n\t\t\tpreProcessor:  node.nodeInfo.preProcessor,\n\t\t\tpostProcessor: node.nodeInfo.postProcessor,\n\t\t}\n\n\t\tbranches := g.branches[name]\n\t\tif len(branches) > 0 {\n\t\t\tbranchRuns := make([]*GraphBranch, 0, len(branches))\n\t\t\tbranchRuns = append(branchRuns, branches...)\n\n\t\t\tchCall.writeToBranches = branchRuns\n\t\t}\n\n\t\tchanSubscribeTo[name] = chCall\n\t}\n\n\tdataPredecessors := make(map[string][]string)\n\tcontrolPredecessors := make(map[string][]string)\n\tfor start, ends := range g.controlEdges {\n\t\tfor _, end := range ends {\n\t\t\tif _, ok := controlPredecessors[end]; !ok {\n\t\t\t\tcontrolPredecessors[end] = []string{start}\n\t\t\t} else {\n\t\t\t\tcontrolPredecessors[end] = append(controlPredecessors[end], start)\n\t\t\t}\n\t\t}\n\t}\n\tfor start, ends := range g.dataEdges {\n\t\tfor _, end := range ends {\n\t\t\tif _, ok := dataPredecessors[end]; !ok {\n\t\t\t\tdataPredecessors[end] = []string{start}\n\t\t\t} else {\n\t\t\t\tdataPredecessors[end] = append(dataPredecessors[end], start)\n\t\t\t}\n\t\t}\n\t}\n\tfor start, branches := range g.branches {\n\t\tfor _, branch := range branches {\n\t\t\tfor end := range branch.endNodes {\n\t\t\t\tif _, ok := controlPredecessors[end]; !ok {\n\t\t\t\t\tcontrolPredecessors[end] = []string{start}\n\t\t\t\t} else {\n\t\t\t\t\tcontrolPredecessors[end] = append(controlPredecessors[end], start)\n\t\t\t\t}\n\n\t\t\t\tif !branch.noDataFlow {\n\t\t\t\t\tif _, ok := dataPredecessors[end]; !ok {\n\t\t\t\t\t\tdataPredecessors[end] = []string{start}\n\t\t\t\t\t} else {\n\t\t\t\t\t\tdataPredecessors[end] = append(dataPredecessors[end], start)\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tinputChannels := &chanCall{\n\t\twriteTo:         g.dataEdges[START],\n\t\tcontrols:        g.controlEdges[START],\n\t\twriteToBranches: make([]*GraphBranch, len(g.branches[START])),\n\t}\n\tcopy(inputChannels.writeToBranches, g.branches[START])\n\n\tvar mergeConfigs map[string]FanInMergeConfig\n\tif opt != nil {\n\t\tmergeConfigs = opt.mergeConfigs\n\t}\n\tif mergeConfigs == nil {\n\t\tmergeConfigs = make(map[string]FanInMergeConfig)\n\t}\n\n\tr := &runner{\n\t\tchanSubscribeTo:     chanSubscribeTo,\n\t\tcontrolPredecessors: controlPredecessors,\n\t\tdataPredecessors:    dataPredecessors,\n\n\t\tinputChannels: inputChannels,\n\n\t\teager: eager,\n\n\t\tchanBuilder: cb,\n\n\t\tinputType:     g.inputType(),\n\t\toutputType:    g.outputType(),\n\t\tgenericHelper: g.genericHelper,\n\n\t\tpreBranchHandlerManager: &preBranchHandlerManager{h: g.handlerPreBranch},\n\t\tpreNodeHandlerManager:   &preNodeHandlerManager{h: g.handlerPreNode},\n\t\tedgeHandlerManager:      &edgeHandlerManager{h: g.handlerOnEdges},\n\n\t\tmergeConfigs: mergeConfigs,\n\t}\n\n\tsuccessors := make(map[string][]string)\n\tfor ch := range r.chanSubscribeTo {\n\t\tsuccessors[ch] = getSuccessors(r.chanSubscribeTo[ch])\n\t}\n\tr.successors = successors\n\n\tif g.stateGenerator != nil {\n\t\tr.runCtx = func(ctx context.Context) context.Context {\n\t\t\tvar parent *internalState\n\t\t\tif p, ok := ctx.Value(stateKey{}).(*internalState); ok {\n\t\t\t\tparent = p\n\t\t\t}\n\n\t\t\treturn context.WithValue(ctx, stateKey{}, &internalState{\n\t\t\t\tstate:  g.stateGenerator(ctx),\n\t\t\t\tparent: parent,\n\t\t\t})\n\t\t}\n\t}\n\n\tif runType == runTypeDAG {\n\t\terr := validateDAG(r.chanSubscribeTo, controlPredecessors)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tr.dag = true\n\t}\n\n\tif opt != nil {\n\t\tinputPairs := make(map[string]streamConvertPair)\n\t\toutputPairs := make(map[string]streamConvertPair)\n\t\tfor key, c := range r.chanSubscribeTo {\n\t\t\tinputPairs[key] = c.action.inputStreamConvertPair\n\t\t\toutputPairs[key] = c.action.outputStreamConvertPair\n\t\t}\n\t\tinputPairs[END] = r.outputConvertStreamPair\n\t\toutputPairs[START] = r.inputConvertStreamPair\n\t\tr.checkPointer = newCheckPointer(inputPairs, outputPairs, opt.checkPointStore, opt.serializer)\n\n\t\tr.interruptBeforeNodes = opt.interruptBeforeNodes\n\t\tr.interruptAfterNodes = opt.interruptAfterNodes\n\t\tr.options = *opt\n\t}\n\n\t// default options\n\tif r.dag && r.options.maxRunSteps > 0 {\n\t\treturn nil, fmt.Errorf(\"cannot set max run steps in dag mode\")\n\t} else if !r.dag && r.options.maxRunSteps == 0 {\n\t\tr.options.maxRunSteps = len(r.chanSubscribeTo) + 10\n\t}\n\n\tg.compiled = true\n\n\tg.onCompileFinish(ctx, opt, key2SubGraphs)\n\n\treturn r.toComposableRunnable(), nil\n}\n\nfunc getSuccessors(c *chanCall) []string {\n\tret := make([]string, len(c.writeTo))\n\tcopy(ret, c.writeTo)\n\tret = append(ret, c.controls...)\n\tfor _, branch := range c.writeToBranches {\n\t\tfor node := range branch.endNodes {\n\t\t\tret = append(ret, node)\n\t\t}\n\t}\n\treturn uniqueSlice(ret)\n}\n\nfunc uniqueSlice(s []string) []string {\n\tseen := make(map[string]struct{}, len(s))\n\tcur := 0\n\tfor i := range s {\n\t\tif _, ok := seen[s[i]]; !ok {\n\t\t\tseen[s[i]] = struct{}{}\n\t\t\ts[cur] = s[i]\n\t\t\tcur++\n\t\t}\n\t}\n\treturn s[:cur]\n}\n\ntype subGraphCompileCallback struct {\n\tclosure func(ctx context.Context, info *GraphInfo)\n}\n\n// OnFinish is called when the graph is compiled.\nfunc (s *subGraphCompileCallback) OnFinish(ctx context.Context, info *GraphInfo) {\n\ts.closure(ctx, info)\n}\n\nfunc (g *graph) beforeChildGraphsCompile(opt *graphCompileOptions) map[string]*GraphInfo {\n\tif opt == nil || len(opt.callbacks) == 0 {\n\t\treturn nil\n\t}\n\n\treturn make(map[string]*GraphInfo)\n}\n\nfunc (gn *graphNode) beforeChildGraphCompile(nodeKey string, key2SubGraphs map[string]*GraphInfo) {\n\tif gn.g == nil || key2SubGraphs == nil {\n\t\treturn\n\t}\n\n\tsubGraphCallback := func(ctx2 context.Context, subGraph *GraphInfo) {\n\t\tkey2SubGraphs[nodeKey] = subGraph\n\t}\n\n\tgn.nodeInfo.compileOption.callbacks = append(gn.nodeInfo.compileOption.callbacks, &subGraphCompileCallback{closure: subGraphCallback})\n}\n\nfunc (g *graph) toGraphInfo(opt *graphCompileOptions, key2SubGraphs map[string]*GraphInfo) *GraphInfo {\n\tgInfo := &GraphInfo{\n\t\tCompileOptions: opt.origOpts,\n\t\tNodes:          make(map[string]GraphNodeInfo, len(g.nodes)),\n\t\tEdges:          gmap.Clone(g.controlEdges),\n\t\tDataEdges:      gmap.Clone(g.dataEdges),\n\t\tBranches: gmap.Map(g.branches, func(startNode string, branches []*GraphBranch) (string, []GraphBranch) {\n\t\t\tbranchInfo := make([]GraphBranch, 0, len(branches))\n\t\t\tfor _, b := range branches {\n\t\t\t\tbranchInfo = append(branchInfo, GraphBranch{\n\t\t\t\t\tinvoke:        b.invoke,\n\t\t\t\t\tcollect:       b.collect,\n\t\t\t\t\tinputType:     b.inputType,\n\t\t\t\t\tgenericHelper: b.genericHelper,\n\t\t\t\t\tendNodes:      gmap.Clone(b.endNodes),\n\t\t\t\t})\n\t\t\t}\n\t\t\treturn startNode, branchInfo\n\t\t}),\n\t\tInputType:       g.expectedInputType,\n\t\tOutputType:      g.expectedOutputType,\n\t\tName:            opt.graphName,\n\t\tGenStateFn:      g.stateGenerator,\n\t\tNewGraphOptions: g.newOpts,\n\t}\n\n\tfor key := range g.nodes {\n\t\tgNode := g.nodes[key]\n\t\tif gNode.executorMeta.component == ComponentOfPassthrough {\n\t\t\tgInfo.Nodes[key] = GraphNodeInfo{\n\t\t\t\tComponent:        gNode.executorMeta.component,\n\t\t\t\tGraphAddNodeOpts: gNode.opts,\n\t\t\t\tInputType:        gNode.cr.inputType,\n\t\t\t\tOutputType:       gNode.cr.outputType,\n\t\t\t\tName:             gNode.nodeInfo.name,\n\t\t\t\tInputKey:         gNode.cr.nodeInfo.inputKey,\n\t\t\t\tOutputKey:        gNode.cr.nodeInfo.outputKey,\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\n\t\tgNodeInfo := &GraphNodeInfo{\n\t\t\tComponent:        gNode.executorMeta.component,\n\t\t\tInstance:         gNode.instance,\n\t\t\tGraphAddNodeOpts: gNode.opts,\n\t\t\tInputType:        gNode.cr.inputType,\n\t\t\tOutputType:       gNode.cr.outputType,\n\t\t\tName:             gNode.nodeInfo.name,\n\t\t\tInputKey:         gNode.cr.nodeInfo.inputKey,\n\t\t\tOutputKey:        gNode.cr.nodeInfo.outputKey,\n\t\t\tMappings:         g.fieldMappingRecords[key],\n\t\t}\n\n\t\tif gi, ok := key2SubGraphs[key]; ok {\n\t\t\tgNodeInfo.GraphInfo = gi\n\t\t}\n\n\t\tgInfo.Nodes[key] = *gNodeInfo\n\t}\n\n\treturn gInfo\n}\n\nfunc (g *graph) onCompileFinish(ctx context.Context, opt *graphCompileOptions, key2SubGraphs map[string]*GraphInfo) {\n\tif opt == nil {\n\t\treturn\n\t}\n\n\tif len(opt.callbacks) == 0 {\n\t\treturn\n\t}\n\n\tgInfo := g.toGraphInfo(opt, key2SubGraphs)\n\n\tfor _, cb := range opt.callbacks {\n\t\tcb.OnFinish(ctx, gInfo)\n\t}\n}\n\nfunc (g *graph) getGenericHelper() *genericHelper {\n\treturn g.genericHelper\n}\n\nfunc (g *graph) GetType() string {\n\treturn \"\"\n}\n\nfunc transferTask(script [][]string, invertedEdges map[string][]string) [][]string {\n\tutilMap := map[string]bool{}\n\tfor i := len(script) - 1; i >= 0; i-- {\n\t\tfor j := 0; j < len(script[i]); j++ {\n\t\t\t// deduplicate\n\t\t\tif _, ok := utilMap[script[i][j]]; ok {\n\t\t\t\tscript[i] = append(script[i][:j], script[i][j+1:]...)\n\t\t\t\tj--\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tutilMap[script[i][j]] = true\n\n\t\t\ttarget := i\n\t\t\tfor k := i + 1; k < len(script); k++ {\n\t\t\t\thasDependencies := false\n\t\t\t\tfor l := range script[k] {\n\t\t\t\t\tfor _, dependency := range invertedEdges[script[i][j]] {\n\t\t\t\t\t\tif script[k][l] == dependency {\n\t\t\t\t\t\t\thasDependencies = true\n\t\t\t\t\t\t\tbreak\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\tif hasDependencies {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\tif hasDependencies {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\ttarget = k\n\t\t\t}\n\t\t\tif target != i {\n\t\t\t\tscript[target] = append(script[target], script[i][j])\n\t\t\t\tscript[i] = append(script[i][:j], script[i][j+1:]...)\n\t\t\t\tj--\n\t\t\t}\n\t\t}\n\t}\n\n\treturn script\n}\n\nfunc validateDAG(chanSubscribeTo map[string]*chanCall, controlPredecessors map[string][]string) error {\n\tm := map[string]int{}\n\tfor node := range chanSubscribeTo {\n\t\tif edges, ok := controlPredecessors[node]; ok {\n\t\t\tm[node] = len(edges)\n\t\t\tfor _, pre := range edges {\n\t\t\t\tif pre == START {\n\t\t\t\t\tm[node] -= 1\n\t\t\t\t}\n\t\t\t}\n\t\t} else {\n\t\t\tm[node] = 0\n\t\t}\n\t}\n\thasChanged := true\n\tfor hasChanged {\n\t\thasChanged = false\n\t\tfor node := range m {\n\t\t\tif m[node] == 0 {\n\t\t\t\thasChanged = true\n\t\t\t\tfor _, subNode := range chanSubscribeTo[node].controls {\n\t\t\t\t\tif subNode == END {\n\t\t\t\t\t\tcontinue\n\t\t\t\t\t}\n\t\t\t\t\tm[subNode]--\n\t\t\t\t}\n\t\t\t\tfor _, subBranch := range chanSubscribeTo[node].writeToBranches {\n\t\t\t\t\tfor subNode := range subBranch.endNodes {\n\t\t\t\t\t\tif subNode == END {\n\t\t\t\t\t\t\tcontinue\n\t\t\t\t\t\t}\n\t\t\t\t\t\tm[subNode]--\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\tm[node] = -1\n\t\t\t}\n\t\t}\n\t}\n\n\tvar loopStarts []string\n\tfor k, v := range m {\n\t\tif v > 0 {\n\t\t\tloopStarts = append(loopStarts, k)\n\t\t}\n\t}\n\tif len(loopStarts) > 0 {\n\t\treturn fmt.Errorf(\"%w: %s\", DAGInvalidLoopErr, formatLoops(findLoops(loopStarts, chanSubscribeTo)))\n\t}\n\treturn nil\n}\n\n// DAGInvalidLoopErr indicates the graph contains a cycle and is invalid.\nvar DAGInvalidLoopErr = errors.New(\"DAG is invalid, has loop\")\n\nfunc findLoops(startNodes []string, chanCalls map[string]*chanCall) [][]string {\n\tcontrolSuccessors := map[string][]string{}\n\tfor node, ch := range chanCalls {\n\t\tcontrolSuccessors[node] = append(controlSuccessors[node], ch.controls...)\n\t\tfor _, b := range ch.writeToBranches {\n\t\t\tfor end := range b.endNodes {\n\t\t\t\tcontrolSuccessors[node] = append(controlSuccessors[node], end)\n\t\t\t}\n\t\t}\n\t}\n\n\tvisited := map[string]bool{}\n\tvar dfs func(path []string) [][]string\n\tdfs = func(path []string) [][]string {\n\t\tvar ret [][]string\n\t\tpathEnd := path[len(path)-1]\n\t\tsuccessors, ok := controlSuccessors[pathEnd]\n\t\tif !ok {\n\t\t\treturn nil\n\t\t}\n\t\tfor _, successor := range successors {\n\t\t\tvisited[successor] = true\n\n\t\t\tif successor == END {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tvar looped bool\n\t\t\tfor i, node := range path {\n\t\t\t\tif node == successor {\n\t\t\t\t\tret = append(ret, append(path[i:], successor))\n\t\t\t\t\tlooped = true\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t\tif looped {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tret = append(ret, dfs(append(path, successor))...)\n\t\t}\n\t\treturn ret\n\t}\n\n\tvar ret [][]string\n\tfor _, node := range startNodes {\n\t\tif !visited[node] {\n\t\t\tret = append(ret, dfs([]string{node})...)\n\t\t}\n\t}\n\treturn ret\n}\n\nfunc formatLoops(loops [][]string) string {\n\tsb := strings.Builder{}\n\tfor _, loop := range loops {\n\t\tif len(loop) == 0 {\n\t\t\tcontinue\n\t\t}\n\t\tsb.WriteString(\"[\")\n\t\tsb.WriteString(loop[0])\n\t\tfor i := 1; i < len(loop); i++ {\n\t\t\tsb.WriteString(\"->\")\n\t\t\tsb.WriteString(loop[i])\n\t\t}\n\t\tsb.WriteString(\"]\")\n\t}\n\treturn sb.String()\n}\n\n// NewNodePath specifies a path to a node in the graph, which is composed of node keys.\n// Starting from the top graph,\n// following this set of node keys can lead to a specific node in the top graph or a subgraph.\n//\n// e.g.\n// NewNodePath(\"sub_graph_node_key\", \"node_key_within_sub_graph\")\nfunc NewNodePath(nodeKeyPath ...string) *NodePath {\n\treturn &NodePath{path: nodeKeyPath}\n}\n\n// NodePath represents a path composed of node keys to locate a node.\ntype NodePath struct {\n\tpath []string\n}\n\n// GetPath returns the sequence of node keys in the path.\nfunc (p *NodePath) GetPath() []string {\n\treturn p.path\n}\n"
  },
  {
    "path": "compose/graph_add_node_options.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\ntype graphAddNodeOpts struct {\n\tnodeOptions *nodeOptions\n\tprocessor   *processorOpts\n\n\tneedState bool\n}\n\n// GraphAddNodeOpt is a functional option type for adding a node to a graph.\n// e.g.\n//\n//\tgraph.AddNode(\"node_name\", node, compose.WithInputKey(\"input_key\"), compose.WithOutputKey(\"output_key\"))\ntype GraphAddNodeOpt func(o *graphAddNodeOpts)\n\ntype nodeOptions struct {\n\tnodeName string\n\n\tnodeKey string\n\n\tinputKey  string\n\toutputKey string\n\n\tgraphCompileOption []GraphCompileOption // when this node is itself an AnyGraph, this option will be used to compile the node as a nested graph\n}\n\n// WithNodeName sets the name of the node.\nfunc WithNodeName(n string) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.nodeOptions.nodeName = n\n\t}\n}\n\n// WithNodeKey set the node key, which is used to identify the node in the chain.\n// only for use in Chain/StateChain.\nfunc WithNodeKey(key string) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.nodeOptions.nodeKey = key\n\t}\n}\n\n// WithInputKey sets the input key of the node.\n// this will change the input value of the node, for example, if the pre node's output is map[string]any{\"key01\": \"value01\"},\n// and the current node's input key is \"key01\", then the current node's input value will be \"value01\".\nfunc WithInputKey(k string) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.nodeOptions.inputKey = k\n\t}\n}\n\n// WithOutputKey sets the output key of the node.\n// this will change the output value of the node, for example, if the current node's output key is \"key01\",\n// then the node's output value will be map[string]any{\"key01\": value}.\nfunc WithOutputKey(k string) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.nodeOptions.outputKey = k\n\t}\n}\n\n// WithGraphCompileOptions when the node is an AnyGraph, use this option to set compile option for the node.\n// e.g.\n//\n//\tgraph.AddNode(\"node_name\", node, compose.WithGraphCompileOptions(compose.WithGraphName(\"my_sub_graph\")))\nfunc WithGraphCompileOptions(opts ...GraphCompileOption) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.nodeOptions.graphCompileOption = opts\n\t}\n}\n\n// WithStatePreHandler modify node's input of I according to state S and input or store input information into state, and it's thread-safe.\n// notice: this option requires Graph to be created with WithGenLocalState option.\n// I: input type of the Node like ChatModel, Lambda, Retriever etc.\n// S: state type defined in WithGenLocalState\nfunc WithStatePreHandler[I, S any](pre StatePreHandler[I, S]) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.processor.statePreHandler = convertPreHandler(pre)\n\t\to.processor.preStateType = generic.TypeOf[S]()\n\t\to.needState = true\n\t}\n}\n\n// WithStatePostHandler modify node's output of O according to state S and output or store output information into state, and it's thread-safe.\n// notice: this option requires Graph to be created with WithGenLocalState option.\n// O: output type of the Node like ChatModel, Lambda, Retriever etc.\n// S: state type defined in WithGenLocalState\nfunc WithStatePostHandler[O, S any](post StatePostHandler[O, S]) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.processor.statePostHandler = convertPostHandler(post)\n\t\to.processor.postStateType = generic.TypeOf[S]()\n\t\to.needState = true\n\t}\n}\n\n// WithStreamStatePreHandler modify node's streaming input of I according to state S and input or store input information into state, and it's thread-safe.\n// notice: this option requires Graph to be created with WithGenLocalState option.\n// when to use: when upstream node's output is an actual stream, and you want the current node's input to remain an actual stream after state pre handler.\n// caution: while StreamStatePreHandler is thread safe, modifying state within your own goroutine is NOT.\n// I: input type of the Node like ChatModel, Lambda, Retriever etc.\n// S: state type defined in WithGenLocalState\nfunc WithStreamStatePreHandler[I, S any](pre StreamStatePreHandler[I, S]) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.processor.statePreHandler = streamConvertPreHandler(pre)\n\t\to.processor.preStateType = generic.TypeOf[S]()\n\t\to.needState = true\n\t}\n}\n\n// WithStreamStatePostHandler modify node's streaming output of O according to state S and output or store output information into state, and it's thread-safe.\n// notice: this option requires Graph to be created with WithGenLocalState option.\n// when to use: when current node's output is an actual stream, and you want the downstream node's input to remain an actual stream after state post handler.\n// caution: while StreamStatePostHandler is thread safe, modifying state within your own goroutine is NOT.\n// O: output type of the Node like ChatModel, Lambda, Retriever etc.\n// S: state type defined in WithGenLocalState\nfunc WithStreamStatePostHandler[O, S any](post StreamStatePostHandler[O, S]) GraphAddNodeOpt {\n\treturn func(o *graphAddNodeOpts) {\n\t\to.processor.statePostHandler = streamConvertPostHandler(post)\n\t\to.processor.postStateType = generic.TypeOf[S]()\n\t\to.needState = true\n\t}\n}\n\ntype processorOpts struct {\n\tstatePreHandler  *composableRunnable\n\tpreStateType     reflect.Type // used for type validation\n\tstatePostHandler *composableRunnable\n\tpostStateType    reflect.Type // used for type validation\n}\n\nfunc getGraphAddNodeOpts(opts ...GraphAddNodeOpt) *graphAddNodeOpts {\n\topt := &graphAddNodeOpts{\n\t\tnodeOptions: &nodeOptions{\n\t\t\tnodeName: \"\",\n\t\t\tnodeKey:  \"\",\n\t\t},\n\t\tprocessor: &processorOpts{\n\t\t\tstatePreHandler:  nil,\n\t\t\tstatePostHandler: nil,\n\t\t},\n\t}\n\n\tfor _, fn := range opts {\n\t\tfn(opt)\n\t}\n\n\treturn opt\n}\n"
  },
  {
    "path": "compose/graph_call_options.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"reflect\"\n\t\"time\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n)\n\ntype graphCancelChanKey struct{}\ntype graphCancelChanVal struct {\n\tch chan *time.Duration\n}\n\ntype graphInterruptOptions struct {\n\ttimeout *time.Duration\n}\n\n// GraphInterruptOption configures behavior when interrupting a running graph.\ntype GraphInterruptOption func(o *graphInterruptOptions)\n\n// WithGraphInterruptTimeout specifies the max waiting time before generating an interrupt.\n// After the max waiting time, the graph will force an interrupt. Any unfinished tasks will be re-run when the graph is resumed.\nfunc WithGraphInterruptTimeout(timeout time.Duration) GraphInterruptOption {\n\treturn func(o *graphInterruptOptions) {\n\t\to.timeout = &timeout\n\t}\n}\n\n// WithGraphInterrupt creates a context with graph cancellation support.\n// When the returned context is used to invoke a graph or workflow, calling the interrupt function will trigger an interrupt.\n// The graph will wait for current tasks to complete by default.\n//\n// Input Persistence: When WithGraphInterrupt is used, ALL nodes (in both root graph and subgraphs) will automatically\n// persist their inputs (both streaming and non-streaming) before execution. If the graph is interrupted, these inputs\n// are restored when the graph resumes from a checkpoint, ensuring interrupted nodes receive their original inputs.\n//\n// This behavior differs from internal interrupts triggered via compose.Interrupt() within a node's function body.\n// Internal interrupts do NOT automatically persist inputs - the node author must manage input persistence manually,\n// either by saving it in the global graph state or using compose.StatefulInterrupt() to store it in local interrupt state.\n// WithGraphInterrupt enables automatic input persistence because external interrupts can occur at any point during\n// node execution, making it impossible for the node to prepare for the interrupt.\n//\n// Why input persistence is not enabled by default for internal interrupts: Enabling it universally would break\n// existing code that relies on checking \"input == nil\" to determine whether the node is running for the first time\n// or resuming from an interrupt. The recommended approach is to use compose.GetInterruptState() to explicitly\n// determine whether the current execution is a first run or a resume.\nfunc WithGraphInterrupt(parent context.Context) (ctx context.Context, interrupt func(opts ...GraphInterruptOption)) {\n\tch := make(chan *time.Duration, 1)\n\tctx = context.WithValue(parent, graphCancelChanKey{}, &graphCancelChanVal{\n\t\tch: ch,\n\t})\n\treturn ctx, func(opts ...GraphInterruptOption) {\n\t\to := &graphInterruptOptions{}\n\t\tfor _, opt := range opts {\n\t\t\topt(o)\n\t\t}\n\t\tch <- o.timeout\n\t\tclose(ch)\n\t}\n}\n\nfunc getGraphCancel(ctx context.Context) *graphCancelChanVal {\n\tval, ok := ctx.Value(graphCancelChanKey{}).(*graphCancelChanVal)\n\tif !ok {\n\t\treturn nil\n\t}\n\treturn val\n}\n\n// Option is a functional option type for calling a graph.\ntype Option struct {\n\toptions []any\n\thandler []callbacks.Handler\n\n\tpaths []*NodePath\n\n\tmaxRunSteps         int\n\tcheckPointID        *string\n\twriteToCheckPointID *string\n\tforceNewRun         bool\n\tstateModifier       StateModifier\n}\n\nfunc (o Option) deepCopy() Option {\n\tnOptions := make([]any, len(o.options))\n\tcopy(nOptions, o.options)\n\tnHandler := make([]callbacks.Handler, len(o.handler))\n\tcopy(nHandler, o.handler)\n\tnPaths := make([]*NodePath, len(o.paths))\n\tfor i, path := range o.paths {\n\t\tnPath := *path\n\t\tnPaths[i] = &nPath\n\t}\n\treturn Option{\n\t\toptions:     nOptions,\n\t\thandler:     nHandler,\n\t\tpaths:       nPaths,\n\t\tmaxRunSteps: o.maxRunSteps,\n\t}\n}\n\n// DesignateNode sets the key of the node to which the option will be applied.\n// notice: only effective at the top graph.\n// e.g.\n//\n// embeddingOption := compose.WithEmbeddingOption(embedding.WithModel(\"text-embedding-3-small\"))\n// runnable.Invoke(ctx, \"input\", embeddingOption.DesignateNode(\"embedding_node_key\"))\nfunc (o Option) DesignateNode(nodeKey ...string) Option {\n\tnKeys := make([]*NodePath, len(nodeKey))\n\tfor i, k := range nodeKey {\n\t\tnKeys[i] = NewNodePath(k)\n\t}\n\treturn o.DesignateNodeWithPath(nKeys...)\n}\n\n// DesignateNodeWithPath sets the path of the node(s) to which the option will be applied.\n// You can specify a node in the subgraph through `NodePath` to make the option only take effect at this node.\n//\n// e.g.\n// nodePath := NewNodePath(\"sub_graph_node_key\", \"node_key_within_sub_graph\")\n// DesignateNodeWithPath(nodePath)\nfunc (o Option) DesignateNodeWithPath(path ...*NodePath) Option {\n\to.paths = append(o.paths, path...)\n\treturn o\n}\n\n// WithEmbeddingOption is a functional option type for embedding component.\n// e.g.\n//\n//\tembeddingOption := compose.WithEmbeddingOption(embedding.WithModel(\"text-embedding-3-small\"))\n//\trunnable.Invoke(ctx, \"input\", embeddingOption)\nfunc WithEmbeddingOption(opts ...embedding.Option) Option {\n\treturn withComponentOption(opts...)\n}\n\n// WithRetrieverOption is a functional option type for retriever component.\n// e.g.\n//\n//\tretrieverOption := compose.WithRetrieverOption(retriever.WithIndex(\"my_index\"))\n//\trunnable.Invoke(ctx, \"input\", retrieverOption)\nfunc WithRetrieverOption(opts ...retriever.Option) Option {\n\treturn withComponentOption(opts...)\n}\n\n// WithLoaderOption is a functional option type for loader component.\n// e.g.\n//\n//\tloaderOption := compose.WithLoaderOption(document.WithCollection(\"my_collection\"))\n//\trunnable.Invoke(ctx, \"input\", loaderOption)\nfunc WithLoaderOption(opts ...document.LoaderOption) Option {\n\treturn withComponentOption(opts...)\n}\n\n// WithDocumentTransformerOption is a functional option type for document transformer component.\nfunc WithDocumentTransformerOption(opts ...document.TransformerOption) Option {\n\treturn withComponentOption(opts...)\n}\n\n// WithIndexerOption is a functional option type for indexer component.\n// e.g.\n//\n//\tindexerOption := compose.WithIndexerOption(indexer.WithSubIndexes([]string{\"my_sub_index\"}))\n//\trunnable.Invoke(ctx, \"input\", indexerOption)\nfunc WithIndexerOption(opts ...indexer.Option) Option {\n\treturn withComponentOption(opts...)\n}\n\n// WithChatModelOption is a functional option type for chat model component.\n// e.g.\n//\n//\tchatModelOption := compose.WithChatModelOption(model.WithTemperature(0.7))\n//\trunnable.Invoke(ctx, \"input\", chatModelOption)\nfunc WithChatModelOption(opts ...model.Option) Option {\n\treturn withComponentOption(opts...)\n}\n\n// WithChatTemplateOption is a functional option type for chat template component.\nfunc WithChatTemplateOption(opts ...prompt.Option) Option {\n\treturn withComponentOption(opts...)\n}\n\n// WithToolsNodeOption is a functional option type for tools node component.\nfunc WithToolsNodeOption(opts ...ToolsNodeOption) Option {\n\treturn withComponentOption(opts...)\n}\n\n// WithLambdaOption is a functional option type for lambda component.\nfunc WithLambdaOption(opts ...any) Option {\n\treturn Option{\n\t\toptions: opts,\n\t\tpaths:   make([]*NodePath, 0),\n\t}\n}\n\n// WithCallbacks set callback handlers for all components in a single call.\n// e.g.\n//\n//\trunnable.Invoke(ctx, \"input\", compose.WithCallbacks(&myCallbacks{}))\nfunc WithCallbacks(cbs ...callbacks.Handler) Option {\n\treturn Option{\n\t\thandler: cbs,\n\t}\n}\n\n// WithRuntimeMaxSteps sets the maximum number of steps for the graph runtime.\n// e.g.\n//\n//\trunnable.Invoke(ctx, \"input\", compose.WithRuntimeMaxSteps(20))\nfunc WithRuntimeMaxSteps(maxSteps int) Option {\n\treturn Option{\n\t\tmaxRunSteps: maxSteps,\n\t}\n}\n\nfunc withComponentOption[TOption any](opts ...TOption) Option {\n\to := make([]any, 0, len(opts))\n\tfor i := range opts {\n\t\to = append(o, opts[i])\n\t}\n\treturn Option{\n\t\toptions: o,\n\t\tpaths:   make([]*NodePath, 0),\n\t}\n}\n\nfunc convertOption[TOption any](opts ...any) ([]TOption, error) {\n\tif len(opts) == 0 {\n\t\treturn nil, nil\n\t}\n\tret := make([]TOption, 0, len(opts))\n\tfor i := range opts {\n\t\to, ok := opts[i].(TOption)\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"unexpected component option type, expected:%s, actual:%s\", reflect.TypeOf((*TOption)(nil)).Elem().String(), reflect.TypeOf(opts[i]).String())\n\t\t}\n\t\tret = append(ret, o)\n\t}\n\treturn ret, nil\n}\n"
  },
  {
    "path": "compose/graph_call_options_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\tmockDocument \"github.com/cloudwego/eino/internal/mock/components/document\"\n\tmockEmbedding \"github.com/cloudwego/eino/internal/mock/components/embedding\"\n\tmockRetriever \"github.com/cloudwego/eino/internal/mock/components/retriever\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nvar optionSuccess = true\nvar idx int\n\nfunc checkOption(opts ...model.Option) bool {\n\tif len(opts) != 2 {\n\t\treturn false\n\t}\n\to := model.GetCommonOptions(&model.Options{}, opts...)\n\tif o.TopP == nil || *o.TopP != 1.0 {\n\t\treturn false\n\t}\n\tif o.Model == nil {\n\t\treturn false\n\t}\n\tif idx == 0 {\n\t\tidx = 1\n\t\tif o.Model == nil || *o.Model != \"123\" {\n\t\t\treturn false\n\t\t}\n\t} else {\n\t\tidx = 0\n\t\tif o.Model == nil || *o.Model != \"456\" {\n\t\t\treturn false\n\t\t}\n\t}\n\n\treturn true\n}\n\ntype testModel struct{}\n\nfunc (t *testModel) BindTools(tools []*schema.ToolInfo) error {\n\treturn nil\n}\n\nfunc (t *testModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tif !checkOption(opts...) {\n\t\toptionSuccess = false\n\t}\n\treturn &schema.Message{}, nil\n}\n\nfunc (t *testModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tif !checkOption(opts...) {\n\t\toptionSuccess = false\n\t}\n\tsr, sw := schema.Pipe[*schema.Message](1)\n\tsw.Send(nil, nil)\n\tsw.Close()\n\treturn sr, nil\n}\n\nfunc TestCallOption(t *testing.T) {\n\tg := NewGraph[[]*schema.Message, *schema.Message]()\n\terr := g.AddLambdaNode(\"1\", InvokableLambdaWithOption(func(ctx context.Context, input []*schema.Message, opts ...string) (output []*schema.Message, err error) {\n\t\tif len(opts) != 1 || opts[0] != \"1\" {\n\t\t\tt.Fatalf(\"lambda option length isn't 1 or content isn't '1': %v\", opts)\n\t\t}\n\t\treturn input, nil\n\t}))\n\tassert.Nil(t, err)\n\n\terr = g.AddChatModelNode(\"2\", &testModel{})\n\tassert.Nil(t, err)\n\n\terr = g.AddLambdaNode(\"-\", InvokableLambda(func(ctx context.Context, input *schema.Message) (output []*schema.Message, err error) {\n\t\treturn []*schema.Message{input}, nil\n\t}))\n\tassert.Nil(t, err)\n\n\terr = g.AddChatModelNode(\"3\", &testModel{})\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.Nil(t, err)\n\n\terr = g.AddEdge(\"2\", \"-\")\n\tassert.Nil(t, err)\n\n\terr = g.AddEdge(\"-\", \"3\")\n\tassert.Nil(t, err)\n\n\terr = g.AddEdge(\"3\", END)\n\tassert.Nil(t, err)\n\n\tctx := context.Background()\n\n\tr, err := g.Compile(ctx)\n\tassert.Nil(t, err)\n\n\tsessionKey := struct{}{}\n\tstartCnt := 0\n\tendCnt := 0\n\topts := []Option{\n\t\tWithChatModelOption(\n\t\t\tmodel.WithModel(\"123\"),\n\t\t).DesignateNode(\"2\"),\n\t\tWithChatModelOption(\n\t\t\tmodel.WithModel(\"456\"),\n\t\t).DesignateNode(\"3\"),\n\t\tWithChatModelOption(\n\t\t\tmodel.WithTopP(1.0),\n\t\t),\n\t\tWithCallbacks(callbacks.NewHandlerBuilder().\n\t\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\t\tstartCnt++\n\t\t\t\treturn context.WithValue(ctx, sessionKey, \"start\")\n\t\t\t}).\n\t\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\t\tif ctx.Value(sessionKey).(string) == \"start\" {\n\t\t\t\t\tendCnt++\n\t\t\t\t\treturn context.WithValue(ctx, sessionKey, \"end\")\n\t\t\t\t}\n\t\t\t\treturn ctx\n\t\t\t}).Build()).DesignateNode(\"3\"),\n\t\tWithLambdaOption(\"1\").DesignateNode(\"1\"),\n\t}\n\n\t_, err = r.Invoke(ctx, []*schema.Message{},\n\t\topts...)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif !optionSuccess {\n\t\tt.Fatal(\"invoke option fail\")\n\t}\n\tif startCnt != 1 {\n\t\tt.Fatal(\"node callback fail\")\n\t}\n\tif endCnt != 1 {\n\t\tt.Fatal(\"node callback fail\")\n\t}\n\t_, err = r.Stream(ctx, []*schema.Message{},\n\t\topts...)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif !optionSuccess {\n\t\tt.Fatal(\"stream option fail\")\n\t}\n\n\tsrOfCollect, swOfCollect := schema.Pipe[[]*schema.Message](1)\n\tswOfCollect.Send([]*schema.Message{}, nil)\n\tswOfCollect.Close()\n\t_, err = r.Collect(ctx, srOfCollect, opts...)\n\tassert.Nil(t, err)\n\n\tif !optionSuccess {\n\t\tt.Fatal(\"collect option fail\")\n\t}\n\n\tsrOfTransform, swOfTransform := schema.Pipe[[]*schema.Message](1)\n\tswOfTransform.Send([]*schema.Message{}, nil)\n\tswOfTransform.Close()\n\t_, err = r.Transform(ctx, srOfTransform, opts...)\n\tassert.Nil(t, err)\n\n\tif !optionSuccess {\n\t\tt.Fatal(\"transform option fail\")\n\t}\n}\n\nfunc TestCallOptionsOneByOne(t *testing.T) {\n\tctx := context.Background()\n\tt.Run(\"common_option\", func(t *testing.T) {\n\t\ttype option struct {\n\t\t\tuid int64\n\t\t}\n\n\t\topt := withComponentOption(&option{uid: 100})\n\t\tassert.Len(t, opt.options, 1)\n\t\tassert.IsType(t, &option{}, opt.options[0])\n\t\tassert.Equal(t, &option{uid: 100}, opt.options[0])\n\t})\n\n\tt.Run(\"embedding_option\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tinst := mockEmbedding.NewMockEmbedder(ctrl)\n\t\tvar opt *embedding.Options\n\t\tinst.EXPECT().EmbedStrings(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) {\n\t\t\t\topt = embedding.GetCommonOptions(&embedding.Options{}, opts...)\n\t\t\t\treturn nil, nil\n\t\t\t}).Times(1)\n\t\tch := NewChain[map[string]any, map[string]any]()\n\t\tch.AppendEmbedding(inst, WithInputKey(\"input\"), WithOutputKey(\"output\"))\n\t\tr, err := ch.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\touts, err := r.Invoke(ctx,\n\t\t\tmap[string]any{\"input\": []string{}},\n\t\t\tWithEmbeddingOption(embedding.WithModel(\"123\")),\n\t\t)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, outs, \"output\")\n\n\t\tassert.NotNil(t, opt.Model)\n\t\tassert.Equal(t, \"123\", *opt.Model)\n\t})\n\n\tt.Run(\"retriever_option\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tinst := mockRetriever.NewMockRetriever(ctrl)\n\t\tvar opt *retriever.Options\n\t\tinst.EXPECT().Retrieve(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n\t\t\t\topt = retriever.GetCommonOptions(&retriever.Options{}, opts...)\n\t\t\t\treturn nil, nil\n\t\t\t}).\n\t\t\tTimes(1)\n\t\tch := NewChain[map[string]any, map[string]any]()\n\t\tch.AppendRetriever(inst, WithInputKey(\"input\"), WithOutputKey(\"output\"))\n\t\tr, err := ch.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\touts, err := r.Invoke(ctx,\n\t\t\tmap[string]any{\"input\": \"hi\"},\n\t\t\tWithRetrieverOption(retriever.WithIndex(\"123\")),\n\t\t)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, outs, \"output\")\n\n\t\tassert.NotNil(t, opt.Index)\n\t\tassert.Equal(t, \"123\", *opt.Index)\n\t})\n\n\tt.Run(\"loader_option\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tinst := mockDocument.NewMockLoader(ctrl)\n\t\ttype implOption struct {\n\t\t\tuid int64\n\t\t}\n\n\t\ttype implOptFn func(o *implOption)\n\n\t\twithUID := func(uid int64) document.LoaderOption {\n\t\t\treturn document.WrapLoaderImplSpecificOptFn[implOption](func(i *implOption) {\n\t\t\t\ti.uid = uid\n\t\t\t})\n\t\t}\n\n\t\tvar opt *implOption\n\n\t\tinst.EXPECT().Load(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, src document.Source, opts ...document.LoaderOption) ([]*schema.Document, error) {\n\t\t\t\topt = document.GetLoaderImplSpecificOptions[implOption](&implOption{uid: 1}, opts...)\n\t\t\t\treturn nil, nil\n\t\t\t}).\n\t\t\tTimes(1)\n\t\tch := NewChain[map[string]any, map[string]any]()\n\t\tch.AppendLoader(inst, WithInputKey(\"input\"), WithOutputKey(\"output\"))\n\t\tr, err := ch.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\touts, err := r.Invoke(ctx,\n\t\t\tmap[string]any{\"input\": document.Source{}},\n\t\t\tWithLoaderOption(withUID(123)),\n\t\t)\n\t\tassert.NoError(t, err)\n\t\tassert.Contains(t, outs, \"output\")\n\n\t\tassert.Equal(t, int64(123), opt.uid)\n\t})\n}\n\nfunc TestCallOptionInSubGraph(t *testing.T) {\n\tctx := context.Background()\n\n\ttype child1Option string\n\ttype child2Option string\n\ttype parentOption string\n\ttype grandparentOption string\n\n\tchild1 := NewGraph[string, string]()\n\terr := child1.AddLambdaNode(\"1\", InvokableLambdaWithOption(func(ctx context.Context, input string, opts ...child1Option) (output string, err error) {\n\t\tif len(opts) != 1 || opts[0] != \"child1-1\" {\n\t\t\tt.Fatal(\"child1-1 option error\")\n\t\t}\n\t\treturn input + \" child1-1\", nil\n\t}), WithNodeName(\"child1-1\"))\n\tassert.NoError(t, err)\n\terr = child1.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = child1.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\n\tchild2 := NewGraph[string, string]()\n\terr = child2.AddLambdaNode(\"1\", InvokableLambdaWithOption(func(ctx context.Context, input string, opts ...child2Option) (output string, err error) {\n\t\tif len(opts) != 1 || opts[0] != \"child2-1\" {\n\t\t\tt.Fatal(\"child2-1 option error\")\n\t\t}\n\t\treturn input + \" child2-1\", nil\n\t}), WithNodeName(\"child2-1\"))\n\tassert.NoError(t, err)\n\terr = child2.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = child2.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\n\tparent := NewGraph[string, string]()\n\terr = parent.AddLambdaNode(\"1\", InvokableLambdaWithOption(func(ctx context.Context, input string, opts ...parentOption) (output string, err error) {\n\t\tif len(opts) != 1 || opts[0] != \"parent-1\" {\n\t\t\tt.Fatal(\"parent-1 option error\")\n\t\t}\n\t\treturn input + \" parent-1\", nil\n\t}), WithNodeName(\"parent-1\"))\n\tassert.NoError(t, err)\n\terr = parent.AddGraphNode(\"2\", child1, WithNodeName(\"child1\"))\n\tassert.NoError(t, err)\n\terr = parent.AddGraphNode(\"3\", child2, WithNodeName(\"child2\"))\n\tassert.NoError(t, err)\n\terr = parent.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = parent.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = parent.AddEdge(\"2\", \"3\")\n\tassert.NoError(t, err)\n\terr = parent.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\n\tgrandParent := NewGraph[string, string]()\n\terr = grandParent.AddLambdaNode(\"1\", InvokableLambdaWithOption(func(ctx context.Context, input string, opts ...grandparentOption) (output string, err error) {\n\t\tif len(opts) != 1 || opts[0] != \"grandparent-1\" {\n\t\t\tt.Fatal(\"grandparent-1 option error\")\n\t\t}\n\t\treturn input + \" grandparent-1\", nil\n\t}), WithNodeName(\"grandparent-1\"))\n\tassert.NoError(t, err)\n\terr = grandParent.AddGraphNode(\"2\", parent, WithNodeName(\"parent\"))\n\tassert.NoError(t, err)\n\terr = grandParent.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = grandParent.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = grandParent.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\n\tr, err := grandParent.Compile(ctx, WithGraphName(\"grandparent\"))\n\tassert.NoError(t, err)\n\n\tgrandCommonTimes := 0\n\tgrandCommonCB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\tswitch grandCommonTimes {\n\t\tcase 0:\n\t\t\tif info.Name != \"grandparent\" || info.Component != ComponentOfGraph {\n\t\t\t\tt.Fatal(\"grandparent common callback 0 error\")\n\t\t\t}\n\t\tcase 1:\n\t\t\tif info.Name != \"grandparent-1\" {\n\t\t\t\tt.Fatal(\"grandparent common callback 1 error\")\n\t\t\t}\n\t\tcase 2:\n\t\t\tif info.Name != \"parent\" {\n\t\t\t\tt.Fatal(\"grandparent common callback 2 error\")\n\t\t\t}\n\t\tcase 3:\n\t\t\tif info.Name != \"parent-1\" {\n\t\t\t\tt.Fatal(\"grandparent common callback 3 error\")\n\t\t\t}\n\t\tcase 4:\n\t\t\tif info.Name != \"child1\" {\n\t\t\t\tt.Fatal(\"grandparent common callback 4 error\")\n\t\t\t}\n\t\tcase 5:\n\t\t\tif info.Name != \"child1-1\" {\n\t\t\t\tt.Fatal(\"grandparent common callback 5 error\")\n\t\t\t}\n\t\tcase 6:\n\t\t\tif info.Name != \"child2\" {\n\t\t\t\tt.Fatal(\"grandparent common callback 6 error\")\n\t\t\t}\n\t\tcase 7:\n\t\t\tif info.Name != \"child2-1\" {\n\t\t\t\tt.Fatal(\"grandparent common callback 7 error\")\n\t\t\t}\n\t\tdefault:\n\t\t\tt.Fatal(\"grandparent common callback too many\")\n\t\t}\n\t\tgrandCommonTimes++\n\t\treturn ctx\n\t}).Build()\n\tgrand1CB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\tif info.Name != \"grandparent-1\" {\n\t\t\tt.Fatal(\"grandparent common callback 0 error\")\n\t\t}\n\t\treturn ctx\n\t}).Build()\n\tparentCommonCBTimes := 0\n\tparentCommonCB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\tswitch parentCommonCBTimes {\n\t\tcase 0:\n\t\t\tif info.Name != \"parent\" {\n\t\t\t\tt.Fatal(\"parent common callback 0 error\")\n\t\t\t}\n\t\tcase 1:\n\t\t\tif info.Name != \"parent-1\" {\n\t\t\t\tt.Fatal(\"parent common callback 1 error\")\n\t\t\t}\n\t\tcase 2:\n\t\t\tif info.Name != \"child1\" {\n\t\t\t\tt.Fatal(\"parent common callback 2 error\")\n\t\t\t}\n\t\tcase 3:\n\t\t\tif info.Name != \"child1-1\" {\n\t\t\t\tt.Fatal(\"parent common callback 3 error\")\n\t\t\t}\n\t\tcase 4:\n\t\t\tif info.Name != \"child2\" {\n\t\t\t\tt.Fatal(\"parent common callback 4 error\")\n\t\t\t}\n\t\tcase 5:\n\t\t\tif info.Name != \"child2-1\" {\n\t\t\t\tt.Fatal(\"parent common callback 5 error\")\n\t\t\t}\n\t\tdefault:\n\t\t\tt.Fatal(\"parent common callback too many\")\n\t\t}\n\t\tparentCommonCBTimes++\n\t\treturn ctx\n\t}).Build()\n\tchild1CommonCBTimes := 0\n\tchild1CommonCB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\tswitch child1CommonCBTimes {\n\t\tcase 0:\n\t\t\tif info.Name != \"child1\" {\n\t\t\t\tt.Fatal(\"child1 common callback 0 error\")\n\t\t\t}\n\t\tcase 1:\n\t\t\tif info.Name != \"child1-1\" {\n\t\t\t\tt.Fatal(\"child1 common callback 1 error\")\n\t\t\t}\n\t\tdefault:\n\t\t\tt.Fatal(\"child1 common callback too many\")\n\t\t}\n\t\tchild1CommonCBTimes++\n\t\treturn ctx\n\t}).Build()\n\tchild2CB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\tif info.Name != \"child2-1\" {\n\t\t\tt.Fatal(\"child2-1 common callback 0 error\")\n\t\t}\n\t\treturn ctx\n\t}).Build()\n\n\tresult, err := r.Invoke(ctx, \"input\",\n\t\tWithCallbacks(grandCommonCB),\n\t\tWithCallbacks(parentCommonCB).DesignateNodeWithPath(NewNodePath(\"2\")),\n\t\tWithCallbacks(grand1CB).DesignateNode(\"1\"),\n\t\tWithCallbacks(child1CommonCB).DesignateNodeWithPath(NewNodePath(\"2\", \"2\")),\n\t\tWithCallbacks(child2CB).DesignateNodeWithPath(NewNodePath(\"2\", \"3\", \"1\")),\n\t\tWithLambdaOption(grandparentOption(\"grandparent-1\")).DesignateNodeWithPath(NewNodePath(\"1\")),\n\t\tWithLambdaOption(parentOption(\"parent-1\")).DesignateNodeWithPath(NewNodePath(\"2\", \"1\")),\n\t\tWithLambdaOption(child1Option(\"child1-1\")).DesignateNodeWithPath(NewNodePath(\"2\", \"2\", \"1\")),\n\t\tWithLambdaOption(child2Option(\"child2-1\")).DesignateNodeWithPath(NewNodePath(\"2\", \"3\", \"1\")),\n\t)\n\tassert.NoError(t, err)\n\tassert.Equal(t, result, \"input grandparent-1 parent-1 child1-1 child2-1\")\n}\n"
  },
  {
    "path": "compose/graph_compile_options.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\ntype graphCompileOptions struct {\n\tmaxRunSteps     int\n\tgraphName       string\n\tnodeTriggerMode NodeTriggerMode // default to AnyPredecessor (pregel)\n\n\tcallbacks []GraphCompileCallback\n\n\torigOpts []GraphCompileOption\n\n\tcheckPointStore      CheckPointStore\n\tserializer           Serializer\n\tinterruptBeforeNodes []string\n\tinterruptAfterNodes  []string\n\n\teagerDisabled bool\n\n\tmergeConfigs map[string]FanInMergeConfig\n}\n\nfunc newGraphCompileOptions(opts ...GraphCompileOption) *graphCompileOptions {\n\toption := &graphCompileOptions{}\n\n\tfor _, o := range opts {\n\t\to(option)\n\t}\n\n\toption.origOpts = opts\n\n\treturn option\n}\n\n// GraphCompileOption options for compiling AnyGraph.\ntype GraphCompileOption func(*graphCompileOptions)\n\n// WithMaxRunSteps sets the maximum number of steps that a graph can run.\n// This is useful to prevent infinite loops in graphs with cycles.\n// If the number of steps exceeds maxSteps, the graph execution will be terminated with an error.\nfunc WithMaxRunSteps(maxSteps int) GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\to.maxRunSteps = maxSteps\n\t}\n}\n\n// WithGraphName sets a name for the graph.\n// The name is used for debugging and logging purposes.\n// If not set, a default name will be used.\nfunc WithGraphName(graphName string) GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\to.graphName = graphName\n\t}\n}\n\n// WithEagerExecution enables the eager execution mode for the graph.\n// In eager mode, nodes will be executed immediately once they are ready to run,\n// without waiting for the completion of a super step, ref: https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/#runtime-engine\n// Note: Eager mode is not allowed when the graph's trigger mode is set to AnyPredecessor.\n// Workflow uses eager mode by default.\n// Deprecated: Eager execution is automatically enabled by default when a node's trigger mode is set to AllPredecessor.\n// If you were using this option previously, it can be safely removed without changing behavior.\nfunc WithEagerExecution() GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\treturn\n\t}\n}\n\n// WithEagerExecutionDisabled disables the eager execution mode for the graph.\n// By default, eager execution is enabled for Workflow and Graph with the AllPredecessor trigger mode.\n// After using this option, nodes will wait for the completion of a super step instead of execute immediately once they are ready to run.\n// ref: https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/#runtime-engine\nfunc WithEagerExecutionDisabled() GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\to.eagerDisabled = true\n\t}\n}\n\n// WithNodeTriggerMode sets the trigger mode for nodes in the graph.\n// The trigger mode determines when a node is triggered during graph execution, ref: https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/#runtime-engine\n// AnyPredecessor by default.\nfunc WithNodeTriggerMode(triggerMode NodeTriggerMode) GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\to.nodeTriggerMode = triggerMode\n\t}\n}\n\n// WithGraphCompileCallbacks sets callbacks for graph compilation.\nfunc WithGraphCompileCallbacks(cbs ...GraphCompileCallback) GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\to.callbacks = append(o.callbacks, cbs...)\n\t}\n}\n\n// FanInMergeConfig defines the configuration for fan-in merge operations.\n// It allows specifying how multiple inputs are merged into a single input.\n// StreamMergeWithSourceEOF indicates whether to emit a SourceEOF error for each stream\n// when it ends, before the final merged output is produced. This is useful for\n// tracking the completion of individual input streams in a named stream merge.\ntype FanInMergeConfig struct {\n\tStreamMergeWithSourceEOF bool //indicates whether to emit a SourceEOF error for each stream\n}\n\n// WithFanInMergeConfig sets the fan-in merge configurations\n// for the graph nodes that receive inputs from multiple sources.\nfunc WithFanInMergeConfig(confs map[string]FanInMergeConfig) GraphCompileOption {\n\treturn func(o *graphCompileOptions) {\n\t\to.mergeConfigs = confs\n\t}\n}\n\n// InitGraphCompileCallbacks set global graph compile callbacks,\n// which ONLY will be added to top level graph compile options\nfunc InitGraphCompileCallbacks(cbs []GraphCompileCallback) {\n\tglobalGraphCompileCallbacks = cbs\n}\n\nvar globalGraphCompileCallbacks []GraphCompileCallback\n"
  },
  {
    "path": "compose/graph_manager.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"runtime/debug\"\n\t\"time\"\n\n\t\"github.com/cloudwego/eino/internal\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n)\n\ntype channel interface {\n\treportValues(map[string]any) error\n\treportDependencies([]string)\n\treportSkip([]string) bool\n\tget(bool, string, *edgeHandlerManager) (any, bool, error)\n\tconvertValues(fn func(map[string]any) error) error\n\tload(channel) error\n\n\tsetMergeConfig(FanInMergeConfig)\n}\n\ntype edgeHandlerManager struct {\n\th map[string]map[string][]handlerPair\n}\n\nfunc (e *edgeHandlerManager) handle(from, to string, value any, isStream bool) (any, error) {\n\tif _, ok := e.h[from]; !ok {\n\t\treturn value, nil\n\t}\n\tif _, ok := e.h[from][to]; !ok {\n\t\treturn value, nil\n\t}\n\tif isStream {\n\t\tfor _, v := range e.h[from][to] {\n\t\t\tvalue = v.transform(value.(streamReader))\n\t\t}\n\t} else {\n\t\tfor _, v := range e.h[from][to] {\n\t\t\tvar err error\n\t\t\tvalue, err = v.invoke(value)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\treturn value, nil\n}\n\ntype preNodeHandlerManager struct {\n\th map[string][]handlerPair\n}\n\nfunc (p *preNodeHandlerManager) handle(nodeKey string, value any, isStream bool) (any, error) {\n\tif _, ok := p.h[nodeKey]; !ok {\n\t\treturn value, nil\n\t}\n\tif isStream {\n\t\tfor _, v := range p.h[nodeKey] {\n\t\t\tvalue = v.transform(value.(streamReader))\n\t\t}\n\t} else {\n\t\tfor _, v := range p.h[nodeKey] {\n\t\t\tvar err error\n\t\t\tvalue, err = v.invoke(value)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\treturn value, nil\n}\n\ntype preBranchHandlerManager struct {\n\th map[string][][]handlerPair\n}\n\nfunc (p *preBranchHandlerManager) handle(nodeKey string, idx int, value any, isStream bool) (any, error) {\n\tif _, ok := p.h[nodeKey]; !ok {\n\t\treturn value, nil\n\t}\n\tif isStream {\n\t\tfor _, v := range p.h[nodeKey][idx] {\n\t\t\tvalue = v.transform(value.(streamReader))\n\t\t}\n\t} else {\n\t\tfor _, v := range p.h[nodeKey][idx] {\n\t\t\tvar err error\n\t\t\tvalue, err = v.invoke(value)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\treturn value, nil\n}\n\ntype channelManager struct {\n\tisStream bool\n\tchannels map[string]channel\n\n\tsuccessors          map[string][]string\n\tdataPredecessors    map[string]map[string]struct{}\n\tcontrolPredecessors map[string]map[string]struct{}\n\n\tedgeHandlerManager    *edgeHandlerManager\n\tpreNodeHandlerManager *preNodeHandlerManager\n}\n\nfunc (c *channelManager) loadChannels(channels map[string]channel) error {\n\tfor key, ch := range c.channels {\n\t\tif nCh, ok := channels[key]; ok {\n\t\t\tif err := ch.load(nCh); err != nil {\n\t\t\t\treturn fmt.Errorf(\"load channel[%s] fail: %w\", key, err)\n\t\t\t}\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc (c *channelManager) updateValues(_ context.Context, values map[string] /*to*/ map[string] /*from*/ any) error {\n\tfor target, fromMap := range values {\n\t\ttoChannel, ok := c.channels[target]\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"target channel doesn't existed: %s\", target)\n\t\t}\n\t\tdps, ok := c.dataPredecessors[target]\n\t\tif !ok {\n\t\t\tdps = map[string]struct{}{}\n\t\t}\n\t\tnFromMap := make(map[string]any, len(fromMap))\n\t\tfor from, value := range fromMap {\n\t\t\tif _, ok = dps[from]; ok {\n\t\t\t\tnFromMap[from] = fromMap[from]\n\t\t\t} else {\n\t\t\t\tif sr, okk := value.(streamReader); okk {\n\t\t\t\t\tsr.close()\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\terr := toChannel.reportValues(nFromMap)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"update target channel[%s] fail: %w\", target, err)\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc (c *channelManager) updateDependencies(_ context.Context, dependenciesMap map[string][]string) error {\n\tfor target, dependencies := range dependenciesMap {\n\t\ttoChannel, ok := c.channels[target]\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"target channel doesn't existed: %s\", target)\n\t\t}\n\t\tcps, ok := c.controlPredecessors[target]\n\t\tif !ok {\n\t\t\tcps = map[string]struct{}{}\n\t\t}\n\t\tvar deps []string\n\t\tfor _, from := range dependencies {\n\t\t\tif _, ok = cps[from]; ok {\n\t\t\t\tdeps = append(deps, from)\n\t\t\t}\n\t\t}\n\n\t\ttoChannel.reportDependencies(deps)\n\t}\n\treturn nil\n}\n\nfunc (c *channelManager) getFromReadyChannels(_ context.Context) (map[string]any, error) {\n\tresult := make(map[string]any)\n\tfor target, ch := range c.channels {\n\t\tv, ready, err := ch.get(c.isStream, target, c.edgeHandlerManager)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"get value from ready channel[%s] fail: %w\", target, err)\n\t\t}\n\t\tif ready {\n\t\t\tv, err = c.preNodeHandlerManager.handle(target, v, c.isStream)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tresult[target] = v\n\t\t}\n\t}\n\treturn result, nil\n}\n\nfunc (c *channelManager) updateAndGet(ctx context.Context, values map[string]map[string]any, dependencies map[string][]string) (map[string]any, error) {\n\terr := c.updateValues(ctx, values)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"update channel fail: %w\", err)\n\t}\n\terr = c.updateDependencies(ctx, dependencies)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"update channel fail: %w\", err)\n\t}\n\treturn c.getFromReadyChannels(ctx)\n}\n\nfunc (c *channelManager) reportBranch(from string, skippedNodes []string) error {\n\tvar nKeys []string\n\tfor _, node := range skippedNodes {\n\t\tskipped := c.channels[node].reportSkip([]string{from})\n\t\tif skipped {\n\t\t\tnKeys = append(nKeys, node)\n\t\t}\n\t}\n\n\tfor i := 0; i < len(nKeys); i++ {\n\t\tkey := nKeys[i]\n\n\t\tif key == END {\n\t\t\tcontinue\n\t\t}\n\t\tif _, ok := c.successors[key]; !ok {\n\t\t\treturn fmt.Errorf(\"unknown node: %s\", key)\n\t\t}\n\t\tfor _, successor := range c.successors[key] {\n\t\t\tskipped := c.channels[successor].reportSkip([]string{key})\n\t\t\tif skipped {\n\t\t\t\tnKeys = appendIfNotExist(nKeys, successor)\n\t\t\t}\n\t\t\t// todo: detect if end node has been skipped?\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc appendIfNotExist(s []string, elem string) []string {\n\tfor _, i := range s {\n\t\tif i == elem {\n\t\t\treturn s\n\t\t}\n\t}\n\treturn append(s, elem)\n}\n\ntype task struct {\n\tctx            context.Context\n\tnodeKey        string\n\tcall           *chanCall\n\tinput          any\n\toriginalInput  any\n\toutput         any\n\toption         []any\n\terr            error\n\tskipPreHandler bool\n}\n\ntype taskManager struct {\n\trunWrapper runnableCallWrapper\n\topts       []Option\n\tneedAll    bool\n\n\tnum          uint32\n\tdone         *internal.UnboundedChan[*task]\n\trunningTasks map[string]*task\n\n\tcancelCh chan *time.Duration\n\tcanceled bool\n\tdeadline *time.Time\n\n\tpersistRerunInput bool\n}\n\nfunc (t *taskManager) execute(currentTask *task) {\n\tdefer func() {\n\t\tpanicInfo := recover()\n\t\tif panicInfo != nil {\n\t\t\tcurrentTask.output = nil\n\t\t\tcurrentTask.err = safe.NewPanicErr(panicInfo, debug.Stack())\n\t\t}\n\n\t\tt.done.Send(currentTask)\n\t}()\n\n\tctx := initNodeCallbacks(currentTask.ctx, currentTask.nodeKey, currentTask.call.action.nodeInfo, currentTask.call.action.meta, t.opts...)\n\tcurrentTask.output, currentTask.err = t.runWrapper(ctx, currentTask.call.action, currentTask.input, currentTask.option...)\n}\n\nfunc (t *taskManager) submit(tasks []*task) error {\n\tif len(tasks) == 0 {\n\t\treturn nil\n\t}\n\n\t// synchronously execute one task, if there are no other tasks in the task pool and meet one of the following conditions：\n\t// 1. the new task is the only one\n\t// 2. the task manager mode is set to needAll\n\tfor i := 0; i < len(tasks); i++ {\n\t\tcurrentTask := tasks[i]\n\n\t\tif t.persistRerunInput {\n\t\t\tif sr, ok := currentTask.input.(streamReader); ok {\n\t\t\t\tcopies := sr.copy(2)\n\t\t\t\tcurrentTask.originalInput, currentTask.input = copies[0], copies[1]\n\t\t\t} else {\n\t\t\t\tcurrentTask.originalInput = currentTask.input\n\t\t\t}\n\t\t}\n\n\t\terr := runPreHandler(currentTask, t.runWrapper)\n\t\tif err != nil {\n\t\t\t// pre-handler error, regarded as a failure of the task itself\n\t\t\tcurrentTask.err = err\n\t\t\ttasks = append(tasks[:i], tasks[i+1:]...)\n\t\t\ti--\n\t\t\tt.num++\n\t\t\tt.done.Send(currentTask)\n\t\t}\n\n\t\tt.runningTasks[currentTask.nodeKey] = currentTask\n\t}\n\tif len(tasks) == 0 {\n\t\t// all tasks' pre-handler failed\n\t\treturn nil\n\t}\n\n\tvar syncTask *task\n\tif t.num == 0 && (len(tasks) == 1 || t.needAll) && t.cancelCh == nil /*if graph can be interrupted by user, shouldn't sync run task*/ {\n\t\tsyncTask = tasks[0]\n\t\ttasks = tasks[1:]\n\t}\n\tfor _, currentTask := range tasks {\n\t\tt.num += 1\n\t\tgo t.execute(currentTask)\n\t}\n\tif syncTask != nil {\n\t\tt.num += 1\n\t\tt.execute(syncTask)\n\t}\n\treturn nil\n}\n\nfunc (t *taskManager) wait() (tasks []*task, canceled bool, canceledTasks []*task) {\n\tif t.needAll {\n\t\ttasks, canceledTasks = t.waitAll()\n\t\treturn tasks, t.canceled, canceledTasks\n\t}\n\n\tta, success, canceled := t.waitOne()\n\tif canceled {\n\t\t// has canceled and timeout, return canceled tasks\n\t\tfor _, rta := range t.runningTasks {\n\t\t\tcanceledTasks = append(canceledTasks, rta)\n\t\t}\n\t\tt.runningTasks = make(map[string]*task)\n\t\tt.num = 0\n\t\treturn nil, true, canceledTasks\n\t}\n\tif t.canceled {\n\t\t// has canceled, but not timeout, wait all\n\t\ttasks, canceledTasks = t.waitAll()\n\t\treturn append(tasks, ta), true, canceledTasks\n\t}\n\tif !success {\n\t\treturn []*task{}, t.canceled, nil\n\t}\n\n\treturn []*task{ta}, t.canceled, nil\n}\n\nfunc (t *taskManager) waitOne() (ta *task, success bool, canceled bool) {\n\tif t.num == 0 {\n\t\treturn nil, false, false\n\t}\n\n\tif t.cancelCh == nil {\n\t\tta, _ = t.done.Receive()\n\t} else {\n\t\tta, _, canceled = t.receive(t.done.Receive)\n\t}\n\n\tt.num--\n\n\tif canceled {\n\t\treturn nil, false, true\n\t}\n\n\tdelete(t.runningTasks, ta.nodeKey)\n\n\tif ta.originalInput != nil && (ta.err == nil || !isInterruptError(ta.err)) {\n\t\tif sr, ok := ta.originalInput.(streamReader); ok {\n\t\t\tsr.close()\n\t\t}\n\t\tta.originalInput = nil\n\t}\n\n\tif ta.err != nil {\n\t\t// biz error, jump post processor\n\t\treturn ta, true, false\n\t}\n\trunPostHandler(ta, t.runWrapper)\n\treturn ta, true, false\n}\n\nfunc (t *taskManager) waitAll() (successTasks []*task, canceledTasks []*task) {\n\tresult := make([]*task, 0, t.num)\n\tfor {\n\t\tta, success, canceled := t.waitOne()\n\t\tif canceled {\n\t\t\tfor _, rt := range t.runningTasks {\n\t\t\t\tcanceledTasks = append(canceledTasks, rt)\n\t\t\t}\n\t\t\tt.runningTasks = make(map[string]*task)\n\t\t\tt.num = 0\n\t\t\treturn result, canceledTasks\n\t\t}\n\t\tif !success {\n\t\t\treturn result, nil\n\t\t}\n\t\tresult = append(result, ta)\n\t}\n}\n\nfunc (t *taskManager) receive(recv func() (*task, bool)) (ta *task, closed bool, canceled bool) {\n\tif t.deadline != nil {\n\t\t// have canceled, receive in a certain time\n\t\treturn receiveWithDeadline(recv, *t.deadline)\n\t}\n\tif t.canceled {\n\t\t// canceled without timeout\n\t\tta, closed = recv()\n\t\treturn ta, closed, false\n\t}\n\tif t.cancelCh != nil {\n\t\t// have not canceled, receive while listening\n\t\tta, closed, canceled, t.canceled, t.deadline = receiveWithListening(recv, t.cancelCh)\n\t\treturn ta, closed, canceled\n\t}\n\t// won't cancel\n\tta, closed = recv()\n\treturn ta, closed, false\n}\n\nfunc receiveWithDeadline(recv func() (*task, bool), deadline time.Time) (ta *task, closed bool, canceled bool) {\n\tnow := time.Now()\n\tif deadline.Before(now) {\n\t\treturn nil, false, true\n\t}\n\n\ttimeout := deadline.Sub(now)\n\n\tresultCh := make(chan struct{}, 1)\n\n\tgo func() {\n\t\tta, closed = recv()\n\t\tresultCh <- struct{}{}\n\t}()\n\n\ttimeoutCh := time.After(timeout)\n\n\tselect {\n\tcase <-resultCh:\n\t\treturn ta, closed, false\n\tcase <-timeoutCh:\n\t\treturn nil, false, true\n\t}\n}\n\nfunc receiveWithListening(recv func() (*task, bool), cancel chan *time.Duration) (*task, bool, bool, bool, *time.Time) {\n\ttype pair struct {\n\t\tta     *task\n\t\tclosed bool\n\t}\n\tresultCh := make(chan pair, 1)\n\tvar timeoutCh <-chan time.Time\n\n\tvar deadline *time.Time\n\tcanceled := false\n\tgo func() {\n\t\tta, closed := recv()\n\t\tresultCh <- pair{ta, closed}\n\t}()\n\n\tselect {\n\tcase p := <-resultCh:\n\t\treturn p.ta, p.closed, false, false, nil\n\tcase timeout, ok := <-cancel:\n\t\tif !ok {\n\t\t\t// unreachable\n\t\t\tbreak\n\t\t}\n\t\tcanceled = true\n\t\tif timeout == nil {\n\t\t\t// canceled without timeout\n\t\t\tbreak\n\t\t}\n\t\ttimeoutCh = time.After(*timeout)\n\t\tdt := time.Now().Add(*timeout)\n\t\tdeadline = &dt\n\t}\n\n\tif timeoutCh != nil {\n\t\tselect {\n\t\tcase p := <-resultCh:\n\t\t\treturn p.ta, p.closed, false, canceled, deadline\n\t\tcase <-timeoutCh:\n\t\t\treturn nil, false, true, canceled, deadline\n\t\t}\n\t}\n\tp := <-resultCh\n\treturn p.ta, p.closed, false, canceled, nil\n}\n\nfunc runPreHandler(ta *task, runWrapper runnableCallWrapper) (err error) {\n\tdefer func() {\n\t\tif e := recover(); e != nil {\n\t\t\terr = safe.NewPanicErr(fmt.Errorf(\"panic in pre handler: %v\", e), debug.Stack())\n\t\t}\n\t}()\n\tif ta.call.preProcessor != nil && !ta.skipPreHandler {\n\t\tnInput, err := runWrapper(ta.ctx, ta.call.preProcessor, ta.input, ta.option...)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"run node[%s] pre processor fail: %w\", ta.nodeKey, err)\n\t\t}\n\t\tta.input = nInput\n\t}\n\treturn nil\n}\n\nfunc runPostHandler(ta *task, runWrapper runnableCallWrapper) {\n\tdefer func() {\n\t\tif e := recover(); e != nil {\n\t\t\tta.err = safe.NewPanicErr(fmt.Errorf(\"panic in post handler: %v\", e), debug.Stack())\n\t\t}\n\t}()\n\tif ta.call.postProcessor != nil {\n\t\tnOutput, err := runWrapper(ta.ctx, ta.call.postProcessor, ta.output, ta.option...)\n\t\tif err != nil {\n\t\t\tta.err = fmt.Errorf(\"run node[%s] post processor fail: %w\", ta.nodeKey, err)\n\t\t}\n\t\tta.output = nOutput\n\t}\n}\n"
  },
  {
    "path": "compose/graph_node.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\n// the info of most original executable object directly provided by the user\ntype executorMeta struct {\n\n\t// automatically identified based on the way of addNode\n\tcomponent component\n\n\t// indicates whether the executable object user provided could execute the callback aspect itself.\n\t// if it could, the callback in the corresponding graph node won't be executed\n\t// for components, the value comes from callbacks.Checker\n\tisComponentCallbackEnabled bool\n\n\t// for components, the value comes from components.Typer\n\t// for lambda, the value comes from the user's explicit config\n\t// if componentImplType is empty, then the class name or func name in the instance will be inferred, but no guarantee.\n\tcomponentImplType string\n}\n\ntype nodeInfo struct {\n\n\t// the name of graph node for display purposes, not unique.\n\t// passed from WithNodeName()\n\tname string\n\n\tinputKey  string\n\toutputKey string\n\n\tpreProcessor, postProcessor *composableRunnable\n\n\tcompileOption *graphCompileOptions // if the node is an AnyGraph, it will need compile options of its own\n}\n\n// graphNode the complete information of the node in graph\ntype graphNode struct {\n\tcr *composableRunnable\n\n\tg AnyGraph\n\n\tnodeInfo     *nodeInfo\n\texecutorMeta *executorMeta\n\n\tinstance any\n\topts     []GraphAddNodeOpt\n}\n\nfunc (gn *graphNode) getGenericHelper() *genericHelper {\n\tvar ret *genericHelper\n\tif gn.g != nil {\n\t\tret = gn.g.getGenericHelper()\n\t} else if gn.cr != nil {\n\t\tret = gn.cr.genericHelper\n\t} else {\n\t\treturn nil\n\t}\n\n\tif gn.nodeInfo != nil {\n\t\tif len(gn.nodeInfo.inputKey) > 0 {\n\t\t\tret = ret.forMapInput()\n\t\t}\n\t\tif len(gn.nodeInfo.outputKey) > 0 {\n\t\t\tret = ret.forMapOutput()\n\t\t}\n\t}\n\treturn ret\n}\n\nfunc (gn *graphNode) inputType() reflect.Type {\n\tif gn.nodeInfo != nil && len(gn.nodeInfo.inputKey) != 0 {\n\t\treturn generic.TypeOf[map[string]any]()\n\t}\n\t// priority follow compile\n\tif gn.g != nil {\n\t\treturn gn.g.inputType()\n\t} else if gn.cr != nil {\n\t\treturn gn.cr.inputType\n\t}\n\n\treturn nil\n}\n\nfunc (gn *graphNode) outputType() reflect.Type {\n\tif gn.nodeInfo != nil && len(gn.nodeInfo.outputKey) != 0 {\n\t\treturn generic.TypeOf[map[string]any]()\n\t}\n\t// priority follow compile\n\tif gn.g != nil {\n\t\treturn gn.g.outputType()\n\t} else if gn.cr != nil {\n\t\treturn gn.cr.outputType\n\t}\n\n\treturn nil\n}\n\nfunc (gn *graphNode) compileIfNeeded(ctx context.Context) (*composableRunnable, error) {\n\tvar r *composableRunnable\n\tif gn.g != nil {\n\t\tcr, err := gn.g.compile(ctx, gn.nodeInfo.compileOption)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tr = cr\n\t\tgn.cr = cr\n\t} else if gn.cr != nil {\n\t\tr = gn.cr\n\t} else {\n\t\treturn nil, errors.New(\"no graph or component provided\")\n\t}\n\n\tr.meta = gn.executorMeta\n\tr.nodeInfo = gn.nodeInfo\n\n\tif gn.nodeInfo.outputKey != \"\" {\n\t\tr = outputKeyedComposableRunnable(gn.nodeInfo.outputKey, r)\n\t}\n\n\tif gn.nodeInfo.inputKey != \"\" {\n\t\tr = inputKeyedComposableRunnable(gn.nodeInfo.inputKey, r)\n\t}\n\n\treturn r, nil\n}\n\nfunc parseExecutorInfoFromComponent(c component, executor any) *executorMeta {\n\n\tcomponentImplType, ok := components.GetType(executor)\n\tif !ok {\n\t\tcomponentImplType = generic.ParseTypeName(reflect.ValueOf(executor))\n\t}\n\n\treturn &executorMeta{\n\t\tcomponent:                  c,\n\t\tisComponentCallbackEnabled: components.IsCallbacksEnabled(executor),\n\t\tcomponentImplType:          componentImplType,\n\t}\n}\n\nfunc getNodeInfo(opts ...GraphAddNodeOpt) (*nodeInfo, *graphAddNodeOpts) {\n\n\topt := getGraphAddNodeOpts(opts...)\n\n\treturn &nodeInfo{\n\t\tname:          opt.nodeOptions.nodeName,\n\t\tinputKey:      opt.nodeOptions.inputKey,\n\t\toutputKey:     opt.nodeOptions.outputKey,\n\t\tpreProcessor:  opt.processor.statePreHandler,\n\t\tpostProcessor: opt.processor.statePostHandler,\n\t\tcompileOption: newGraphCompileOptions(opt.nodeOptions.graphCompileOption...),\n\t}, opt\n}\n"
  },
  {
    "path": "compose/graph_run.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"reflect\"\n\t\"strings\"\n\n\t\"github.com/cloudwego/eino/internal\"\n\t\"github.com/cloudwego/eino/internal/core\"\n\t\"github.com/cloudwego/eino/internal/serialization\"\n)\n\ntype chanCall struct {\n\taction          *composableRunnable\n\twriteTo         []string\n\twriteToBranches []*GraphBranch\n\n\tcontrols []string // branch must control\n\n\tpreProcessor, postProcessor *composableRunnable\n}\n\ntype chanBuilder func(dependencies []string, indirectDependencies []string, zeroValue func() any, emptyStream func() streamReader) channel\n\ntype runner struct {\n\tchanSubscribeTo map[string]*chanCall\n\n\tsuccessors          map[string][]string\n\tdataPredecessors    map[string][]string\n\tcontrolPredecessors map[string][]string\n\n\tinputChannels *chanCall\n\n\tchanBuilder chanBuilder // could be nil\n\teager       bool\n\tdag         bool\n\n\trunCtx func(ctx context.Context) context.Context\n\n\toptions graphCompileOptions\n\n\tinputType  reflect.Type\n\toutputType reflect.Type\n\n\t// take effect as a subgraph through toComposableRunnable\n\tinputStreamFilter                               streamMapFilter\n\tinputConverter                                  handlerPair\n\tinputFieldMappingConverter                      handlerPair\n\tinputConvertStreamPair, outputConvertStreamPair streamConvertPair\n\n\t*genericHelper\n\n\t// checks need to do because cannot check at compile\n\truntimeCheckEdges    map[string]map[string]bool\n\truntimeCheckBranches map[string][]bool\n\n\tedgeHandlerManager      *edgeHandlerManager\n\tpreNodeHandlerManager   *preNodeHandlerManager\n\tpreBranchHandlerManager *preBranchHandlerManager\n\n\tcheckPointer         *checkPointer\n\tinterruptBeforeNodes []string\n\tinterruptAfterNodes  []string\n\n\tmergeConfigs map[string]FanInMergeConfig\n}\n\nfunc (r *runner) invoke(ctx context.Context, input any, opts ...Option) (any, error) {\n\treturn r.run(ctx, false, input, opts...)\n}\n\nfunc (r *runner) transform(ctx context.Context, input streamReader, opts ...Option) (streamReader, error) {\n\ts, err := r.run(ctx, true, input, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn s.(streamReader), nil\n}\n\ntype runnableCallWrapper func(context.Context, *composableRunnable, any, ...any) (any, error)\n\nfunc runnableInvoke(ctx context.Context, r *composableRunnable, input any, opts ...any) (any, error) {\n\treturn r.i(ctx, input, opts...)\n}\n\nfunc runnableTransform(ctx context.Context, r *composableRunnable, input any, opts ...any) (any, error) {\n\treturn r.t(ctx, input.(streamReader), opts...)\n}\n\nfunc (r *runner) run(ctx context.Context, isStream bool, input any, opts ...Option) (result any, err error) {\n\thaveOnStart := false // delay triggering onGraphStart until state initialization is complete, so that the state can be accessed within onGraphStart.\n\tdefer func() {\n\t\tif !haveOnStart {\n\t\t\tctx, input = onGraphStart(ctx, input, isStream)\n\t\t}\n\t\tif err != nil {\n\t\t\tctx, err = onGraphError(ctx, err)\n\t\t} else {\n\t\t\tctx, result = onGraphEnd(ctx, result, isStream)\n\t\t}\n\t}()\n\n\tvar runWrapper runnableCallWrapper\n\trunWrapper = runnableInvoke\n\tif isStream {\n\t\trunWrapper = runnableTransform\n\t}\n\n\t// Initialize channel and task managers.\n\tcm := r.initChannelManager(isStream)\n\ttm := r.initTaskManager(runWrapper, getGraphCancel(ctx), opts...)\n\tmaxSteps := r.options.maxRunSteps\n\n\tmaxSteps, err = r.resolveMaxSteps(maxSteps, opts)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Extract and validate options for each node.\n\toptMap, extractErr := extractOption(r.chanSubscribeTo, opts...)\n\tif extractErr != nil {\n\t\treturn nil, newGraphRunError(fmt.Errorf(\"graph extract option fail: %w\", extractErr))\n\t}\n\n\t// Extract CheckPointID\n\tcheckPointID, writeToCheckPointID, stateModifier, forceNewRun := getCheckPointInfo(opts...)\n\tif checkPointID != nil && r.checkPointer.store == nil {\n\t\treturn nil, newGraphRunError(fmt.Errorf(\"receive checkpoint id but have not set checkpoint store\"))\n\t}\n\n\t// Extract subgraph\n\tpath, isSubGraph := getNodePath(ctx)\n\n\t// load checkpoint from ctx/store or init graph\n\tinitialized := false\n\tvar nextTasks []*task\n\tif cp := getCheckPointFromCtx(ctx); cp != nil {\n\t\t// in subgraph, try to load checkpoint from ctx\n\t\tinitialized = true\n\n\t\tctx, err = r.restoreCheckPointState(ctx, *path, getStateModifier(ctx), cp, isStream, cm)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tctx, input = onGraphStart(ctx, input, isStream)\n\t\thaveOnStart = true\n\n\t\tnextTasks, err = r.restoreTasks(ctx, cp.Inputs, cp.SkipPreHandler, cp.RerunNodes, isStream, optMap)\n\t\tif err != nil {\n\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"restore tasks fail: %w\", err))\n\t\t}\n\t} else if checkPointID != nil && !forceNewRun {\n\t\tcp, err = getCheckPointFromStore(ctx, *checkPointID, r.checkPointer)\n\t\tif err != nil {\n\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"load checkpoint from store fail: %w\", err))\n\t\t}\n\t\tif cp != nil {\n\t\t\t// load checkpoint from store\n\t\t\tinitialized = true\n\n\t\t\tctx = setStateModifier(ctx, stateModifier)\n\t\t\tctx = setCheckPointToCtx(ctx, cp)\n\n\t\t\tctx, err = r.restoreCheckPointState(ctx, *NewNodePath(), stateModifier, cp, isStream, cm)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\tctx, input = onGraphStart(ctx, input, isStream)\n\t\t\thaveOnStart = true\n\n\t\t\tnextTasks, err = r.restoreTasks(ctx, cp.Inputs, cp.SkipPreHandler, cp.RerunNodes, isStream, optMap)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"restore tasks fail: %w\", err))\n\t\t\t}\n\t\t}\n\t}\n\tif !initialized {\n\t\t// have not inited from checkpoint\n\t\tif r.runCtx != nil {\n\t\t\tctx = r.runCtx(ctx)\n\t\t}\n\n\t\tctx, input = onGraphStart(ctx, input, isStream)\n\t\thaveOnStart = true\n\n\t\tvar isEnd bool\n\t\tnextTasks, result, isEnd, err = r.calculateNextTasks(ctx, []*task{{\n\t\t\tnodeKey: START,\n\t\t\tcall:    r.inputChannels,\n\t\t\toutput:  input,\n\t\t}}, isStream, cm, optMap)\n\t\tif err != nil {\n\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"calculate next tasks fail: %w\", err))\n\t\t}\n\t\tif isEnd {\n\t\t\treturn result, nil\n\t\t}\n\t\tif len(nextTasks) == 0 {\n\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"no tasks to execute after graph start\"))\n\t\t}\n\n\t\tif keys := getHitKey(nextTasks, r.interruptBeforeNodes); len(keys) > 0 {\n\t\t\ttempInfo := newInterruptTempInfo()\n\t\t\ttempInfo.interruptBeforeNodes = append(tempInfo.interruptBeforeNodes, keys...)\n\t\t\treturn nil, r.handleInterrupt(ctx,\n\t\t\t\ttempInfo,\n\t\t\t\tnextTasks,\n\t\t\t\tcm.channels,\n\t\t\t\tisStream,\n\t\t\t\tisSubGraph,\n\t\t\t\twriteToCheckPointID,\n\t\t\t)\n\t\t}\n\t}\n\n\t// used to reporting NoTask error\n\tvar lastCompletedTask []*task\n\n\t// Main execution loop.\n\tfor step := 0; ; step++ {\n\t\t// Check for context cancellation.\n\t\tselect {\n\t\tcase <-ctx.Done():\n\t\t\t_, _ = tm.waitAll()\n\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"context has been canceled: %w\", ctx.Err()))\n\t\tdefault:\n\t\t}\n\t\tif !r.dag && step >= maxSteps {\n\t\t\treturn nil, newGraphRunError(ErrExceedMaxSteps)\n\t\t}\n\n\t\t// 1. submit next tasks\n\t\t// 2. get completed tasks\n\t\t// 3. calculate next tasks\n\n\t\terr = tm.submit(nextTasks)\n\t\tif err != nil {\n\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"failed to submit tasks: %w\", err))\n\t\t}\n\n\t\tvar totalCanceledTasks []*task\n\n\t\tcompletedTasks, canceled, canceledTasks := tm.wait()\n\t\ttotalCanceledTasks = append(totalCanceledTasks, canceledTasks...)\n\t\ttempInfo := newInterruptTempInfo()\n\t\ttempInfo.collectCanceledInfo(canceled, canceledTasks, completedTasks)\n\n\t\terr = r.resolveInterruptCompletedTasks(tempInfo, completedTasks)\n\t\tif err != nil {\n\t\t\treturn nil, err // err has been wrapped\n\t\t}\n\n\t\tif len(tempInfo.subGraphInterrupts)+len(tempInfo.interruptRerunNodes) > 0 {\n\t\t\tvar newCompletedTasks []*task\n\t\t\tnewCompletedTasks, canceledTasks = tm.waitAll()\n\t\t\ttotalCanceledTasks = append(totalCanceledTasks, canceledTasks...)\n\t\t\tfor _, ct := range canceledTasks {\n\t\t\t\t// handle timeout tasks as rerun\n\t\t\t\ttempInfo.interruptRerunNodes = append(tempInfo.interruptRerunNodes, ct.nodeKey)\n\t\t\t}\n\n\t\t\terr = r.resolveInterruptCompletedTasks(tempInfo, newCompletedTasks)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err // err has been wrapped\n\t\t\t}\n\n\t\t\t// subgraph has interrupted\n\t\t\t// save other completed tasks to channel\n\t\t\t// save interrupted subgraph as next task with SkipPreHandler\n\t\t\t// report current graph interrupt info\n\t\t\treturn nil, r.handleInterruptWithSubGraphAndRerunNodes(\n\t\t\t\tctx,\n\t\t\t\ttempInfo,\n\t\t\t\tappend(append(completedTasks, newCompletedTasks...), totalCanceledTasks...), // canceled tasks are handled as rerun\n\t\t\t\twriteToCheckPointID,\n\t\t\t\tisSubGraph,\n\t\t\t\tcm,\n\t\t\t\tisStream,\n\t\t\t)\n\t\t}\n\n\t\tif len(completedTasks) == 0 {\n\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"no tasks to execute, last completed nodes: %v\", printTask(lastCompletedTask)))\n\t\t}\n\t\tlastCompletedTask = completedTasks\n\n\t\tvar isEnd bool\n\t\tnextTasks, result, isEnd, err = r.calculateNextTasks(ctx, completedTasks, isStream, cm, optMap)\n\t\tif err != nil {\n\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"failed to calculate next tasks: %w\", err))\n\t\t}\n\t\tif isEnd {\n\t\t\treturn result, nil\n\t\t}\n\n\t\ttempInfo.interruptBeforeNodes = getHitKey(nextTasks, r.interruptBeforeNodes)\n\n\t\tif len(tempInfo.interruptBeforeNodes) > 0 || len(tempInfo.interruptAfterNodes) > 0 {\n\t\t\tvar newCompletedTasks []*task\n\t\t\tnewCompletedTasks, canceledTasks = tm.waitAll()\n\t\t\ttotalCanceledTasks = append(totalCanceledTasks, canceledTasks...)\n\t\t\tfor _, ct := range canceledTasks {\n\t\t\t\ttempInfo.interruptRerunNodes = append(tempInfo.interruptRerunNodes, ct.nodeKey)\n\t\t\t}\n\n\t\t\terr = r.resolveInterruptCompletedTasks(tempInfo, newCompletedTasks)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err // err has been wrapped\n\t\t\t}\n\n\t\t\tif len(tempInfo.subGraphInterrupts)+len(tempInfo.interruptRerunNodes) > 0 {\n\t\t\t\treturn nil, r.handleInterruptWithSubGraphAndRerunNodes(\n\t\t\t\t\tctx,\n\t\t\t\t\ttempInfo,\n\t\t\t\t\tappend(append(completedTasks, newCompletedTasks...), totalCanceledTasks...),\n\t\t\t\t\twriteToCheckPointID,\n\t\t\t\t\tisSubGraph,\n\t\t\t\t\tcm,\n\t\t\t\t\tisStream,\n\t\t\t\t)\n\t\t\t}\n\n\t\t\tvar newNextTasks []*task\n\t\t\tnewNextTasks, result, isEnd, err = r.calculateNextTasks(ctx, newCompletedTasks, isStream, cm, optMap)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, newGraphRunError(fmt.Errorf(\"failed to calculate next tasks: %w\", err))\n\t\t\t}\n\n\t\t\tif isEnd {\n\t\t\t\treturn result, nil\n\t\t\t}\n\n\t\t\ttempInfo.interruptBeforeNodes = append(tempInfo.interruptBeforeNodes, getHitKey(newNextTasks, r.interruptBeforeNodes)...)\n\n\t\t\t// simple interrupt\n\t\t\treturn nil, r.handleInterrupt(ctx, tempInfo, append(nextTasks, newNextTasks...), cm.channels, isStream, isSubGraph, writeToCheckPointID)\n\t\t}\n\t}\n}\n\nfunc (r *runner) resolveMaxSteps(maxSteps int, opts []Option) (int, error) {\n\tif r.dag {\n\t\tfor i := range opts {\n\t\t\tif opts[i].maxRunSteps > 0 {\n\t\t\t\treturn 0, newGraphRunError(fmt.Errorf(\"cannot set max run steps in dag\"))\n\t\t\t}\n\t\t}\n\t\treturn maxSteps, nil\n\t}\n\tfor i := range opts {\n\t\tif opts[i].maxRunSteps > 0 {\n\t\t\tmaxSteps = opts[i].maxRunSteps\n\t\t}\n\t}\n\tif maxSteps < 1 {\n\t\treturn 0, newGraphRunError(errors.New(\"max run steps limit must be at least 1\"))\n\t}\n\treturn maxSteps, nil\n}\n\nfunc (r *runner) restoreCheckPointState(\n\tctx context.Context,\n\tpath NodePath,\n\tsm StateModifier,\n\tcp *checkpoint,\n\tisStream bool,\n\tcm *channelManager,\n) (context.Context, error) {\n\terr := r.checkPointer.restoreCheckPoint(cp, isStream)\n\tif err != nil {\n\t\treturn ctx, newGraphRunError(fmt.Errorf(\"restore checkpoint fail: %w\", err))\n\t}\n\n\terr = cm.loadChannels(cp.Channels)\n\tif err != nil {\n\t\treturn ctx, newGraphRunError(err)\n\t}\n\tif sm != nil && cp.State != nil {\n\t\terr = sm(ctx, path, cp.State)\n\t\tif err != nil {\n\t\t\treturn ctx, newGraphRunError(fmt.Errorf(\"state modifier fail: %w\", err))\n\t\t}\n\t}\n\tif cp.State != nil {\n\t\tisResumeTarget, hasData, data := GetResumeContext[any](ctx)\n\t\tif isResumeTarget && hasData {\n\t\t\tcp.State = data\n\t\t}\n\n\t\tvar parent *internalState\n\t\tif prev := ctx.Value(stateKey{}); prev != nil {\n\t\t\tif p, ok := prev.(*internalState); ok {\n\t\t\t\tparent = p\n\t\t\t}\n\t\t}\n\n\t\tctx = context.WithValue(ctx, stateKey{}, &internalState{state: cp.State, parent: parent})\n\t}\n\n\treturn ctx, nil\n}\n\nfunc newInterruptTempInfo() *interruptTempInfo {\n\treturn &interruptTempInfo{\n\t\tsubGraphInterrupts:  map[string]*subGraphInterruptError{},\n\t\tinterruptRerunExtra: map[string]any{},\n\t}\n}\n\ntype interruptTempInfo struct {\n\tsubGraphInterrupts   map[string]*subGraphInterruptError\n\tinterruptRerunNodes  []string\n\tinterruptBeforeNodes []string\n\tinterruptAfterNodes  []string\n\tinterruptRerunExtra  map[string]any\n\n\tsignals []*core.InterruptSignal\n}\n\nfunc (ti *interruptTempInfo) collectCanceledInfo(canceled bool, canceledTasks, completedTasks []*task) {\n\tif !canceled {\n\t\treturn\n\t}\n\tif len(canceledTasks) > 0 {\n\t\tfor _, t := range canceledTasks {\n\t\t\tti.interruptRerunNodes = append(ti.interruptRerunNodes, t.nodeKey)\n\t\t}\n\t} else {\n\t\tfor _, t := range completedTasks {\n\t\t\tti.interruptAfterNodes = append(ti.interruptAfterNodes, t.nodeKey)\n\t\t}\n\t}\n}\n\nfunc (r *runner) resolveInterruptCompletedTasks(tempInfo *interruptTempInfo, completedTasks []*task) (err error) {\n\tfor _, completedTask := range completedTasks {\n\t\tif completedTask.err != nil {\n\t\t\tif info := isSubGraphInterrupt(completedTask.err); info != nil {\n\t\t\t\ttempInfo.subGraphInterrupts[completedTask.nodeKey] = info\n\t\t\t\ttempInfo.signals = append(tempInfo.signals, info.signal)\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tire := &core.InterruptSignal{}\n\t\t\tif errors.As(completedTask.err, &ire) {\n\t\t\t\ttempInfo.interruptRerunNodes = append(tempInfo.interruptRerunNodes, completedTask.nodeKey)\n\t\t\t\tif ire.Info != nil {\n\t\t\t\t\ttempInfo.interruptRerunExtra[completedTask.nodeKey] = ire.InterruptInfo.Info\n\t\t\t\t}\n\n\t\t\t\ttempInfo.signals = append(tempInfo.signals, ire)\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\treturn wrapGraphNodeError(completedTask.nodeKey, completedTask.err)\n\t\t}\n\n\t\tfor _, key := range r.interruptAfterNodes {\n\t\t\tif key == completedTask.nodeKey {\n\t\t\t\ttempInfo.interruptAfterNodes = append(tempInfo.interruptAfterNodes, key)\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc getHitKey(tasks []*task, keys []string) []string {\n\tvar ret []string\n\tfor _, t := range tasks {\n\t\tfor _, key := range keys {\n\t\t\tif key == t.nodeKey {\n\t\t\t\tret = append(ret, t.nodeKey)\n\t\t\t}\n\t\t}\n\t}\n\treturn ret\n}\n\nfunc (r *runner) handleInterrupt(\n\tctx context.Context,\n\ttempInfo *interruptTempInfo,\n\tnextTasks []*task,\n\tchannels map[string]channel,\n\tisStream bool,\n\tisSubGraph bool,\n\tcheckPointID *string,\n) error {\n\tcp := &checkpoint{\n\t\tChannels:       channels,\n\t\tInputs:         make(map[string]any),\n\t\tSkipPreHandler: map[string]bool{},\n\t}\n\tif r.runCtx != nil {\n\t\t// current graph has enable state\n\t\tif state, ok := ctx.Value(stateKey{}).(*internalState); ok {\n\t\t\tcp.State = state.state\n\t\t}\n\t}\n\n\tintInfo := &InterruptInfo{\n\t\tState:           cp.State,\n\t\tAfterNodes:      tempInfo.interruptAfterNodes,\n\t\tBeforeNodes:     tempInfo.interruptBeforeNodes,\n\t\tRerunNodes:      tempInfo.interruptRerunNodes,\n\t\tRerunNodesExtra: tempInfo.interruptRerunExtra,\n\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t}\n\n\tvar info any\n\tif cp.State != nil {\n\t\tcopiedState, err := deepCopyState(cp.State)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"failed to copy state: %w\", err)\n\t\t}\n\t\tinfo = copiedState\n\t}\n\n\tis, err := core.Interrupt(ctx, info, nil, tempInfo.signals)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to interrupt: %w\", err)\n\t}\n\n\tcp.InterruptID2Addr, cp.InterruptID2State = core.SignalToPersistenceMaps(is)\n\n\tfor _, t := range nextTasks {\n\t\tcp.Inputs[t.nodeKey] = t.input\n\t}\n\terr = r.checkPointer.convertCheckPoint(cp, isStream)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to convert checkpoint: %w\", err)\n\t}\n\tif isSubGraph {\n\t\treturn &subGraphInterruptError{\n\t\t\tInfo:       intInfo,\n\t\t\tCheckPoint: cp,\n\t\t\tsignal:     is,\n\t\t}\n\t} else if checkPointID != nil {\n\t\terr := r.checkPointer.set(ctx, *checkPointID, cp)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"failed to set checkpoint: %w, checkPointID: %s\", err, *checkPointID)\n\t\t}\n\t}\n\n\tintInfo.InterruptContexts = core.ToInterruptContexts(is, nil)\n\treturn &interruptError{Info: intInfo}\n}\n\n// deepCopyState creates a deep copy of the state using serialization\nfunc deepCopyState(state any) (any, error) {\n\tif state == nil {\n\t\treturn nil, nil\n\t}\n\tserializer := &serialization.InternalSerializer{}\n\tdata, err := serializer.Marshal(state)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to marshal state: %w\", err)\n\t}\n\n\t// Create new instance of the same type\n\tstateType := reflect.TypeOf(state)\n\tif stateType.Kind() == reflect.Ptr {\n\t\tstateType = stateType.Elem()\n\t}\n\tnewState := reflect.New(stateType).Interface()\n\n\tif err := serializer.Unmarshal(data, newState); err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to unmarshal state: %w\", err)\n\t}\n\treturn newState, nil\n}\n\nfunc (r *runner) handleInterruptWithSubGraphAndRerunNodes(\n\tctx context.Context,\n\ttempInfo *interruptTempInfo,\n\tcompleteTasks []*task,\n\tcheckPointID *string,\n\tisSubGraph bool,\n\tcm *channelManager,\n\tisStream bool,\n) error {\n\tvar rerunTasks, subgraphTasks, otherTasks []*task\n\tskipPreHandler := map[string]bool{}\n\tfor _, t := range completeTasks {\n\t\tif _, ok := tempInfo.subGraphInterrupts[t.nodeKey]; ok {\n\t\t\tsubgraphTasks = append(subgraphTasks, t)\n\t\t\tskipPreHandler[t.nodeKey] = true // subgraph won't run pre-handler again, but rerun nodes will\n\t\t\tcontinue\n\t\t}\n\t\trerun := false\n\t\tfor _, key := range tempInfo.interruptRerunNodes {\n\t\t\tif key == t.nodeKey {\n\t\t\t\trerunTasks = append(rerunTasks, t)\n\t\t\t\trerun = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tif !rerun {\n\t\t\totherTasks = append(otherTasks, t)\n\t\t}\n\t}\n\n\t// forward completed tasks\n\ttoValue, controls, err := r.resolveCompletedTasks(ctx, otherTasks, isStream, cm)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to resolve completed tasks in interrupt: %w\", err)\n\t}\n\terr = cm.updateValues(ctx, toValue)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to update values in interrupt: %w\", err)\n\t}\n\terr = cm.updateDependencies(ctx, controls)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to update dependencies in interrupt: %w\", err)\n\t}\n\n\tcp := &checkpoint{\n\t\tChannels:       cm.channels,\n\t\tInputs:         make(map[string]any),\n\t\tSkipPreHandler: skipPreHandler,\n\t\tSubGraphs:      make(map[string]*checkpoint),\n\t}\n\tif r.runCtx != nil {\n\t\t// current graph has enable state\n\t\tif state, ok := ctx.Value(stateKey{}).(*internalState); ok {\n\t\t\tcp.State = state.state\n\t\t}\n\t}\n\n\tintInfo := &InterruptInfo{\n\t\tState:           cp.State,\n\t\tBeforeNodes:     tempInfo.interruptBeforeNodes,\n\t\tAfterNodes:      tempInfo.interruptAfterNodes,\n\t\tRerunNodes:      tempInfo.interruptRerunNodes,\n\t\tRerunNodesExtra: tempInfo.interruptRerunExtra,\n\t\tSubGraphs:       make(map[string]*InterruptInfo),\n\t}\n\n\tvar info any\n\tif cp.State != nil {\n\t\tcopiedState, err_ := deepCopyState(cp.State)\n\t\tif err_ != nil {\n\t\t\treturn fmt.Errorf(\"failed to copy state: %w\", err_)\n\t\t}\n\t\tinfo = copiedState\n\t}\n\n\tis, err := core.Interrupt(ctx, info, nil, tempInfo.signals)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to interrupt: %w\", err)\n\t}\n\n\tcp.InterruptID2Addr, cp.InterruptID2State = core.SignalToPersistenceMaps(is)\n\n\tfor _, t := range subgraphTasks {\n\t\tcp.RerunNodes = append(cp.RerunNodes, t.nodeKey)\n\t\tcp.SubGraphs[t.nodeKey] = tempInfo.subGraphInterrupts[t.nodeKey].CheckPoint\n\t\tintInfo.SubGraphs[t.nodeKey] = tempInfo.subGraphInterrupts[t.nodeKey].Info\n\t}\n\tfor _, t := range rerunTasks {\n\t\tcp.RerunNodes = append(cp.RerunNodes, t.nodeKey)\n\t\tif t.originalInput != nil {\n\t\t\tcp.Inputs[t.nodeKey] = t.originalInput\n\t\t}\n\t}\n\terr = r.checkPointer.convertCheckPoint(cp, isStream)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to convert checkpoint: %w\", err)\n\t}\n\tif isSubGraph {\n\t\treturn &subGraphInterruptError{\n\t\t\tInfo:       intInfo,\n\t\t\tCheckPoint: cp,\n\t\t\tsignal:     is,\n\t\t}\n\t} else if checkPointID != nil {\n\t\terr = r.checkPointer.set(ctx, *checkPointID, cp)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"failed to set checkpoint: %w, checkPointID: %s\", err, *checkPointID)\n\t\t}\n\t}\n\tintInfo.InterruptContexts = core.ToInterruptContexts(is, nil)\n\treturn &interruptError{Info: intInfo}\n}\n\nfunc (r *runner) calculateNextTasks(ctx context.Context, completedTasks []*task, isStream bool, cm *channelManager, optMap map[string][]any) ([]*task, any, bool, error) {\n\twriteChannelValues, controls, err := r.resolveCompletedTasks(ctx, completedTasks, isStream, cm)\n\tif err != nil {\n\t\treturn nil, nil, false, err\n\t}\n\tnodeMap, err := cm.updateAndGet(ctx, writeChannelValues, controls)\n\tif err != nil {\n\t\treturn nil, nil, false, fmt.Errorf(\"failed to update and get channels: %w\", err)\n\t}\n\tvar nextTasks []*task\n\tif len(nodeMap) > 0 {\n\t\t// Check if we've reached the END node.\n\t\tif v, ok := nodeMap[END]; ok {\n\t\t\treturn nil, v, true, nil\n\t\t}\n\n\t\t// Create and submit the next batch of tasks.\n\t\tnextTasks, err = r.createTasks(ctx, nodeMap, optMap)\n\t\tif err != nil {\n\t\t\treturn nil, nil, false, fmt.Errorf(\"failed to create tasks: %w\", err)\n\t\t}\n\t}\n\treturn nextTasks, nil, false, nil\n}\n\nfunc (r *runner) createTasks(ctx context.Context, nodeMap map[string]any, optMap map[string][]any) ([]*task, error) {\n\tvar nextTasks []*task\n\tfor nodeKey, nodeInput := range nodeMap {\n\t\tcall, ok := r.chanSubscribeTo[nodeKey]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"node[%s] has not been registered\", nodeKey)\n\t\t}\n\n\t\tif call.action.nodeInfo != nil && call.action.nodeInfo.compileOption != nil {\n\t\t\tctx = forwardCheckPoint(ctx, nodeKey)\n\t\t}\n\n\t\tnextTasks = append(nextTasks, &task{\n\t\t\tctx:     AppendAddressSegment(ctx, AddressSegmentNode, nodeKey),\n\t\t\tnodeKey: nodeKey,\n\t\t\tcall:    call,\n\t\t\tinput:   nodeInput,\n\t\t\toption:  optMap[nodeKey],\n\t\t})\n\t}\n\treturn nextTasks, nil\n}\n\nfunc getCheckPointInfo(opts ...Option) (checkPointID *string, writeToCheckPointID *string, stateModifier StateModifier, forceNewRun bool) {\n\tfor _, opt := range opts {\n\t\tif opt.checkPointID != nil {\n\t\t\tcheckPointID = opt.checkPointID\n\t\t}\n\t\tif opt.writeToCheckPointID != nil {\n\t\t\twriteToCheckPointID = opt.writeToCheckPointID\n\t\t}\n\t\tif opt.stateModifier != nil {\n\t\t\tstateModifier = opt.stateModifier\n\t\t}\n\t\tforceNewRun = opt.forceNewRun\n\t}\n\tif writeToCheckPointID == nil {\n\t\twriteToCheckPointID = checkPointID\n\t}\n\treturn\n}\n\nfunc (r *runner) restoreTasks(\n\tctx context.Context,\n\tinputs map[string]any,\n\tskipPreHandler map[string]bool,\n\trerunNodes []string,\n\tisStream bool,\n\toptMap map[string][]any) ([]*task, error) {\n\tret := make([]*task, 0, len(inputs))\n\tfor _, key := range rerunNodes {\n\t\tif _, hasInput := inputs[key]; hasInput {\n\t\t\tcontinue\n\t\t}\n\n\t\tcall, ok := r.chanSubscribeTo[key]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"channel[%s] from checkpoint is not registered\", key)\n\t\t}\n\t\tif isStream {\n\t\t\tinputs[key] = call.action.inputEmptyStream()\n\t\t} else {\n\t\t\tinputs[key] = call.action.inputZeroValue()\n\t\t}\n\t}\n\tfor key, input := range inputs {\n\t\tcall, ok := r.chanSubscribeTo[key]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"channel[%s] from checkpoint is not registered\", key)\n\t\t}\n\n\t\tif call.action.nodeInfo != nil && call.action.nodeInfo.compileOption != nil {\n\t\t\t// sub graph\n\t\t\tctx = forwardCheckPoint(ctx, key)\n\t\t}\n\n\t\tnewTask := &task{\n\t\t\tctx:            AppendAddressSegment(ctx, AddressSegmentNode, key),\n\t\t\tnodeKey:        key,\n\t\t\tcall:           call,\n\t\t\tinput:          input,\n\t\t\toption:         nil,\n\t\t\tskipPreHandler: skipPreHandler[key],\n\t\t}\n\t\tif opt, ok := optMap[key]; ok {\n\t\t\tnewTask.option = opt\n\t\t}\n\n\t\tret = append(ret, newTask)\n\t}\n\treturn ret, nil\n}\n\nfunc (r *runner) resolveCompletedTasks(ctx context.Context, completedTasks []*task, isStream bool, cm *channelManager) (map[string]map[string]any, map[string][]string, error) {\n\twriteChannelValues := make(map[string]map[string]any)\n\tnewDependencies := make(map[string][]string)\n\tfor _, t := range completedTasks {\n\t\tfor _, key := range t.call.controls {\n\t\t\tnewDependencies[key] = append(newDependencies[key], t.nodeKey)\n\t\t}\n\n\t\t// update channel & new_next_tasks\n\t\tvs := copyItem(t.output, len(t.call.writeTo)+len(t.call.writeToBranches)*2)\n\t\tnextNodeKeys, err := r.calculateBranch(ctx, t.nodeKey, t.call,\n\t\t\tvs[len(t.call.writeTo)+len(t.call.writeToBranches):], isStream, cm)\n\t\tif err != nil {\n\t\t\treturn nil, nil, fmt.Errorf(\"calculate next step fail, node: %s, error: %w\", t.nodeKey, err)\n\t\t}\n\n\t\tfor _, key := range nextNodeKeys {\n\t\t\tnewDependencies[key] = append(newDependencies[key], t.nodeKey)\n\t\t}\n\t\tnextNodeKeys = append(nextNodeKeys, t.call.writeTo...)\n\n\t\t// If branches generates more than one successor, the inputs need to be copied accordingly.\n\t\tif len(nextNodeKeys) > 0 {\n\t\t\ttoCopyNum := len(nextNodeKeys) - len(t.call.writeTo) - len(t.call.writeToBranches)\n\t\t\tnVs := copyItem(vs[len(t.call.writeTo)+len(t.call.writeToBranches)-1], toCopyNum+1)\n\t\t\tvs = append(vs[:len(t.call.writeTo)+len(t.call.writeToBranches)-1], nVs...)\n\n\t\t\tfor i, next := range nextNodeKeys {\n\t\t\t\tif _, ok := writeChannelValues[next]; !ok {\n\t\t\t\t\twriteChannelValues[next] = make(map[string]any)\n\t\t\t\t}\n\t\t\t\twriteChannelValues[next][t.nodeKey] = vs[i]\n\t\t\t}\n\t\t}\n\t}\n\treturn writeChannelValues, newDependencies, nil\n}\n\nfunc (r *runner) calculateBranch(ctx context.Context, curNodeKey string, startChan *chanCall, input []any, isStream bool, cm *channelManager) ([]string, error) {\n\tif len(input) < len(startChan.writeToBranches) {\n\t\t// unreachable\n\t\treturn nil, errors.New(\"calculate next input length is shorter than branches\")\n\t}\n\n\tret := make([]string, 0, len(startChan.writeToBranches))\n\n\tskippedNodes := make(map[string]struct{})\n\tfor i, branch := range startChan.writeToBranches {\n\t\t// check branch input type if needed\n\t\tvar err error\n\t\tinput[i], err = r.preBranchHandlerManager.handle(curNodeKey, i, input[i], isStream)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"branch[%s]-[%d] pre handler fail: %w\", curNodeKey, branch.idx, err)\n\t\t}\n\n\t\t// process branch output\n\t\tvar ws []string\n\t\tif isStream {\n\t\t\tws, err = branch.collect(ctx, input[i].(streamReader))\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"branch collect run error: %w\", err)\n\t\t\t}\n\t\t} else {\n\t\t\tws, err = branch.invoke(ctx, input[i])\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"branch invoke run error: %w\", err)\n\t\t\t}\n\t\t}\n\n\t\tfor node := range branch.endNodes {\n\t\t\tskipped := true\n\t\t\tfor _, w := range ws {\n\t\t\t\tif node == w {\n\t\t\t\t\tskipped = false\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t\tif skipped {\n\t\t\t\tskippedNodes[node] = struct{}{}\n\t\t\t}\n\t\t}\n\n\t\tret = append(ret, ws...)\n\t}\n\n\t// When a node has multiple branches,\n\t// there may be a situation where a succeeding node is selected by some branches and discarded by the other branches,\n\t// in which case the succeeding node should not be skipped.\n\tvar skippedNodeList []string\n\tfor _, selected := range ret {\n\t\tif _, ok := skippedNodes[selected]; ok {\n\t\t\tdelete(skippedNodes, selected)\n\t\t}\n\t}\n\tfor skipped := range skippedNodes {\n\t\tskippedNodeList = append(skippedNodeList, skipped)\n\t}\n\n\terr := cm.reportBranch(curNodeKey, skippedNodeList)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn ret, nil\n}\n\nfunc (r *runner) initTaskManager(runWrapper runnableCallWrapper, cancelVal *graphCancelChanVal, opts ...Option) *taskManager {\n\ttm := &taskManager{\n\t\trunWrapper:        runWrapper,\n\t\topts:              opts,\n\t\tneedAll:           !r.eager,\n\t\tdone:              internal.NewUnboundedChan[*task](),\n\t\trunningTasks:      make(map[string]*task),\n\t\tpersistRerunInput: cancelVal != nil,\n\t}\n\tif cancelVal != nil {\n\t\ttm.cancelCh = cancelVal.ch\n\t}\n\treturn tm\n}\n\nfunc (r *runner) initChannelManager(isStream bool) *channelManager {\n\tbuilder := r.chanBuilder\n\tif builder == nil {\n\t\tbuilder = pregelChannelBuilder\n\t}\n\n\tchs := make(map[string]channel)\n\tfor ch := range r.chanSubscribeTo {\n\t\tchs[ch] = builder(r.controlPredecessors[ch], r.dataPredecessors[ch], r.chanSubscribeTo[ch].action.inputZeroValue, r.chanSubscribeTo[ch].action.inputEmptyStream)\n\t}\n\n\tchs[END] = builder(r.controlPredecessors[END], r.dataPredecessors[END], r.outputZeroValue, r.outputEmptyStream)\n\n\tdataPredecessors := make(map[string]map[string]struct{})\n\tfor k, vs := range r.dataPredecessors {\n\t\tdataPredecessors[k] = make(map[string]struct{})\n\t\tfor _, v := range vs {\n\t\t\tdataPredecessors[k][v] = struct{}{}\n\t\t}\n\t}\n\tcontrolPredecessors := make(map[string]map[string]struct{})\n\tfor k, vs := range r.controlPredecessors {\n\t\tcontrolPredecessors[k] = make(map[string]struct{})\n\t\tfor _, v := range vs {\n\t\t\tcontrolPredecessors[k][v] = struct{}{}\n\t\t}\n\t}\n\n\tfor k, v := range chs {\n\t\tif cfg, ok := r.mergeConfigs[k]; ok {\n\t\t\tv.setMergeConfig(cfg)\n\t\t}\n\t}\n\n\treturn &channelManager{\n\t\tisStream:            isStream,\n\t\tchannels:            chs,\n\t\tsuccessors:          r.successors,\n\t\tdataPredecessors:    dataPredecessors,\n\t\tcontrolPredecessors: controlPredecessors,\n\n\t\tedgeHandlerManager:    r.edgeHandlerManager,\n\t\tpreNodeHandlerManager: r.preNodeHandlerManager,\n\t}\n}\n\nfunc (r *runner) toComposableRunnable() *composableRunnable {\n\tcr := &composableRunnable{\n\t\ti: func(ctx context.Context, input any, opts ...any) (output any, err error) {\n\t\t\ttos, err := convertOption[Option](opts...)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\treturn r.invoke(ctx, input, tos...)\n\t\t},\n\t\tt: func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) {\n\t\t\ttos, err := convertOption[Option](opts...)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\treturn r.transform(ctx, input, tos...)\n\t\t},\n\n\t\tinputType:     r.inputType,\n\t\toutputType:    r.outputType,\n\t\tgenericHelper: r.genericHelper,\n\t\toptionType:    nil, // if option type is nil, graph will transmit all options.\n\t}\n\n\treturn cr\n}\n\nfunc copyItem(item any, n int) []any {\n\tif n < 2 {\n\t\treturn []any{item}\n\t}\n\n\tret := make([]any, n)\n\tif s, ok := item.(streamReader); ok {\n\t\tss := s.copy(n)\n\t\tfor i := range ret {\n\t\t\tret[i] = ss[i]\n\t\t}\n\n\t\treturn ret\n\t}\n\n\tfor i := range ret {\n\t\tret[i] = item\n\t}\n\n\treturn ret\n}\n\nfunc printTask(ts []*task) string {\n\tif len(ts) == 0 {\n\t\treturn \"[]\"\n\t}\n\tsb := strings.Builder{}\n\tsb.WriteString(\"[\")\n\tfor i := 0; i < len(ts)-1; i++ {\n\t\tsb.WriteString(ts[i].nodeKey)\n\t\tsb.WriteString(\", \")\n\t}\n\tsb.WriteString(ts[len(ts)-1].nodeKey)\n\tsb.WriteString(\"]\")\n\treturn sb.String()\n}\n"
  },
  {
    "path": "compose/graph_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"io\"\n\t\"reflect\"\n\t\"sort\"\n\t\"strconv\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestSingleGraph(t *testing.T) {\n\n\tconst (\n\t\tnodeOfModel  = \"model\"\n\t\tnodeOfPrompt = \"prompt\"\n\t)\n\n\tctx := context.Background()\n\tg := NewGraph[map[string]any, *schema.Message]()\n\n\tpt := prompt.FromMessages(schema.FString,\n\t\tschema.UserMessage(\"what's the weather in {location}?\"),\n\t)\n\n\terr := g.AddChatTemplateNode(\"prompt\", pt)\n\tassert.NoError(t, err)\n\n\tcm := &chatModel{\n\t\tmsgs: []*schema.Message{\n\t\t\t{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tContent: \"the weather is good\",\n\t\t\t},\n\t\t},\n\t}\n\n\terr = g.AddChatModelNode(nodeOfModel, cm, WithNodeName(\"MockChatModel\"))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, nodeOfPrompt)\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(nodeOfPrompt, nodeOfModel)\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(nodeOfModel, END)\n\tassert.NoError(t, err)\n\n\tr, err := g.Compile(context.Background(), WithMaxRunSteps(10))\n\tassert.NoError(t, err)\n\n\tin := map[string]any{\"location\": \"beijing\"}\n\t_, err = r.Invoke(ctx, in)\n\tassert.NoError(t, err)\n\n\t// stream\n\ts, err := r.Stream(ctx, in)\n\tassert.NoError(t, err)\n\n\t_, err = concatStreamReader(s)\n\tassert.NoError(t, err)\n\n\tsr, sw := schema.Pipe[map[string]any](1)\n\t_ = sw.Send(in, nil)\n\tsw.Close()\n\n\t// transform\n\ts, err = r.Transform(ctx, sr)\n\tassert.NoError(t, err)\n\n\t_, err = concatStreamReader(s)\n\tassert.NoError(t, err)\n\n\t// error test\n\tin = map[string]any{\"wrong key\": 1}\n\t_, err = r.Invoke(ctx, in)\n\tassert.Errorf(t, err, \"could not find key: location\")\n\n\t_, err = r.Stream(ctx, in)\n\tassert.Errorf(t, err, \"could not find key: location\")\n\n\tsr, sw = schema.Pipe[map[string]any](1)\n\t_ = sw.Send(in, nil)\n\tsw.Close()\n\n\t_, err = r.Transform(ctx, sr)\n\tassert.Errorf(t, err, \"could not find key: location\")\n}\n\ntype person interface {\n\tSay() string\n}\n\ntype doctor struct {\n\tsay string\n}\n\nfunc (d *doctor) Say() string {\n\treturn d.say\n}\n\nfunc TestGraphWithImplementableType(t *testing.T) {\n\n\tconst (\n\t\tnode1 = \"1st\"\n\t\tnode2 = \"2nd\"\n\t)\n\n\tctx := context.Background()\n\n\tg := NewGraph[string, string]()\n\n\terr := g.AddLambdaNode(node1, InvokableLambda(func(ctx context.Context, input string) (output *doctor, err error) {\n\t\treturn &doctor{say: input}, nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = g.AddLambdaNode(node2, InvokableLambda(func(ctx context.Context, input person) (output string, err error) {\n\t\treturn input.Say(), nil\n\t}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, node1)\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(node1, node2)\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(node2, END)\n\tassert.NoError(t, err)\n\n\tr, err := g.Compile(context.Background(), WithMaxRunSteps(10))\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"how are you\", WithRuntimeMaxSteps(1))\n\tassert.Error(t, err)\n\tassert.ErrorContains(t, err, \"exceeds max steps\")\n\n\t_, err = r.Invoke(ctx, \"how are you\", WithRuntimeMaxSteps(1))\n\tassert.Error(t, err)\n\tassert.ErrorContains(t, err, \"exceeds max steps\")\n\n\tout, err := r.Invoke(ctx, \"how are you\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"how are you\", out)\n\n\toutStream, err := r.Stream(ctx, \"i'm fine\")\n\tassert.NoError(t, err)\n\tdefer outStream.Close()\n\n\tsay, err := outStream.Recv()\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"i'm fine\", say)\n}\n\nfunc TestNestedGraph(t *testing.T) {\n\tconst (\n\t\tnodeOfLambda1  = \"lambda1\"\n\t\tnodeOfLambda2  = \"lambda2\"\n\t\tnodeOfSubGraph = \"sub_graph\"\n\t\tnodeOfModel    = \"model\"\n\t\tnodeOfPrompt   = \"prompt\"\n\t)\n\n\tctx := context.Background()\n\tg := NewGraph[string, *schema.Message]()\n\tsg := NewGraph[map[string]any, *schema.Message]()\n\n\tl1 := InvokableLambda[string, map[string]any](\n\t\tfunc(ctx context.Context, input string) (output map[string]any, err error) {\n\t\t\treturn map[string]any{\"location\": input}, nil\n\t\t})\n\n\tl2 := InvokableLambda[*schema.Message, *schema.Message](\n\t\tfunc(ctx context.Context, input *schema.Message) (output *schema.Message, err error) {\n\t\t\tinput.Content = fmt.Sprintf(\"after lambda 2: %s\", input.Content)\n\t\t\treturn input, nil\n\t\t})\n\n\tpt := prompt.FromMessages(schema.FString,\n\t\tschema.UserMessage(\"what's the weather in {location}?\"),\n\t)\n\n\terr := sg.AddChatTemplateNode(\"prompt\", pt)\n\tassert.NoError(t, err)\n\n\tcm := &chatModel{\n\t\tmsgs: []*schema.Message{\n\t\t\t{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tContent: \"the weather is good\",\n\t\t\t},\n\t\t},\n\t}\n\n\terr = sg.AddChatModelNode(nodeOfModel, cm, WithNodeName(\"MockChatModel\"))\n\tassert.NoError(t, err)\n\n\terr = sg.AddEdge(START, nodeOfPrompt)\n\tassert.NoError(t, err)\n\n\terr = sg.AddEdge(nodeOfPrompt, nodeOfModel)\n\tassert.NoError(t, err)\n\n\terr = sg.AddEdge(nodeOfModel, END)\n\tassert.NoError(t, err)\n\n\terr = g.AddLambdaNode(nodeOfLambda1, l1, WithNodeName(\"Lambda1\"))\n\tassert.NoError(t, err)\n\n\terr = g.AddGraphNode(nodeOfSubGraph, sg, WithNodeName(\"SubGraphName\"))\n\tassert.NoError(t, err)\n\n\terr = g.AddLambdaNode(nodeOfLambda2, l2, WithNodeName(\"Lambda2\"))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, nodeOfLambda1)\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(nodeOfLambda1, nodeOfSubGraph)\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(nodeOfSubGraph, nodeOfLambda2)\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(nodeOfLambda2, END)\n\tassert.NoError(t, err)\n\n\tr, err := g.Compile(context.Background(),\n\t\tWithMaxRunSteps(10),\n\t\tWithGraphName(\"GraphName\"),\n\t)\n\tassert.NoError(t, err)\n\n\tck := \"depth\"\n\tcb := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tv, ok := ctx.Value(ck).(int)\n\t\t\tif ok {\n\t\t\t\tv++\n\t\t\t}\n\n\t\t\treturn context.WithValue(ctx, ck, v)\n\t\t}).\n\t\tOnStartWithStreamInputFn(func(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {\n\t\t\tinput.Close()\n\n\t\t\tv, ok := ctx.Value(ck).(int)\n\t\t\tif ok {\n\t\t\t\tv++\n\t\t\t}\n\n\t\t\treturn context.WithValue(ctx, ck, v)\n\t\t}).Build()\n\n\t// invoke\n\t_, err = r.Invoke(ctx, \"london\", WithCallbacks(cb))\n\tassert.NoError(t, err)\n\n\t// stream\n\trs, err := r.Stream(ctx, \"london\", WithCallbacks(cb))\n\tassert.NoError(t, err)\n\tfor {\n\t\t_, err = rs.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\n\t\tassert.NoError(t, err)\n\t}\n\n\t// collect\n\tsr, sw := schema.Pipe[string](5)\n\t_ = sw.Send(\"london\", nil)\n\tsw.Close()\n\n\t_, err = r.Collect(ctx, sr, WithCallbacks(cb))\n\tassert.NoError(t, err)\n\n\t// transform\n\tsr, sw = schema.Pipe[string](5)\n\t_ = sw.Send(\"london\", nil)\n\tsw.Close()\n\n\trt, err := r.Transform(ctx, sr, WithCallbacks(cb))\n\tassert.NoError(t, err)\n\tfor {\n\t\t_, err = rt.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\n\t\tassert.NoError(t, err)\n\t}\n}\n\ntype chatModel struct {\n\tmsgs []*schema.Message\n}\n\nfunc (c *chatModel) BindTools(tools []*schema.ToolInfo) error {\n\treturn nil\n}\n\nfunc (c *chatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\treturn c.msgs[0], nil\n}\n\nfunc (c *chatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tsr, sw := schema.Pipe[*schema.Message](len(c.msgs))\n\tgo func() {\n\t\tfor _, msg := range c.msgs {\n\t\t\tsw.Send(msg, nil)\n\t\t}\n\t\tsw.Close()\n\t}()\n\treturn sr, nil\n}\n\nfunc TestValidate(t *testing.T) {\n\t// test unmatched nodes\n\tg := NewGraph[string, string]()\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return \"\", nil }))\n\tassert.NoError(t, err)\n\n\terr = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input int) (output string, err error) { return \"\", nil }))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.ErrorContains(t, err, \"graph edge[1]-[2]: start node's output type[string] and end node's input type[int] mismatch\")\n\n\t// test unmatched passthrough node\n\tg = NewGraph[string, string]()\n\terr = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return \"\", nil }))\n\tassert.NoError(t, err)\n\n\terr = g.AddPassthroughNode(\"2\")\n\tassert.NoError(t, err)\n\n\terr = g.AddLambdaNode(\"3\", InvokableLambda(func(ctx context.Context, input int) (output string, err error) { return \"\", nil }))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(\"2\", \"3\")\n\tassert.ErrorContains(t, err, \"graph edge[2]-[3]: start node's output type[string] and end node's input type[int] mismatch\")\n\n\t// test may matched passthrough\n\tg2 := NewGraph[any, string]()\n\terr = g2.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input any) (output any, err error) { return input, nil }))\n\tassert.NoError(t, err)\n\terr = g2.AddPassthroughNode(\"2\")\n\tassert.NoError(t, err)\n\terr = g2.AddLambdaNode(\"3\", InvokableLambda(func(ctx context.Context, input int) (output string, err error) { return strconv.Itoa(input), nil }))\n\tassert.NoError(t, err)\n\terr = g2.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g2.AddEdge(\"2\", \"3\")\n\tassert.NoError(t, err)\n\terr = g2.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = g2.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\tru, err := g2.Compile(context.Background())\n\tassert.NoError(t, err)\n\t// success\n\tresult, err := ru.Invoke(context.Background(), 1)\n\tassert.NoError(t, err)\n\tassert.Equal(t, result, \"1\")\n\t// fail\n\t_, err = ru.Invoke(context.Background(), \"1\")\n\tassert.ErrorContains(t, err, \"runtime type check\")\n\n\t// test unmatched graph type\n\tg = NewGraph[string, string]()\n\terr = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input int) (output string, err error) { return \"\", nil }))\n\tassert.NoError(t, err)\n\n\terr = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output int, err error) { return 0, nil }))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(START, \"1\")\n\tassert.ErrorContains(t, err, \"graph edge[start]-[1]: start node's output type[string] and end node's input type[int] mismatch\")\n\n\t// sub graph implement\n\ttype A interface {\n\t\tA()\n\t}\n\ttype B interface {\n\t\tB()\n\t}\n\n\ttype AB interface {\n\t\tA\n\t\tB\n\t}\n\tlA := InvokableLambda(func(ctx context.Context, input A) (output string, err error) { return \"\", nil })\n\tlB := InvokableLambda(func(ctx context.Context, input B) (output string, err error) { return \"\", nil })\n\tlAB := InvokableLambda(func(ctx context.Context, input string) (output AB, err error) { return nil, nil })\n\n\tp := NewParallel().AddLambda(\"1\", lA).AddLambda(\"2\", lB)\n\tc := NewChain[string, map[string]any]().AppendLambda(lAB).AppendParallel(p)\n\t_, err = c.Compile(context.Background())\n\tassert.NoError(t, err)\n\n\t// error usage\n\tp = NewParallel().AddLambda(\"1\", lA).AddLambda(\"2\", lAB)\n\tc = NewChain[string, map[string]any]().AppendParallel(p)\n\t_, err = c.Compile(context.Background())\n\tassert.ErrorContains(t, err, \"add parallel edge failed, from=start, to=node_0_parallel_0, err: graph edge[start]-[node_0_parallel_0]: start node's output type[string] and end node's input type[compose.A] mismatch\")\n\n\t// test graph output type check\n\tgg := NewGraph[string, A]()\n\terr = gg.AddLambdaNode(\"nodeA\", InvokableLambda(func(ctx context.Context, input string) (output A, err error) { return nil, nil }))\n\tassert.NoError(t, err)\n\n\terr = gg.AddLambdaNode(\"nodeA2\", InvokableLambda(func(ctx context.Context, input string) (output A, err error) { return nil, nil }))\n\tassert.NoError(t, err)\n\n\terr = gg.AddLambdaNode(\"nodeB\", InvokableLambda(func(ctx context.Context, input string) (output B, err error) { return nil, nil }))\n\tassert.NoError(t, err)\n\n\terr = gg.AddEdge(\"nodeA\", END)\n\tassert.NoError(t, err)\n\n\terr = gg.AddEdge(\"nodeB\", END)\n\tassert.ErrorContains(t, err, \"graph edge[nodeB]-[end]: start node's output type[compose.B] and end node's input type[compose.A] mismatch\")\n\n\terr = gg.AddEdge(\"nodeA2\", END)\n\tassert.ErrorContains(t, err, \"graph edge[nodeB]-[end]: start node's output type[compose.B] and end node's input type[compose.A] mismatch\")\n\n\t// test any type\n\tanyG := NewGraph[any, string]()\n\terr = anyG.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node1\", nil }))\n\tassert.NoError(t, err)\n\n\terr = anyG.AddLambdaNode(\"node2\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node2\", nil }))\n\tassert.NoError(t, err)\n\n\terr = anyG.AddEdge(START, \"node1\")\n\tassert.NoError(t, err)\n\n\terr = anyG.AddEdge(\"node1\", \"node2\")\n\tassert.NoError(t, err)\n\n\terr = anyG.AddEdge(\"node2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tr, err := anyG.Compile(context.Background())\n\tassert.NoError(t, err)\n\n\tresult, err = r.Invoke(context.Background(), \"start\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"startnode1node2\", result)\n\n\tstreamResult, err := r.Stream(context.Background(), \"start\")\n\tassert.NoError(t, err)\n\n\tresult = \"\"\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err != nil {\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t}\n\t\tresult += chunk\n\t}\n\n\tassert.Equal(t, \"startnode1node2\", result)\n\n\t// test any type runtime error\n\tanyG = NewGraph[any, string]()\n\terr = anyG.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return 123, nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = anyG.AddLambdaNode(\"node2\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node2\", nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = anyG.AddEdge(START, \"node1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = anyG.AddEdge(\"node1\", \"node2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = anyG.AddEdge(\"node2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tr, err = anyG.Compile(context.Background())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\t_, err = r.Invoke(context.Background(), \"start\")\n\tif err == nil || !strings.Contains(err.Error(), \"runtime\") {\n\t\tt.Fatal(\"test any type runtime error fail, error is nil or error doesn't contain key word runtime\")\n\t}\n\t_, err = r.Stream(context.Background(), \"start\")\n\tif err == nil || !strings.Contains(err.Error(), \"runtime\") {\n\t\tt.Fatal(\"test any type runtime error fail, error is nil or error doesn't contain key word runtime\")\n\t}\n\n\t// test branch any type\n\t// success\n\tg = NewGraph[string, string]()\n\terr = g.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node1\", nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node2\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node2\", nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node3\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node3\", nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddBranch(\"node1\", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\treturn \"node2\", nil\n\t}, map[string]bool{\"node2\": true, \"node3\": true}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"node1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node3\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\trr, err := g.Compile(context.Background())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tret, err := rr.Invoke(context.Background(), \"start\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif ret != \"startnode1node2\" {\n\t\tt.Fatal(\"test branch any type fail, result is unexpected\")\n\t}\n\tstreamResult, err = rr.Stream(context.Background(), \"start\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tret, err = concatStreamReader(streamResult)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif ret != \"startnode1node2\" {\n\t\tt.Fatal(\"test branch any type fail, result is unexpected\")\n\t}\n\t// fail\n\tg = NewGraph[string, string]()\n\terr = g.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return 1 /*error type*/, nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node2\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node2\", nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node3\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node3\", nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddBranch(\"node1\", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\treturn \"node2\", nil\n\t}, map[string]bool{\"node2\": true, \"node3\": true}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"node1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node3\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\trr, err = g.Compile(context.Background())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\t_, err = rr.Invoke(context.Background(), \"start\")\n\tif err == nil || !strings.Contains(err.Error(), \"runtime\") {\n\t\tt.Fatal(\"test branch any type fail, haven't report runtime error\")\n\t}\n\t_, err = rr.Stream(context.Background(), \"start\")\n\tif err == nil || !strings.Contains(err.Error(), \"runtime\") {\n\t\tt.Fatal(\"test branch any type fail, haven't report runtime error\")\n\t}\n}\n\nfunc TestValidateMultiAnyValueBranch(t *testing.T) {\n\t// success\n\tg := NewGraph[string, map[string]any]()\n\terr := g.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node1\", nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node2\", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\treturn map[string]any{\"node2\": true}, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node3\", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\treturn map[string]any{\"node3\": true}, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node4\", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\treturn map[string]any{\"node4\": true}, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node5\", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\treturn map[string]any{\"node5\": true}, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddBranch(\"node1\", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\treturn \"node2\", nil\n\t}, map[string]bool{\"node2\": true, \"node3\": true}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddBranch(\"node1\", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\treturn \"node4\", nil\n\t}, map[string]bool{\"node4\": true, \"node5\": true}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"node1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node3\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node4\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node5\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\trr, err := g.Compile(context.Background())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tret, err := rr.Invoke(context.Background(), \"start\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif !ret[\"node2\"].(bool) || !ret[\"node4\"].(bool) {\n\t\tt.Fatal(\"test branch any type fail, result is unexpected\")\n\t}\n\tstreamResult, err := rr.Stream(context.Background(), \"start\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tret, err = concatStreamReader(streamResult)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif !ret[\"node2\"].(bool) || !ret[\"node4\"].(bool) {\n\t\tt.Fatal(\"test branch any type fail, result is unexpected\")\n\t}\n\n\t// fail\n\tg = NewGraph[string, map[string]any]()\n\terr = g.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node1\", nil }))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node2\", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\treturn map[string]any{\"node2\": true}, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node3\", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\treturn map[string]any{\"node3\": true}, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node4\", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\treturn map[string]any{\"node4\": true}, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node5\", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\treturn map[string]any{\"node5\": true}, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddBranch(\"node1\", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\treturn \"node2\", nil\n\t}, map[string]bool{\"node2\": true, \"node3\": true}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddBranch(\"node1\", NewGraphBranch(func(ctx context.Context, in int /*error type*/) (endNode string, err error) {\n\t\treturn \"node4\", nil\n\t}, map[string]bool{\"node4\": true, \"node5\": true}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"node1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node3\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node4\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node5\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\trr, err = g.Compile(context.Background())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\t_, err = rr.Invoke(context.Background(), \"start\")\n\tif err == nil || !strings.Contains(err.Error(), \"runtime\") {\n\t\tt.Fatal(\"test multi branch any type fail, haven't report runtime error\")\n\t}\n\t_, err = rr.Stream(context.Background(), \"start\")\n\tif err == nil || !strings.Contains(err.Error(), \"runtime\") {\n\t\tt.Fatal(\"test multi branch any type fail, haven't report runtime error\")\n\t}\n}\n\nfunc TestAnyTypeWithKey(t *testing.T) {\n\tg := NewGraph[any, map[string]any]()\n\terr := g.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node1\", nil }), WithInputKey(\"node1\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node2\", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + \"node2\", nil }), WithOutputKey(\"node2\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"node1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node1\", \"node2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tr, err := g.Compile(context.Background())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tresult, err := r.Invoke(context.Background(), map[string]any{\"node1\": \"start\"})\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif result[\"node2\"] != \"startnode1node2\" {\n\t\tt.Fatal(\"test any type with key fail, result is unexpected\")\n\t}\n\n\tstreamResult, err := r.Stream(context.Background(), map[string]any{\"node1\": \"start\"})\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tret, err := concatStreamReader(streamResult)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif ret[\"node2\"] != \"startnode1node2\" {\n\t\tt.Fatal(\"test any type with key fail, result is unexpected\")\n\t}\n}\n\nfunc TestInputKey(t *testing.T) {\n\tg := NewGraph[map[string]any, map[string]any]()\n\terr := g.AddChatTemplateNode(\"1\", prompt.FromMessages(schema.FString, schema.UserMessage(\"{var1}\")), WithOutputKey(\"1\"), WithInputKey(\"1\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddChatTemplateNode(\"2\", prompt.FromMessages(schema.FString, schema.UserMessage(\"{var2}\")), WithOutputKey(\"2\"), WithInputKey(\"2\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddChatTemplateNode(\"3\", prompt.FromMessages(schema.FString, schema.UserMessage(\"{var3}\")), WithOutputKey(\"3\"), WithInputKey(\"3\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"3\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"1\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"3\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tr, err := g.Compile(context.Background(), WithMaxRunSteps(100))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tctx := context.Background()\n\tresult, err := r.Invoke(ctx, map[string]any{\n\t\t\"1\": map[string]any{\"var1\": \"a\"},\n\t\t\"2\": map[string]any{\"var2\": \"b\"},\n\t\t\"3\": map[string]any{\"var3\": \"c\"},\n\t})\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif result[\"1\"].([]*schema.Message)[0].Content != \"a\" ||\n\t\tresult[\"2\"].([]*schema.Message)[0].Content != \"b\" ||\n\t\tresult[\"3\"].([]*schema.Message)[0].Content != \"c\" {\n\t\tt.Fatal(\"invoke different\")\n\t}\n\n\tsr, sw := schema.Pipe[map[string]any](10)\n\tsw.Send(map[string]any{\"1\": map[string]any{\"var1\": \"a\"}}, nil)\n\tsw.Send(map[string]any{\"2\": map[string]any{\"var2\": \"b\"}}, nil)\n\tsw.Send(map[string]any{\"3\": map[string]any{\"var3\": \"c\"}}, nil)\n\tsw.Close()\n\n\tstreamResult, err := r.Transform(ctx, sr)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tdefer streamResult.Close()\n\n\tresult = make(map[string]any)\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tfor k, v := range chunk {\n\t\t\tresult[k] = v\n\t\t}\n\t}\n\tif result[\"1\"].([]*schema.Message)[0].Content != \"a\" ||\n\t\tresult[\"2\"].([]*schema.Message)[0].Content != \"b\" ||\n\t\tresult[\"3\"].([]*schema.Message)[0].Content != \"c\" {\n\t\tt.Fatal(\"transform different\")\n\t}\n}\n\nfunc TestTransferTask(t *testing.T) {\n\tin := [][]string{\n\t\t{\n\t\t\t\"1\",\n\t\t\t\"2\",\n\t\t},\n\t\t{\n\t\t\t\"3\",\n\t\t\t\"4\",\n\t\t\t\"5\",\n\t\t\t\"6\",\n\t\t},\n\t\t{\n\t\t\t\"5\",\n\t\t\t\"6\",\n\t\t\t\"7\",\n\t\t},\n\t\t{\n\t\t\t\"7\",\n\t\t\t\"8\",\n\t\t},\n\t\t{\n\t\t\t\"8\",\n\t\t},\n\t}\n\tinvertedEdges := map[string][]string{\n\t\t\"1\": {\"3\", \"4\"},\n\t\t\"2\": {\"5\", \"6\"},\n\t\t\"3\": {\"5\"},\n\t\t\"4\": {\"6\"},\n\t\t\"5\": {\"7\"},\n\t\t\"7\": {\"8\"},\n\t}\n\tin = transferTask(in, invertedEdges)\n\n\tif !reflect.DeepEqual(\n\t\t[][]string{\n\t\t\t{\n\t\t\t\t\"1\",\n\t\t\t},\n\t\t\t{\n\t\t\t\t\"3\",\n\t\t\t\t\"2\",\n\t\t\t},\n\t\t\t{\n\t\t\t\t\"5\",\n\t\t\t},\n\t\t\t{\n\t\t\t\t\"7\",\n\t\t\t\t\"4\",\n\t\t\t},\n\t\t\t{\n\t\t\t\t\"8\",\n\t\t\t\t\"6\",\n\t\t\t},\n\t\t}, in) {\n\t\tt.Fatal(\"not equal\")\n\t}\n}\n\nfunc TestPregelEnd(t *testing.T) {\n\tg := NewGraph[string, string]()\n\terr := g.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"node1\", nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddLambdaNode(\"node2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"node2\", nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"node1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node1\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node1\", \"node2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node2\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\trunner, err := g.Compile(context.Background())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tout, err := runner.Invoke(context.Background(), \"\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif out != \"node1\" {\n\t\tt.Fatal(\"graph output is unexpected\")\n\t}\n}\n\ntype cb struct {\n\tgInfo *GraphInfo\n}\n\nfunc (c *cb) OnFinish(ctx context.Context, info *GraphInfo) {\n\tc.gInfo = info\n}\n\nfunc TestGraphCompileCallback(t *testing.T) {\n\tt.Run(\"graph compile callback\", func(t *testing.T) {\n\t\ttype s struct{}\n\n\t\tg := NewGraph[map[string]any, map[string]any](WithGenLocalState(func(ctx context.Context) *s { return &s{} }))\n\n\t\tlambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\treturn \"node1\", nil\n\t\t})\n\t\tlambdaOpts := []GraphAddNodeOpt{WithNodeName(\"lambda_1\"), WithInputKey(\"input_key\")}\n\t\terr := g.AddLambdaNode(\"node1\", lambda, lambdaOpts...)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddPassthroughNode(\"pass1\")\n\t\tassert.NoError(t, err)\n\t\terr = g.AddPassthroughNode(\"pass2\")\n\t\tassert.NoError(t, err)\n\n\t\tcondition := func(ctx context.Context, input string) (string, error) {\n\t\t\treturn input, nil\n\t\t}\n\n\t\tbranch := NewGraphBranch(condition, map[string]bool{\"pass1\": true, \"pass2\": true})\n\t\terr = g.AddBranch(\"node1\", branch)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddEdge(START, \"node1\")\n\t\tassert.NoError(t, err)\n\n\t\tlambda2 := InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\treturn \"node2\", nil\n\t\t})\n\t\tlambdaOpts2 := []GraphAddNodeOpt{WithNodeName(\"lambda_2\")}\n\t\tsubSubGraph := NewGraph[string, string]()\n\t\terr = subSubGraph.AddLambdaNode(\"sub1\", lambda2, lambdaOpts2...)\n\t\tassert.NoError(t, err)\n\t\terr = subSubGraph.AddEdge(START, \"sub1\")\n\t\tassert.NoError(t, err)\n\t\terr = subSubGraph.AddEdge(\"sub1\", END)\n\t\tassert.NoError(t, err)\n\n\t\tsubGraph := NewGraph[string, string]()\n\t\tvar ssGraphCompileOpts []GraphCompileOption\n\t\tssGraphOpts := []GraphAddNodeOpt{WithGraphCompileOptions(ssGraphCompileOpts...)}\n\t\terr = subGraph.AddGraphNode(\"sub_sub_1\", subSubGraph, ssGraphOpts...)\n\t\tassert.NoError(t, err)\n\t\terr = subGraph.AddEdge(START, \"sub_sub_1\")\n\t\tassert.NoError(t, err)\n\t\terr = subGraph.AddEdge(\"sub_sub_1\", END)\n\t\tassert.NoError(t, err)\n\n\t\tsubGraphCompileOpts := []GraphCompileOption{WithMaxRunSteps(2), WithGraphName(\"sub_graph\")}\n\t\tsubGraphOpts := []GraphAddNodeOpt{WithGraphCompileOptions(subGraphCompileOpts...)}\n\t\terr = g.AddGraphNode(\"sub_graph\", subGraph, subGraphOpts...)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddEdge(\"pass1\", \"sub_graph\")\n\t\tassert.NoError(t, err)\n\t\terr = g.AddEdge(\"pass2\", \"sub_graph\")\n\t\tassert.NoError(t, err)\n\n\t\tlambda3 := InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\treturn \"node3\", nil\n\t\t})\n\t\tlambdaOpts3 := []GraphAddNodeOpt{WithNodeName(\"lambda_3\"), WithOutputKey(\"lambda_3\")}\n\t\terr = g.AddLambdaNode(\"node3\", lambda3, lambdaOpts3...)\n\t\tassert.NoError(t, err)\n\n\t\tlambda4 := InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\treturn \"node4\", nil\n\t\t})\n\t\tlambdaOpts4 := []GraphAddNodeOpt{WithNodeName(\"lambda_4\"), WithOutputKey(\"lambda_4\")}\n\t\terr = g.AddLambdaNode(\"node4\", lambda4, lambdaOpts4...)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddEdge(\"sub_graph\", \"node3\")\n\t\tassert.NoError(t, err)\n\t\terr = g.AddEdge(\"sub_graph\", \"node4\")\n\t\tassert.NoError(t, err)\n\t\terr = g.AddEdge(\"node3\", END)\n\t\tassert.NoError(t, err)\n\t\terr = g.AddEdge(\"node4\", END)\n\t\tassert.NoError(t, err)\n\n\t\tc := &cb{}\n\t\topt := []GraphCompileOption{WithGraphCompileCallbacks(c), WithGraphName(\"top_level\")}\n\t\t_, err = g.Compile(context.Background(), opt...)\n\t\tassert.NoError(t, err)\n\t\texpected := &GraphInfo{\n\t\t\tCompileOptions: opt,\n\t\t\tNodes: map[string]GraphNodeInfo{\n\t\t\t\t\"node1\": {\n\t\t\t\t\tComponent:        ComponentOfLambda,\n\t\t\t\t\tInstance:         lambda,\n\t\t\t\t\tGraphAddNodeOpts: lambdaOpts,\n\t\t\t\t\tInputType:        reflect.TypeOf(\"\"),\n\t\t\t\t\tOutputType:       reflect.TypeOf(\"\"),\n\t\t\t\t\tName:             \"lambda_1\",\n\t\t\t\t\tInputKey:         \"input_key\",\n\t\t\t\t},\n\t\t\t\t\"pass1\": {\n\t\t\t\t\tComponent:  ComponentOfPassthrough,\n\t\t\t\t\tInputType:  reflect.TypeOf(\"\"),\n\t\t\t\t\tOutputType: reflect.TypeOf(\"\"),\n\t\t\t\t\tName:       \"\",\n\t\t\t\t},\n\t\t\t\t\"pass2\": {\n\t\t\t\t\tComponent:  ComponentOfPassthrough,\n\t\t\t\t\tInputType:  reflect.TypeOf(\"\"),\n\t\t\t\t\tOutputType: reflect.TypeOf(\"\"),\n\t\t\t\t\tName:       \"\",\n\t\t\t\t},\n\t\t\t\t\"sub_graph\": {\n\t\t\t\t\tComponent:        ComponentOfGraph,\n\t\t\t\t\tInstance:         subGraph,\n\t\t\t\t\tGraphAddNodeOpts: subGraphOpts,\n\t\t\t\t\tInputType:        reflect.TypeOf(\"\"),\n\t\t\t\t\tOutputType:       reflect.TypeOf(\"\"),\n\t\t\t\t\tName:             \"\",\n\t\t\t\t\tGraphInfo: &GraphInfo{\n\t\t\t\t\t\tCompileOptions: subGraphCompileOpts,\n\t\t\t\t\t\tNodes: map[string]GraphNodeInfo{\n\t\t\t\t\t\t\t\"sub_sub_1\": {\n\t\t\t\t\t\t\t\tComponent:        ComponentOfGraph,\n\t\t\t\t\t\t\t\tInstance:         subSubGraph,\n\t\t\t\t\t\t\t\tGraphAddNodeOpts: ssGraphOpts,\n\t\t\t\t\t\t\t\tInputType:        reflect.TypeOf(\"\"),\n\t\t\t\t\t\t\t\tOutputType:       reflect.TypeOf(\"\"),\n\t\t\t\t\t\t\t\tName:             \"\",\n\t\t\t\t\t\t\t\tGraphInfo: &GraphInfo{\n\t\t\t\t\t\t\t\t\tCompileOptions: ssGraphCompileOpts,\n\t\t\t\t\t\t\t\t\tNodes: map[string]GraphNodeInfo{\n\t\t\t\t\t\t\t\t\t\t\"sub1\": {\n\t\t\t\t\t\t\t\t\t\t\tComponent:        ComponentOfLambda,\n\t\t\t\t\t\t\t\t\t\t\tInstance:         lambda2,\n\t\t\t\t\t\t\t\t\t\t\tGraphAddNodeOpts: lambdaOpts2,\n\t\t\t\t\t\t\t\t\t\t\tInputType:        reflect.TypeOf(\"\"),\n\t\t\t\t\t\t\t\t\t\t\tOutputType:       reflect.TypeOf(\"\"),\n\t\t\t\t\t\t\t\t\t\t\tName:             \"lambda_2\",\n\t\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\tEdges: map[string][]string{\n\t\t\t\t\t\t\t\t\t\tSTART:  {\"sub1\"},\n\t\t\t\t\t\t\t\t\t\t\"sub1\": {END},\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\tDataEdges: map[string][]string{\n\t\t\t\t\t\t\t\t\t\tSTART:  {\"sub1\"},\n\t\t\t\t\t\t\t\t\t\t\"sub1\": {END},\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\tBranches:   map[string][]GraphBranch{},\n\t\t\t\t\t\t\t\t\tInputType:  reflect.TypeOf(\"\"),\n\t\t\t\t\t\t\t\t\tOutputType: reflect.TypeOf(\"\"),\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tEdges: map[string][]string{\n\t\t\t\t\t\t\tSTART:       {\"sub_sub_1\"},\n\t\t\t\t\t\t\t\"sub_sub_1\": {END},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tDataEdges: map[string][]string{\n\t\t\t\t\t\t\tSTART:       {\"sub_sub_1\"},\n\t\t\t\t\t\t\t\"sub_sub_1\": {END},\n\t\t\t\t\t\t},\n\t\t\t\t\t\tBranches:   map[string][]GraphBranch{},\n\t\t\t\t\t\tInputType:  reflect.TypeOf(\"\"),\n\t\t\t\t\t\tOutputType: reflect.TypeOf(\"\"),\n\t\t\t\t\t\tName:       \"sub_graph\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t\"node3\": {\n\t\t\t\t\tComponent:        ComponentOfLambda,\n\t\t\t\t\tInstance:         lambda3,\n\t\t\t\t\tGraphAddNodeOpts: lambdaOpts3,\n\t\t\t\t\tInputType:        reflect.TypeOf(\"\"),\n\t\t\t\t\tOutputType:       reflect.TypeOf(\"\"),\n\t\t\t\t\tName:             \"lambda_3\",\n\t\t\t\t\tOutputKey:        \"lambda_3\",\n\t\t\t\t},\n\t\t\t\t\"node4\": {\n\t\t\t\t\tComponent:        ComponentOfLambda,\n\t\t\t\t\tInstance:         lambda4,\n\t\t\t\t\tGraphAddNodeOpts: lambdaOpts4,\n\t\t\t\t\tInputType:        reflect.TypeOf(\"\"),\n\t\t\t\t\tOutputType:       reflect.TypeOf(\"\"),\n\t\t\t\t\tName:             \"lambda_4\",\n\t\t\t\t\tOutputKey:        \"lambda_4\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tEdges: map[string][]string{\n\t\t\t\tSTART:       {\"node1\"},\n\t\t\t\t\"pass1\":     {\"sub_graph\"},\n\t\t\t\t\"pass2\":     {\"sub_graph\"},\n\t\t\t\t\"sub_graph\": {\"node3\", \"node4\"},\n\t\t\t\t\"node3\":     {END},\n\t\t\t\t\"node4\":     {END},\n\t\t\t},\n\t\t\tDataEdges: map[string][]string{\n\t\t\t\tSTART:       {\"node1\"},\n\t\t\t\t\"pass1\":     {\"sub_graph\"},\n\t\t\t\t\"pass2\":     {\"sub_graph\"},\n\t\t\t\t\"sub_graph\": {\"node3\", \"node4\"},\n\t\t\t\t\"node3\":     {END},\n\t\t\t\t\"node4\":     {END},\n\t\t\t},\n\t\t\tBranches: map[string][]GraphBranch{\n\t\t\t\t\"node1\": {*branch},\n\t\t\t},\n\t\t\tInputType:  reflect.TypeOf(map[string]any{}),\n\t\t\tOutputType: reflect.TypeOf(map[string]any{}),\n\t\t\tName:       \"top_level\",\n\t\t}\n\n\t\tstateFn := c.gInfo.GenStateFn\n\t\tassert.NotNil(t, stateFn)\n\t\tassert.Equal(t, &s{}, stateFn(context.Background()))\n\n\t\tassert.Equal(t, 1, len(c.gInfo.NewGraphOptions))\n\t\tc.gInfo.NewGraphOptions = nil\n\n\t\tc.gInfo.GenStateFn = nil\n\n\t\tactualCompileOptions := newGraphCompileOptions(c.gInfo.CompileOptions...)\n\t\texpectedCompileOptions := newGraphCompileOptions(expected.CompileOptions...)\n\t\tassert.Equal(t, len(expectedCompileOptions.callbacks), len(actualCompileOptions.callbacks))\n\t\tassert.Same(t, expectedCompileOptions.callbacks[0], actualCompileOptions.callbacks[0])\n\t\tactualCompileOptions.callbacks = nil\n\t\tactualCompileOptions.origOpts = nil\n\t\texpectedCompileOptions.callbacks = nil\n\t\texpectedCompileOptions.origOpts = nil\n\t\tassert.Equal(t, expectedCompileOptions, actualCompileOptions)\n\n\t\tc.gInfo.CompileOptions = nil\n\t\texpected.CompileOptions = nil\n\t\tassert.Equal(t, expected.Branches[\"node1\"][0].endNodes, c.gInfo.Branches[\"node1\"][0].endNodes)\n\t\tassert.Equal(t, expected.Branches[\"node1\"][0].inputType, c.gInfo.Branches[\"node1\"][0].inputType)\n\n\t\texpected.Branches[\"node1\"] = []GraphBranch{}\n\t\tc.gInfo.Branches[\"node1\"] = []GraphBranch{}\n\t\tassert.Equal(t, expected, c.gInfo)\n\t})\n}\n\nfunc TestCheckAddEdge(t *testing.T) {\n\tg := NewGraph[string, string]()\n\terr := g.AddPassthroughNode(\"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddPassthroughNode(\"2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"1\", \"2\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"1\", \"2\")\n\tif err == nil {\n\t\tt.Fatal(\"add edge repeatedly haven't report error\")\n\t}\n}\n\nfunc TestStartWithEnd(t *testing.T) {\n\tg := NewGraph[string, string]()\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\treturn END, nil\n\t}, map[string]bool{\"1\": true, END: true}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tr, err := g.Compile(context.Background())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tsr, sw := schema.Pipe[string](1)\n\tsw.Send(\"test\", nil)\n\tsw.Close()\n\tresult, err := r.Transform(context.Background(), sr)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tfor {\n\t\tchunk, err := result.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tif chunk != \"test\" {\n\t\t\tt.Fatal(\"result is out of expect\")\n\t\t}\n\t}\n}\n\nfunc TestToString(t *testing.T) {\n\tps := runTypePregel.String()\n\tassert.Equal(t, \"Pregel\", ps)\n\n\tds := runTypeDAG\n\tassert.Equal(t, \"DAG\", ds.String())\n}\n\nfunc TestInputKeyError(t *testing.T) {\n\tg := NewGraph[map[string]any, string]()\n\terr := g.AddLambdaNode(\"node1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}), WithInputKey(\"node1\"))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"node1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"node1\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tctx := context.Background()\n\tr, err := g.Compile(ctx)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\t// invoke\n\t_, err = r.Invoke(ctx, map[string]any{\"unknown\": \"123\"})\n\tif err == nil || !strings.Contains(err.Error(), \"cannot find input key: node1\") {\n\t\tt.Fatal(\"cannot report input key error correctly\")\n\t}\n\n\t// transform\n\tsr, sw := schema.Pipe[map[string]any](1)\n\tsw.Send(map[string]any{\"unknown\": \"123\"}, nil)\n\tsw.Close()\n\t_, err = r.Transform(ctx, sr)\n\tif err == nil || !strings.Contains(err.Error(), \"stream reader is empty, concat fail\") {\n\t\tt.Fatal(\"cannot report input key error correctly\")\n\t}\n}\n\nfunc TestContextCancel(t *testing.T) {\n\tctx := context.Background()\n\tg := NewGraph[string, string]()\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"1\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tr, err := g.Compile(ctx)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tctx, cancel := context.WithCancel(ctx)\n\tcancel()\n\t_, err = r.Invoke(ctx, \"test\")\n\tif !strings.Contains(err.Error(), \"context has been canceled\") {\n\t\tt.Fatal(\"graph have not returned canceled error\")\n\t}\n}\n\nfunc TestDAGStart(t *testing.T) {\n\tg := NewGraph[map[string]any, map[string]any]()\n\terr := g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input map[string]any) (output map[string]any, err error) {\n\t\treturn map[string]any{\"1\": \"1\"}, nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input map[string]any) (output map[string]any, err error) {\n\t\treturn input, nil\n\t}))\n\tassert.NoError(t, err)\n\terr = g.AddEdge(START, \"1\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", \"2\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(START, \"2\")\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\tr, err := g.Compile(context.Background(), WithNodeTriggerMode(AllPredecessor))\n\tassert.NoError(t, err)\n\tresult, err := r.Invoke(context.Background(), map[string]any{\"start\": \"start\"})\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\"start\": \"start\", \"1\": \"1\"}, result)\n}\n\nfunc concatLambda(s string) *Lambda {\n\treturn InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input + s, nil })\n}\nfunc mapLambda(k, v string) *Lambda {\n\treturn InvokableLambda(func(ctx context.Context, input map[string]string) (output map[string]string, err error) {\n\t\treturn map[string]string{\n\t\t\tk: v,\n\t\t}, nil\n\t})\n}\n\nfunc TestBaseDAGBranch(t *testing.T) {\n\tg := NewGraph[string, string]()\n\n\terr := g.AddLambdaNode(\"1\", concatLambda(\"1\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", concatLambda(\"2\"))\n\tassert.NoError(t, err)\n\terr = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\tif len(in) > 3 {\n\t\t\treturn \"2\", nil\n\t\t}\n\t\treturn \"1\", nil\n\t}, map[string]bool{\"1\": true, \"2\": true}))\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor))\n\tassert.NoError(t, err)\n\tresult, err := r.Invoke(ctx, \"hi\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"hi1\", result)\n}\n\nfunc TestMultiDAGBranch(t *testing.T) {\n\tg := NewGraph[map[string]string, map[string]string]()\n\n\terr := g.AddLambdaNode(\"1\", mapLambda(\"1\", \"1\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", mapLambda(\"2\", \"2\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"3\", mapLambda(\"3\", \"3\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"4\", mapLambda(\"4\", \"4\"))\n\tassert.NoError(t, err)\n\terr = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in map[string]string) (endNode string, err error) {\n\t\tif len(in[\"input\"]) > 3 {\n\t\t\treturn \"2\", nil\n\t\t}\n\t\treturn \"1\", nil\n\t}, map[string]bool{\"1\": true, \"2\": true}))\n\terr = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in map[string]string) (endNode string, err error) {\n\t\tif len(in[\"input\"]) > 3 {\n\t\t\treturn \"4\", nil\n\t\t}\n\t\treturn \"3\", nil\n\t}, map[string]bool{\"3\": true, \"4\": true}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"4\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor))\n\tassert.NoError(t, err)\n\tresult, err := r.Invoke(ctx, map[string]string{\"input\": \"hi\"})\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]string{\n\t\t\"1\": \"1\",\n\t\t\"3\": \"3\",\n\t}, result)\n}\n\nfunc TestCrossDAGBranch(t *testing.T) {\n\tg := NewGraph[map[string]string, map[string]string]()\n\n\terr := g.AddLambdaNode(\"1\", mapLambda(\"1\", \"1\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", mapLambda(\"2\", \"2\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"3\", mapLambda(\"3\", \"3\"))\n\tassert.NoError(t, err)\n\terr = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in map[string]string) (endNode string, err error) {\n\t\tif len(in[\"input\"]) > 3 {\n\t\t\treturn \"2\", nil\n\t\t}\n\t\treturn \"1\", nil\n\t}, map[string]bool{\"1\": true, \"2\": true}))\n\terr = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in map[string]string) (endNode string, err error) {\n\t\tif len(in[\"input\"]) > 3 {\n\t\t\treturn \"3\", nil\n\t\t}\n\t\treturn \"2\", nil\n\t}, map[string]bool{\"2\": true, \"3\": true}))\n\tassert.NoError(t, err)\n\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"2\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor))\n\tassert.NoError(t, err)\n\tresult, err := r.Invoke(ctx, map[string]string{\"input\": \"hi\"})\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]string{\n\t\t\"1\": \"1\",\n\t\t\"2\": \"2\",\n\t}, result)\n}\n\nfunc TestNestedDAGBranch(t *testing.T) {\n\tg := NewGraph[string, string]()\n\n\terr := g.AddLambdaNode(\"1\", concatLambda(\"1\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"2\", concatLambda(\"2\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"3\", concatLambda(\"3\"))\n\tassert.NoError(t, err)\n\terr = g.AddLambdaNode(\"4\", concatLambda(\"4\"))\n\tassert.NoError(t, err)\n\terr = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\tif len(in) > 3 {\n\t\t\treturn \"2\", nil\n\t\t}\n\t\treturn \"1\", nil\n\t}, map[string]bool{\"1\": true, \"2\": true}))\n\terr = g.AddBranch(\"2\", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) {\n\t\tif len(in) > 10 {\n\t\t\treturn \"4\", nil\n\t\t}\n\t\treturn \"3\", nil\n\t}, map[string]bool{\"3\": true, \"4\": true}))\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"1\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"3\", END)\n\tassert.NoError(t, err)\n\terr = g.AddEdge(\"4\", END)\n\tassert.NoError(t, err)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor))\n\tassert.NoError(t, err)\n\tresult, err := r.Invoke(ctx, \"hello\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"hello23\", result)\n\tresult, err = r.Invoke(ctx, \"hi\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"hi1\", result)\n\tresult, err = r.Invoke(ctx, \"hellohello\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"hellohello24\", result)\n}\n\nfunc TestHandlerTypeValidate(t *testing.T) {\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state string) {\n\t\treturn \"\"\n\t}))\n\t// passthrough pre fail\n\terr := g.AddPassthroughNode(\"1\", WithStatePreHandler(func(ctx context.Context, in string, state string) (string, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.ErrorContains(t, err, \"passthrough node[1]'s pre handler type isn't any\")\n\tg.buildError = nil\n\t// passthrough pre fail with input key\n\terr = g.AddPassthroughNode(\"1\", WithStatePreHandler(func(ctx context.Context, in string, state string) (string, error) {\n\t\treturn \"\", nil\n\t}), WithInputKey(\"input\"))\n\tassert.ErrorContains(t, err, \"node[1]'s pre handler type[string] is different from its input type[map[string]interface {}]\")\n\tg.buildError = nil\n\t// passthrough post fail\n\terr = g.AddPassthroughNode(\"1\", WithStatePostHandler(func(ctx context.Context, in string, state string) (string, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.ErrorContains(t, err, \"passthrough node[1]'s post handler type isn't any\")\n\tg.buildError = nil\n\t// passthrough post fail with input key\n\terr = g.AddPassthroughNode(\"1\", WithStatePostHandler(func(ctx context.Context, in string, state string) (string, error) {\n\t\treturn \"\", nil\n\t}), WithInputKey(\"input\"))\n\tassert.ErrorContains(t, err, \"passthrough node[1]'s post handler type isn't any\")\n\tg.buildError = nil\n\t// passthrough pre success\n\terr = g.AddPassthroughNode(\"1\", WithStatePreHandler(func(ctx context.Context, in any, state string) (any, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.NoError(t, err)\n\t// passthrough pre success with input key\n\terr = g.AddPassthroughNode(\"2\", WithStatePreHandler(func(ctx context.Context, in map[string]any, state string) (map[string]any, error) {\n\t\treturn nil, nil\n\t}), WithInputKey(\"input\"))\n\tassert.NoError(t, err)\n\t// passthrough post success\n\terr = g.AddPassthroughNode(\"3\", WithStatePostHandler(func(ctx context.Context, in any, state string) (any, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.NoError(t, err)\n\t// passthrough post success with output key\n\terr = g.AddPassthroughNode(\"4\", WithStatePostHandler(func(ctx context.Context, in map[string]any, state string) (map[string]any, error) {\n\t\treturn nil, nil\n\t}), WithOutputKey(\"output\"))\n\tassert.NoError(t, err)\n\t// common node pre fail\n\terr = g.AddLambdaNode(\"5\", InvokableLambda(func(ctx context.Context, input int) (output int, err error) {\n\t\treturn 0, nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, state string) (string, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.ErrorContains(t, err, \"node[5]'s pre handler type[string] is different from its input type[int]\")\n\tg.buildError = nil\n\t// common node post fail\n\terr = g.AddLambdaNode(\"5\", InvokableLambda(func(ctx context.Context, input int) (output int, err error) {\n\t\treturn 0, nil\n\t}), WithStatePostHandler(func(ctx context.Context, in string, state string) (string, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.ErrorContains(t, err, \"node[5]'s post handler type[string] is different from its output type[int]\")\n\tg.buildError = nil\n\t// common node pre success\n\terr = g.AddLambdaNode(\"5\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, state string) (string, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.NoError(t, err)\n\t// common node post success\n\terr = g.AddLambdaNode(\"6\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"\", nil\n\t}), WithStatePostHandler(func(ctx context.Context, in string, state string) (string, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.NoError(t, err)\n\t// pre state fail\n\terr = g.AddLambdaNode(\"7\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in string, state int) (string, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.ErrorContains(t, err, \"node[7]'s pre handler state type[int] is different from graph[string]\")\n\tg.buildError = nil\n\t// post state fail\n\terr = g.AddLambdaNode(\"7\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"\", nil\n\t}), WithStatePostHandler(func(ctx context.Context, in string, state int) (string, error) {\n\t\treturn \"\", nil\n\t}))\n\tassert.ErrorContains(t, err, \"node[7]'s post handler state type[int] is different from graph[string]\")\n\tg.buildError = nil\n\t// common pre success with input key\n\terr = g.AddLambdaNode(\"7\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"\", nil\n\t}), WithStatePreHandler(func(ctx context.Context, in map[string]any, state string) (map[string]any, error) {\n\t\treturn nil, nil\n\t}), WithInputKey(\"input\"))\n\tassert.NoError(t, err)\n\t// common post success with output key\n\terr = g.AddLambdaNode(\"8\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn \"\", nil\n\t}), WithStatePostHandler(func(ctx context.Context, in map[string]any, state string) (map[string]any, error) {\n\t\treturn nil, nil\n\t}), WithOutputKey(\"output\"))\n\tassert.NoError(t, err)\n}\n\nfunc TestSetFanInMergeConfig_RealStreamNode(t *testing.T) {\n\tfor _, triggerMode := range []NodeTriggerMode{AnyPredecessor, AllPredecessor} {\n\t\tt.Run(string(triggerMode), func(t *testing.T) {\n\t\t\tg := NewGraph[int, map[string]any]()\n\n\t\t\t// Add two stream nodes that output streams of int slices\n\t\t\terr := g.AddLambdaNode(\"s1\", StreamableLambda(func(ctx context.Context, input int) (*schema.StreamReader[int], error) {\n\t\t\t\tsr, sw := schema.Pipe[int](2)\n\t\t\t\tsw.Send(input+1, nil)\n\t\t\t\tsw.Send(input+2, nil)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t}), WithOutputKey(\"s1\"))\n\t\t\tassert.NoError(t, err)\n\t\t\terr = g.AddLambdaNode(\"s2\", StreamableLambda(func(ctx context.Context, input int) (*schema.StreamReader[int], error) {\n\t\t\t\tsr, sw := schema.Pipe[int](2)\n\t\t\t\tsw.Send(input+10, nil)\n\t\t\t\tsw.Send(input+20, nil)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t}), WithOutputKey(\"s2\"))\n\t\t\tassert.NoError(t, err)\n\n\t\t\t// Connect edges: START -> s1, START -> s2, s1 -> END, s2 -> END\n\t\t\terr = g.AddEdge(START, \"s1\")\n\t\t\tassert.NoError(t, err)\n\t\t\terr = g.AddEdge(START, \"s2\")\n\t\t\tassert.NoError(t, err)\n\t\t\terr = g.AddEdge(\"s1\", END)\n\t\t\tassert.NoError(t, err)\n\t\t\terr = g.AddEdge(\"s2\", END)\n\t\t\tassert.NoError(t, err)\n\n\t\t\tr, err := g.Compile(context.Background(), WithNodeTriggerMode(triggerMode),\n\t\t\t\tWithFanInMergeConfig(map[string]FanInMergeConfig{END: {StreamMergeWithSourceEOF: true}}))\n\t\t\tassert.NoError(t, err)\n\n\t\t\t// Run the graph in stream mode and check for SourceEOF events\n\t\t\tsr, err := r.Stream(context.Background(), 1)\n\t\t\tassert.NoError(t, err)\n\n\t\t\tmerged := make(map[string]map[int]bool)\n\t\t\tvar sourceEOFCount int\n\t\t\tsourceNames := make(map[string]bool)\n\t\t\tfor {\n\t\t\t\tm, e := sr.Recv()\n\t\t\t\tif e != nil {\n\t\t\t\t\tif name, ok := schema.GetSourceName(e); ok {\n\t\t\t\t\t\tsourceEOFCount++\n\t\t\t\t\t\tsourceNames[name] = true\n\t\t\t\t\t\tcontinue\n\t\t\t\t\t}\n\t\t\t\t\tif e == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\tassert.NoError(t, e)\n\t\t\t\t}\n\n\t\t\t\tfor k, v := range m {\n\t\t\t\t\tif merged[k] == nil {\n\t\t\t\t\t\tmerged[k] = make(map[int]bool)\n\t\t\t\t\t}\n\n\t\t\t\t\tmerged[k][v.(int)] = true\n\t\t\t\t}\n\t\t\t}\n\n\t\t\t// The merged map should contain both results\n\t\t\tassert.Equal(t, map[string]map[int]bool{\"s1\": {2: true, 3: true}, \"s2\": {11: true, 21: true}}, merged)\n\t\t\tassert.Equal(t, 2, sourceEOFCount, \"should receive SourceEOF for each input stream when StreamMergeWithSourceEOF is true\")\n\t\t\tassert.True(t, sourceNames[\"s1\"], \"should receive SourceEOF from s1\")\n\t\t\tassert.True(t, sourceNames[\"s2\"], \"should receive SourceEOF from s2\")\n\t\t})\n\t}\n}\n\nfunc TestFindLoops(t *testing.T) {\n\ttests := []struct {\n\t\tname       string\n\t\tstartNodes []string\n\t\tchanCalls  map[string]*chanCall\n\t\texpected   [][]string\n\t}{\n\t\t{\n\t\t\tname:       \"Graph without cycles\",\n\t\t\tstartNodes: []string{\"A\"},\n\t\t\tchanCalls: map[string]*chanCall{\n\t\t\t\t\"A\": {\n\t\t\t\t\tcontrols: []string{\"B\", \"C\"},\n\t\t\t\t},\n\t\t\t\t\"B\": {\n\t\t\t\t\tcontrols: []string{\"D\"},\n\t\t\t\t},\n\t\t\t\t\"C\": {\n\t\t\t\t\tcontrols: []string{\"E\"},\n\t\t\t\t},\n\t\t\t\t\"D\": {\n\t\t\t\t\tcontrols: []string{},\n\t\t\t\t},\n\t\t\t\t\"E\": {\n\t\t\t\t\tcontrols: []string{},\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: [][]string{},\n\t\t},\n\t\t{\n\t\t\tname:       \"Graph with self-loop\",\n\t\t\tstartNodes: []string{\"A\"},\n\t\t\tchanCalls: map[string]*chanCall{\n\t\t\t\t\"A\": {\n\t\t\t\t\tcontrols: []string{\"A\", \"B\"},\n\t\t\t\t},\n\t\t\t\t\"B\": {\n\t\t\t\t\tcontrols: []string{},\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: [][]string{{\"A\", \"A\"}},\n\t\t},\n\t\t{\n\t\t\tname:       \"Graph with simple cycle\",\n\t\t\tstartNodes: []string{\"A\", \"B\", \"C\"},\n\t\t\tchanCalls: map[string]*chanCall{\n\t\t\t\t\"A\": {\n\t\t\t\t\tcontrols: []string{\"B\"},\n\t\t\t\t},\n\t\t\t\t\"B\": {\n\t\t\t\t\tcontrols: []string{\"C\"},\n\t\t\t\t},\n\t\t\t\t\"C\": {\n\t\t\t\t\tcontrols: []string{\"A\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: [][]string{{\"A\", \"B\", \"C\", \"A\"}},\n\t\t},\n\t\t{\n\t\t\tname:       \"Graph with multiple cycles\",\n\t\t\tstartNodes: []string{\"A\", \"B\", \"C\", \"D\", \"E\", \"F\"},\n\t\t\tchanCalls: map[string]*chanCall{\n\t\t\t\t\"A\": {\n\t\t\t\t\tcontrols: []string{\"B\", \"D\"},\n\t\t\t\t},\n\t\t\t\t\"B\": {\n\t\t\t\t\tcontrols: []string{\"C\"},\n\t\t\t\t},\n\t\t\t\t\"C\": {\n\t\t\t\t\tcontrols: []string{\"B\"},\n\t\t\t\t},\n\t\t\t\t\"D\": {\n\t\t\t\t\tcontrols: []string{\"E\"},\n\t\t\t\t},\n\t\t\t\t\"E\": {\n\t\t\t\t\tcontrols: []string{\"F\"},\n\t\t\t\t},\n\t\t\t\t\"F\": {\n\t\t\t\t\tcontrols: []string{\"D\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: [][]string{{\"B\", \"C\", \"B\"}, {\"D\", \"E\", \"F\", \"D\"}},\n\t\t},\n\t\t{\n\t\t\tname:       \"Graph with branch cycle\",\n\t\t\tstartNodes: []string{\"A\", \"C\"},\n\t\t\tchanCalls: map[string]*chanCall{\n\t\t\t\t\"A\": {\n\t\t\t\t\tcontrols: []string{\"B\"},\n\t\t\t\t\twriteToBranches: []*GraphBranch{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tendNodes: map[string]bool{\n\t\t\t\t\t\t\t\t\"C\": true,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t\"B\": {\n\t\t\t\t\tcontrols: []string{},\n\t\t\t\t},\n\t\t\t\t\"C\": {\n\t\t\t\t\tcontrols: []string{\"A\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: [][]string{{\"A\", \"C\", \"A\"}},\n\t\t},\n\t\t{\n\t\t\tname:       \"Empty graph\",\n\t\t\tstartNodes: []string{},\n\t\t\tchanCalls:  map[string]*chanCall{},\n\t\t\texpected:   [][]string{},\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tloops := findLoops(tt.startNodes, tt.chanCalls)\n\n\t\t\tassert.Equal(t, len(tt.expected), len(loops))\n\n\t\t\tif len(tt.expected) > 0 {\n\t\t\t\tnormalizedExpected := normalizeLoops(tt.expected)\n\t\t\t\tnormalizedActual := normalizeLoops(loops)\n\t\t\t\tassert.Equal(t, normalizedExpected, normalizedActual)\n\t\t\t}\n\t\t})\n\t}\n}\nfunc normalizeLoops(loops [][]string) []string {\n\tresult := make([]string, 0, len(loops))\n\n\tfor _, loop := range loops {\n\t\tif len(loop) == 0 {\n\t\t\tcontinue\n\t\t}\n\n\t\tnormalizedLoop := make([]string, len(loop))\n\t\tcopy(normalizedLoop, loop)\n\t\tif normalizedLoop[0] != normalizedLoop[len(normalizedLoop)-1] {\n\t\t\tnormalizedLoop = append(normalizedLoop, normalizedLoop[0])\n\t\t}\n\n\t\tminIdx := 0\n\t\tfor i := 1; i < len(normalizedLoop)-1; i++ {\n\t\t\tif normalizedLoop[i] < normalizedLoop[minIdx] {\n\t\t\t\tminIdx = i\n\t\t\t}\n\t\t}\n\n\t\tcanonicalLoop := \"\"\n\t\tfor i := 0; i < len(normalizedLoop)-1; i++ {\n\t\t\tidx := (minIdx + i) % (len(normalizedLoop) - 1)\n\t\t\tcanonicalLoop += normalizedLoop[idx] + \",\"\n\t\t}\n\t\tcanonicalLoop += normalizedLoop[minIdx]\n\n\t\tresult = append(result, canonicalLoop)\n\t}\n\n\tsort.Strings(result)\n\treturn result\n}\n\nfunc TestPrintTasks(t *testing.T) {\n\tvar ts []*task\n\tassert.Equal(t, \"[]\", printTask(ts))\n\tts = []*task{{nodeKey: \"1\"}}\n\tassert.Equal(t, \"[1]\", printTask(ts))\n\tts = []*task{{nodeKey: \"1\"}, {nodeKey: \"2\"}, {nodeKey: \"3\"}}\n\tassert.Equal(t, \"[1, 2, 3]\", printTask(ts))\n}\n\nfunc TestSkipBranch(t *testing.T) {\n\tg := NewGraph[string, string]()\n\t_ = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}))\n\t_ = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}))\n\t_ = g.AddEdge(START, \"1\")\n\t_ = g.AddBranch(\"1\", NewGraphMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) {\n\t\treturn map[string]bool{}, nil\n\t}, map[string]bool{\"2\": true}))\n\t_ = g.AddEdge(\"2\", END)\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor))\n\tassert.NoError(t, err)\n\t_, err = r.Invoke(ctx, \"input\")\n\tassert.ErrorContains(t, err, \"[GraphRunError] no tasks to execute, last completed nodes: [1]\")\n\n\tg = NewGraph[string, string]()\n\t_ = g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}))\n\t_ = g.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t}))\n\t_ = g.AddEdge(START, \"1\")\n\t_ = g.AddBranch(\"1\", NewGraphMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) {\n\t\treturn map[string]bool{}, nil\n\t}, map[string]bool{\"2\": true}))\n\t_ = g.AddEdge(\"2\", END)\n\t_ = g.AddEdge(START, \"2\")\n\tr, err = g.Compile(ctx, WithNodeTriggerMode(AllPredecessor))\n\tassert.NoError(t, err)\n\tresult, err := r.Invoke(ctx, \"input\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"input\", result)\n}\n\nfunc TestGetStateInGraphCallback(t *testing.T) {\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (s *state) {\n\t\treturn &state{}\n\t}))\n\tassert.NoError(t, g.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t})))\n\tassert.NoError(t, g.AddEdge(START, \"1\"))\n\tassert.NoError(t, g.AddEdge(\"1\", END))\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx)\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(ctx, \"input\", WithCallbacks(&testGraphStateCallbackHandler{t: t}))\n\tassert.NoError(t, err)\n}\n\ntype state struct {\n\tA string\n}\n\ntype testGraphStateCallbackHandler struct {\n\tt *testing.T\n}\n\nfunc (t *testGraphStateCallbackHandler) OnStart(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\tassert.NoError(t.t, ProcessState[*state](ctx, func(ctx context.Context, s *state) error {\n\t\ts.A = \"test\"\n\t\treturn nil\n\t}))\n\treturn ctx\n}\n\nfunc (t *testGraphStateCallbackHandler) OnEnd(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\treturn ctx\n}\n\nfunc (t *testGraphStateCallbackHandler) OnError(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\treturn ctx\n}\n\nfunc (t *testGraphStateCallbackHandler) OnStartWithStreamInput(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {\n\treturn ctx\n}\n\nfunc (t *testGraphStateCallbackHandler) OnEndWithStreamOutput(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context {\n\treturn ctx\n}\n\nfunc TestUniqueSlice(t *testing.T) {\n\tassert.Equal(t, []string{\"a\", \"b\", \"c\"}, uniqueSlice([]string{\"a\", \"b\", \"a\", \"c\", \"b\"}))\n\tassert.Equal(t, []string{}, uniqueSlice([]string{}))\n}\n"
  },
  {
    "path": "compose/interrupt.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\n\t\"github.com/google/uuid\"\n\n\t\"github.com/cloudwego/eino/internal/core\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// WithInterruptBeforeNodes instructs to interrupt before the given nodes.\nfunc WithInterruptBeforeNodes(nodes []string) GraphCompileOption {\n\treturn func(options *graphCompileOptions) {\n\t\toptions.interruptBeforeNodes = nodes\n\t}\n}\n\n// WithInterruptAfterNodes instructs to interrupt after the given nodes.\nfunc WithInterruptAfterNodes(nodes []string) GraphCompileOption {\n\treturn func(options *graphCompileOptions) {\n\t\toptions.interruptAfterNodes = nodes\n\t}\n}\n\n// Deprecated: prefer Interrupt/StatefulInterrupt and CompositeInterrupt.\n// If you need to pass the legacy error into CompositeInterrupt, wrap it using WrapInterruptAndRerunIfNeeded first.\nvar InterruptAndRerun = deprecatedInterruptAndRerun\nvar deprecatedInterruptAndRerun = errors.New(\"interrupt and rerun\")\n\n// NewInterruptAndRerunErr creates a legacy interrupt-and-rerun error.\n// Deprecated: prefer Interrupt(ctx, info) or StatefulInterrupt(ctx, info, state).\n// If passing into CompositeInterrupt, wrap using WrapInterruptAndRerunIfNeeded first.\nfunc NewInterruptAndRerunErr(extra any) error {\n\treturn deprecatedInterruptAndRerunErr(extra)\n}\nfunc deprecatedInterruptAndRerunErr(extra any) error {\n\treturn &core.InterruptSignal{InterruptInfo: core.InterruptInfo{\n\t\tInfo:        extra,\n\t\tIsRootCause: true,\n\t}}\n}\n\ntype wrappedInterruptAndRerun struct {\n\tps    Address\n\tinner error\n}\n\nfunc (w *wrappedInterruptAndRerun) Error() string {\n\treturn fmt.Sprintf(\"interrupt and rerun at address %s: %s\", w.ps.String(), w.inner.Error())\n}\n\nfunc (w *wrappedInterruptAndRerun) Unwrap() error {\n\treturn w.inner\n}\n\n// WrapInterruptAndRerunIfNeeded wraps the deprecated old interrupt errors, with the current execution address.\n// If the error is returned by either Interrupt, StatefulInterrupt or CompositeInterrupt,\n// it will be returned as-is without wrapping\nfunc WrapInterruptAndRerunIfNeeded(ctx context.Context, step AddressSegment, err error) error {\n\taddr := GetCurrentAddress(ctx)\n\tnewAddr := append(append([]AddressSegment{}, addr...), step)\n\tif errors.Is(err, deprecatedInterruptAndRerun) {\n\t\treturn &wrappedInterruptAndRerun{\n\t\t\tps:    newAddr,\n\t\t\tinner: err,\n\t\t}\n\t}\n\n\tire := &core.InterruptSignal{}\n\tif errors.As(err, &ire) {\n\t\tif ire.Address == nil {\n\t\t\treturn &wrappedInterruptAndRerun{\n\t\t\t\tps:    newAddr,\n\t\t\t\tinner: err,\n\t\t\t}\n\t\t}\n\t\treturn ire\n\t}\n\n\treturn fmt.Errorf(\"failed to wrap error as addressed InterruptAndRerun: %w\", err)\n}\n\n// Interrupt creates a special error that signals the execution engine to interrupt\n// the current run at the component's specific address and save a checkpoint.\n//\n// This is the standard way for a single, non-composite component to signal a resumable interruption.\n//\n//   - ctx: The context of the running component, used to retrieve the current execution address.\n//   - info: User-facing information about the interrupt. This is not persisted but is exposed to the\n//     calling application via the InterruptCtx to provide context (e.g., a reason for the pause).\nfunc Interrupt(ctx context.Context, info any) error {\n\tis, err := core.Interrupt(ctx, info, nil, nil)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treturn is\n}\n\n// StatefulInterrupt creates a special error that signals the execution engine to interrupt\n// the current run at the component's specific address and save a checkpoint.\n//\n// This is the standard way for a single, non-composite component to signal a resumable interruption.\n//\n//   - ctx: The context of the running component, used to retrieve the current execution address.\n//   - info: User-facing information about the interrupt. This is not persisted but is exposed to the\n//     calling application via the InterruptCtx to provide context (e.g., a reason for the pause).\n//   - state: The internal state that the interrupting component needs to persist to be able to resume\n//     its work later. This state is saved in the checkpoint and will be provided back to the component\n//     upon resumption via GetInterruptState.\nfunc StatefulInterrupt(ctx context.Context, info any, state any) error {\n\tis, err := core.Interrupt(ctx, info, state, nil)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treturn is\n}\n\n// CompositeInterrupt creates a special error that signals a composite interruption.\n// It is designed for \"composite\" nodes (like ToolsNode) that manage multiple, independent,\n// interruptible sub-processes. It bundles multiple sub-interrupt errors into a single error\n// that the engine can deconstruct into a flat list of resumable points.\n//\n// This function is robust and can handle several types of errors from sub-processes:\n//\n//   - A `Interrupt` or `StatefulInterrupt` error from a simple component.\n//\n//   - A nested `CompositeInterrupt` error from another composite component.\n//\n//   - An error containing `InterruptInfo` returned by a `Runnable` (e.g., a Graph within a lambda node).\n//\n//   - An error returned by \\'WrapInterruptAndRerunIfNeeded\\' for the legacy old interrupt and rerun error,\n//     and for the error returned by the deprecated old interrupt errors.\n//\n// Parameters:\n//\n//   - ctx: The context of the running composite node.\n//\n//   - info: User-facing information for the composite node itself. Can be nil.\n//     This info will be attached to InterruptInfo.RerunNodeExtra.\n//     Provided mainly for compatibility purpose as the composite node itself\n//     is not an interrupt point with interrupt ID,\n//     which means it lacks enough reason to give a user-facing info.\n//\n//   - state: The state for the composite node itself. Can be nil.\n//     This could be useful when the composite node needs to restore state,\n//     such as its input (e.g. ToolsNode).\n//\n//   - errs: a list of errors emitted by sub-processes.\n//\n// NOTE: if the error you passed in is the deprecated old interrupt and rerun err, or an error returned by\n// the deprecated old interrupt function, you must wrap it using WrapInterruptAndRerunIfNeeded first\n// before passing them into this function.\nfunc CompositeInterrupt(ctx context.Context, info any, state any, errs ...error) error {\n\tif len(errs) == 0 {\n\t\treturn StatefulInterrupt(ctx, info, state)\n\t}\n\n\tvar cErrs []*core.InterruptSignal\n\tfor _, err := range errs {\n\t\twrapped := &wrappedInterruptAndRerun{}\n\t\tif errors.As(err, &wrapped) {\n\t\t\tinner := wrapped.Unwrap()\n\t\t\tif errors.Is(inner, deprecatedInterruptAndRerun) {\n\t\t\t\tid := uuid.NewString()\n\t\t\t\tcErrs = append(cErrs, &core.InterruptSignal{\n\t\t\t\t\tID:      id,\n\t\t\t\t\tAddress: wrapped.ps,\n\t\t\t\t\tInterruptInfo: core.InterruptInfo{\n\t\t\t\t\t\tInfo:        nil,\n\t\t\t\t\t\tIsRootCause: true,\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tire := &core.InterruptSignal{}\n\t\t\tif errors.As(err, &ire) {\n\t\t\t\tid := uuid.NewString()\n\t\t\t\tcErrs = append(cErrs, &core.InterruptSignal{\n\t\t\t\t\tID:      id,\n\t\t\t\t\tAddress: wrapped.ps,\n\t\t\t\t\tInterruptInfo: core.InterruptInfo{\n\t\t\t\t\t\tInfo:        ire.InterruptInfo.Info,\n\t\t\t\t\t\tIsRootCause: ire.InterruptInfo.IsRootCause,\n\t\t\t\t\t},\n\t\t\t\t\tInterruptState: core.InterruptState{\n\t\t\t\t\t\tState: ire.InterruptState.State,\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t}\n\n\t\t\tcontinue\n\t\t}\n\n\t\tire := &core.InterruptSignal{}\n\t\tif errors.As(err, &ire) {\n\t\t\tcErrs = append(cErrs, ire)\n\t\t\tcontinue\n\t\t}\n\n\t\tie := &interruptError{}\n\t\tif errors.As(err, &ie) {\n\t\t\tis := core.FromInterruptContexts(ie.Info.InterruptContexts)\n\t\t\tcErrs = append(cErrs, is)\n\t\t\tcontinue\n\t\t}\n\n\t\treturn fmt.Errorf(\"composite interrupt but one of the sub error is not interrupt and rerun error: %w\", err)\n\t}\n\n\tis, err := core.Interrupt(ctx, info, state, cErrs)\n\tif err != nil {\n\t\treturn err\n\t}\n\treturn is\n}\n\n// IsInterruptRerunError reports whether the error represents an interrupt-and-rerun\n// and returns any attached info.\nfunc IsInterruptRerunError(err error) (any, bool) {\n\tinfo, _, ok := isInterruptRerunError(err)\n\treturn info, ok\n}\n\nfunc isInterruptRerunError(err error) (info any, state any, ok bool) {\n\tif errors.Is(err, deprecatedInterruptAndRerun) {\n\t\treturn nil, nil, true\n\t}\n\tire := &core.InterruptSignal{}\n\tif errors.As(err, &ire) {\n\t\treturn ire.Info, ire.State, true\n\t}\n\treturn nil, nil, false\n}\n\n// InterruptInfo aggregates interrupt metadata for composite or nested runs.\ntype InterruptInfo struct {\n\tState             any\n\tBeforeNodes       []string\n\tAfterNodes        []string\n\tRerunNodes        []string\n\tRerunNodesExtra   map[string]any\n\tSubGraphs         map[string]*InterruptInfo\n\tInterruptContexts []*InterruptCtx\n}\n\nfunc init() {\n\tschema.RegisterName[*InterruptInfo](\"_eino_compose_interrupt_info\")\n}\n\n// AddressSegmentType defines the type of a segment in an execution address.\ntype AddressSegmentType = core.AddressSegmentType\n\nconst (\n\t// AddressSegmentNode represents a segment of an address that corresponds to a graph node.\n\tAddressSegmentNode AddressSegmentType = \"node\"\n\t// AddressSegmentTool represents a segment of an address that corresponds to a specific tool call within a ToolsNode.\n\tAddressSegmentTool AddressSegmentType = \"tool\"\n\t// AddressSegmentRunnable represents a segment of an address that corresponds to an instance of the Runnable interface.\n\t// Currently the possible Runnable types are: Graph, Workflow and Chain.\n\t// Note that for sub-graphs added through AddGraphNode to another graph is not a Runnable.\n\t// So a AddressSegmentRunnable indicates a standalone Root level Graph,\n\t// or a Root level Graph inside a node such as Lambda node.\n\tAddressSegmentRunnable AddressSegmentType = \"runnable\"\n)\n\n// Address represents a full, hierarchical address to a point in the execution structure.\ntype Address = core.Address\n\n// AddressSegment represents a single segment in the hierarchical address of an execution point.\n// A sequence of AddressSegments uniquely identifies a location within a potentially nested structure.\ntype AddressSegment = core.AddressSegment\n\n// InterruptCtx provides a complete, user-facing context for a single, resumable interrupt point.\ntype InterruptCtx = core.InterruptCtx\n\n// ExtractInterruptInfo extracts InterruptInfo from an error if present.\nfunc ExtractInterruptInfo(err error) (info *InterruptInfo, existed bool) {\n\tif err == nil {\n\t\treturn nil, false\n\t}\n\tvar iE *interruptError\n\tif errors.As(err, &iE) {\n\t\treturn iE.Info, true\n\t}\n\tvar sIE *subGraphInterruptError\n\tif errors.As(err, &sIE) {\n\t\treturn sIE.Info, true\n\t}\n\treturn nil, false\n}\n\ntype interruptError struct {\n\tInfo *InterruptInfo\n}\n\nfunc (e *interruptError) Error() string {\n\treturn fmt.Sprintf(\"interrupt happened, info: %+v\", e.Info)\n}\n\nfunc (e *interruptError) GetInterruptContexts() []*InterruptCtx {\n\tif e.Info == nil {\n\t\treturn nil\n\t}\n\treturn e.Info.InterruptContexts\n}\n\nfunc isSubGraphInterrupt(err error) *subGraphInterruptError {\n\tif err == nil {\n\t\treturn nil\n\t}\n\tvar iE *subGraphInterruptError\n\tif errors.As(err, &iE) {\n\t\treturn iE\n\t}\n\treturn nil\n}\n\ntype subGraphInterruptError struct {\n\tInfo       *InterruptInfo\n\tCheckPoint *checkpoint\n\n\tsignal *core.InterruptSignal\n}\n\nfunc (e *subGraphInterruptError) Error() string {\n\treturn fmt.Sprintf(\"interrupt happened, info: %+v\", e.Info)\n}\n\nfunc isInterruptError(err error) bool {\n\tif _, ok := ExtractInterruptInfo(err); ok {\n\t\treturn true\n\t}\n\tif info := isSubGraphInterrupt(err); info != nil {\n\t\treturn true\n\t}\n\tif _, ok := IsInterruptRerunError(err); ok {\n\t\treturn true\n\t}\n\n\treturn false\n}\n"
  },
  {
    "path": "compose/introspect.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/components\"\n)\n\n// GraphNodeInfo the info which end users pass in when they are adding nodes to graph.\ntype GraphNodeInfo struct {\n\tComponent             components.Component\n\tInstance              any\n\tGraphAddNodeOpts      []GraphAddNodeOpt\n\tInputType, OutputType reflect.Type // mainly for lambda, whose input and output types cannot be inferred by component type\n\tName                  string\n\tInputKey, OutputKey   string\n\tGraphInfo             *GraphInfo\n\tMappings              []*FieldMapping\n}\n\n// GraphInfo the info which end users pass in when they are compiling a graph.\n// it is used in compile callback for user to get the node info and instance.\n// you may need all details info of the graph for observation.\ntype GraphInfo struct {\n\tCompileOptions        []GraphCompileOption\n\tNodes                 map[string]GraphNodeInfo // node key -> node info\n\tEdges                 map[string][]string      // edge start node key -> edge end node key, control edges\n\tDataEdges             map[string][]string\n\tBranches              map[string][]GraphBranch // branch start node key -> branch\n\tInputType, OutputType reflect.Type\n\tName                  string\n\n\tNewGraphOptions []NewGraphOption\n\tGenStateFn      func(context.Context) any\n}\n\n// GraphCompileCallback is the callback which will be called when graph compilation finishes.\ntype GraphCompileCallback interface {\n\tOnFinish(ctx context.Context, info *GraphInfo)\n}\n"
  },
  {
    "path": "compose/pregel.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport \"fmt\"\n\nfunc pregelChannelBuilder(_ []string, _ []string, _ func() any, _ func() streamReader) channel {\n\treturn &pregelChannel{Values: make(map[string]any)}\n}\n\ntype pregelChannel struct {\n\tValues map[string]any\n\n\tmergeConfig FanInMergeConfig\n}\n\nfunc (ch *pregelChannel) setMergeConfig(cfg FanInMergeConfig) {\n\tch.mergeConfig.StreamMergeWithSourceEOF = cfg.StreamMergeWithSourceEOF\n}\n\nfunc (ch *pregelChannel) load(c channel) error {\n\tdc, ok := c.(*pregelChannel)\n\tif !ok {\n\t\treturn fmt.Errorf(\"load pregel channel fail, got %T, want *pregelChannel\", c)\n\t}\n\tch.Values = dc.Values\n\treturn nil\n}\n\nfunc (ch *pregelChannel) convertValues(fn func(map[string]any) error) error {\n\treturn fn(ch.Values)\n}\n\nfunc (ch *pregelChannel) reportValues(ins map[string]any) error {\n\tfor k, v := range ins {\n\t\tch.Values[k] = v\n\t}\n\treturn nil\n}\n\nfunc (ch *pregelChannel) get(isStream bool, name string, edgeHandler *edgeHandlerManager) (\n\tany, bool, error) {\n\tif len(ch.Values) == 0 {\n\t\treturn nil, false, nil\n\t}\n\tdefer func() { ch.Values = map[string]any{} }()\n\tvalues := make([]any, len(ch.Values))\n\tnames := make([]string, len(ch.Values))\n\ti := 0\n\tfor k, v := range ch.Values {\n\t\tresolvedV, err := edgeHandler.handle(k, name, v, isStream)\n\t\tif err != nil {\n\t\t\treturn nil, false, err\n\t\t}\n\t\tvalues[i] = resolvedV\n\t\tnames[i] = k\n\t\ti++\n\t}\n\n\tif len(values) == 1 {\n\t\treturn values[0], true, nil\n\t}\n\n\t// merge\n\tmergeOpts := &mergeOptions{\n\t\tstreamMergeWithSourceEOF: ch.mergeConfig.StreamMergeWithSourceEOF,\n\t\tnames:                    names,\n\t}\n\tv, err := mergeValues(values, mergeOpts)\n\tif err != nil {\n\t\treturn nil, false, err\n\t}\n\treturn v, true, nil\n}\n\nfunc (ch *pregelChannel) reportSkip(_ []string) bool {\n\treturn false\n}\nfunc (ch *pregelChannel) reportDependencies(_ []string) {\n\treturn\n}\n"
  },
  {
    "path": "compose/resume.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/internal/core\"\n)\n\n// GetInterruptState provides a type-safe way to check for and retrieve the persisted state from a previous interruption.\n// It is the primary function a component should use to understand its past state.\n//\n// It returns three values:\n//   - wasInterrupted (bool): True if the node was part of a previous interruption, regardless of whether state was provided.\n//   - state (T): The typed state object, if it was provided and matches type `T`.\n//   - hasState (bool): True if state was provided during the original interrupt and successfully cast to type `T`.\nfunc GetInterruptState[T any](ctx context.Context) (wasInterrupted bool, hasState bool, state T) {\n\treturn core.GetInterruptState[T](ctx)\n}\n\n// GetResumeContext checks if the current component is the target of a resume operation\n// and retrieves any data provided by the user for that resumption.\n//\n// This function is typically called *after* a component has already determined it is in a\n// resumed state by calling GetInterruptState.\n//\n// It returns three values:\n//   - isResumeFlow: A boolean that is true if the current component's address was explicitly targeted\n//     by a call to Resume() or ResumeWithData().\n//   - hasData: A boolean that is true if data was provided for this component (i.e., not nil).\n//   - data: The typed data provided by the user.\n//\n// ### How to Use This Function: A Decision Framework\n//\n// The correct usage pattern depends on the application's desired resume strategy.\n//\n// #### Strategy 1: Implicit \"Resume All\"\n// In some use cases, any resume operation implies that *all* interrupted points should proceed.\n// For example, if an application's UI only provides a single \"Continue\" button for a set of\n// interruptions. In this model, a component can often just use `GetInterruptState` to see if\n// `wasInterrupted` is true and then proceed with its logic, as it can assume it is an intended target.\n// It may still call `GetResumeContext` to check for optional data, but the `isResumeFlow` flag is less critical.\n//\n// #### Strategy 2: Explicit \"Targeted Resume\" (Most Common)\n// For applications with multiple, distinct interrupt points that must be resumed independently, it is\n// crucial to differentiate which point is being resumed. This is the primary use case for the `isResumeFlow` flag.\n//   - If `isResumeFlow` is `true`: Your component is the explicit target. You should consume\n//     the `data` (if any) and complete your work.\n//   - If `isResumeFlow` is `false`: Another component is the target. You MUST re-interrupt\n//     (e.g., by returning `StatefulInterrupt(...)`) to preserve your state and allow the\n//     resume signal to propagate.\n//\n// ### Guidance for Composite Components\n//\n// Composite components (like `Graph` or other `Runnable`s that contain sub-processes) have a dual role:\n//  1. Check for Self-Targeting: A composite component can itself be the target of a resume\n//     operation, for instance, to modify its internal state. It may call `GetResumeContext`\n//     to check for data targeted at its own address.\n//  2. Act as a Conduit: After checking for itself, its primary role is to re-execute its children,\n//     allowing the resume context to flow down to them. It must not consume a resume signal\n//     intended for one of its descendants.\nfunc GetResumeContext[T any](ctx context.Context) (isResumeFlow bool, hasData bool, data T) {\n\treturn core.GetResumeContext[T](ctx)\n}\n\n// GetCurrentAddress returns the hierarchical address of the currently executing component.\n// The address is a sequence of segments, each identifying a structural part of the execution\n// like an agent, a graph node, or a tool call. This can be useful for logging or debugging.\nfunc GetCurrentAddress(ctx context.Context) Address {\n\treturn core.GetCurrentAddress(ctx)\n}\n\n// Resume prepares a context for an \"Explicit Targeted Resume\" operation by targeting one or more\n// components without providing data. It is a convenience wrapper around BatchResumeWithData.\n//\n// This is useful when the act of resuming is itself the signal, and no extra data is needed.\n// The components at the provided addresses (interrupt IDs) will receive `isResumeFlow = true`\n// when they call `GetResumeContext`.\nfunc Resume(ctx context.Context, interruptIDs ...string) context.Context {\n\tresumeData := make(map[string]any, len(interruptIDs))\n\tfor _, addr := range interruptIDs {\n\t\tresumeData[addr] = nil\n\t}\n\treturn BatchResumeWithData(ctx, resumeData)\n}\n\n// ResumeWithData prepares a context to resume a single, specific component with data.\n// It is the primary function for the \"Explicit Targeted Resume\" strategy when data is required.\n// It is a convenience wrapper around BatchResumeWithData.\n// The `interruptID` parameter is the unique interrupt ID of the target component.\nfunc ResumeWithData(ctx context.Context, interruptID string, data any) context.Context {\n\treturn BatchResumeWithData(ctx, map[string]any{interruptID: data})\n}\n\n// BatchResumeWithData is the core function for preparing a resume context. It injects a map\n// of resume targets and their corresponding data into the context.\n//\n// The `resumeData` map should contain the interrupt IDs (which are the string form of addresses) of the\n// components to be resumed as keys. The value can be the resume data for that component, or `nil`\n// if no data is needed (equivalent to using `Resume`).\n//\n// This function is the foundation for the \"Explicit Targeted Resume\" strategy. Components whose interrupt IDs\n// are present as keys in the map will receive `isResumeFlow = true` when they call `GetResumeContext`.\nfunc BatchResumeWithData(ctx context.Context, resumeData map[string]any) context.Context {\n\treturn core.BatchResumeWithData(ctx, resumeData)\n}\n\nfunc getNodePath(ctx context.Context) (*NodePath, bool) {\n\tcurrentAddress := GetCurrentAddress(ctx)\n\tif len(currentAddress) == 0 {\n\t\treturn nil, false\n\t}\n\n\tnodePath := make([]string, 0, len(currentAddress))\n\tfor _, p := range currentAddress {\n\t\tif p.Type == AddressSegmentRunnable {\n\t\t\tnodePath = []string{}\n\t\t\tcontinue\n\t\t}\n\n\t\tnodePath = append(nodePath, p.ID)\n\t}\n\n\treturn NewNodePath(nodePath...), len(nodePath) > 0\n}\n\n// AppendAddressSegment creates a new execution context for a sub-component (e.g., a graph node or a tool call).\n//\n// It extends the current context's address with a new segment and populates the new context with the\n// appropriate interrupt state and resume data for that specific sub-address.\n//\n//   - ctx: The parent context, typically the one passed into the component's Invoke/Stream method.\n//   - segType: The type of the new address segment (e.g., \"node\", \"tool\").\n//   - segID: The unique ID for the new address segment.\nfunc AppendAddressSegment(ctx context.Context, segType AddressSegmentType, segID string) context.Context {\n\treturn core.AppendAddressSegment(ctx, segType, segID, \"\")\n}\n\nfunc appendToolAddressSegment(ctx context.Context, segID string, subID string) context.Context {\n\treturn core.AppendAddressSegment(ctx, AddressSegmentTool, segID, subID)\n}\n"
  },
  {
    "path": "compose/resume_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype myInterruptState struct {\n\tOriginalInput string\n}\n\ntype myResumeData struct {\n\tMessage string\n}\n\ntype resumeTestState struct {\n\tOnStartCalledOnResume bool `json:\"on_start_called_on_resume\"`\n\tCounter               int  `json:\"counter\"`\n}\n\nfunc init() {\n\tschema.Register[resumeTestState]()\n}\n\nfunc TestInterruptStateAndResumeForRootGraph(t *testing.T) {\n\t// create a graph with a lambda node\n\t// this lambda node will interrupt with a typed state and an info for end-user\n\t// verify the info thrown by the lambda node\n\t// resume with a structured resume data\n\t// within the lambda node, getRunCtx and verify the state and resume data\n\tg := NewGraph[string, string]()\n\n\tlambda := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\twasInterrupted, hasState, state := GetInterruptState[*myInterruptState](ctx)\n\t\tif !wasInterrupted {\n\t\t\t// First run: interrupt with state\n\t\t\treturn \"\", StatefulInterrupt(ctx,\n\t\t\t\tmap[string]any{\"reason\": \"scheduled maintenance\"},\n\t\t\t\t&myInterruptState{OriginalInput: input},\n\t\t\t)\n\t\t}\n\n\t\t// This is a resumed run.\n\t\tassert.True(t, hasState)\n\t\tassert.Equal(t, \"initial input\", state.OriginalInput)\n\n\t\tisResume, hasData, data := GetResumeContext[*myResumeData](ctx)\n\t\tassert.True(t, isResume)\n\t\tassert.True(t, hasData)\n\t\tassert.Equal(t, \"let's continue\", data.Message)\n\n\t\treturn \"Resumed successfully with input: \" + state.OriginalInput, nil\n\t})\n\n\t_ = g.AddLambdaNode(\"lambda\", lambda)\n\t_ = g.AddEdge(START, \"lambda\")\n\t_ = g.AddEdge(\"lambda\", END)\n\n\tgraph, err := g.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()), WithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\n\t// First invocation, which should be interrupted\n\tcheckPointID := \"test-checkpoint-1\"\n\t_, err = graph.Invoke(context.Background(), \"initial input\", WithCheckPointID(checkPointID))\n\n\t// Verify the interrupt error and extracted info\n\tassert.Error(t, err)\n\tinterruptInfo, isInterrupt := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt)\n\tassert.NotNil(t, interruptInfo)\n\n\tinterruptContexts := interruptInfo.InterruptContexts\n\tassert.Equal(t, 1, len(interruptContexts))\n\tassert.Equal(t, \"runnable:root;node:lambda\", interruptContexts[0].Address.String())\n\tassert.Equal(t, map[string]any{\"reason\": \"scheduled maintenance\"}, interruptContexts[0].Info)\n\n\t// Prepare resume data\n\tctx := ResumeWithData(context.Background(), interruptContexts[0].ID,\n\t\t&myResumeData{Message: \"let's continue\"})\n\n\t// Resume execution\n\toutput, err := graph.Invoke(ctx, \"\", WithCheckPointID(checkPointID))\n\n\t// Verify the final result\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"Resumed successfully with input: initial input\", output)\n}\n\nfunc TestProcessStateInOnStartDuringResume(t *testing.T) {\n\tgraphOnStartCallCount := 0\n\tprocessStateErrorOnResume := error(nil)\n\n\tcb := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\tif info.Name == \"test-process-state-onstart\" {\n\t\t\t\tgraphOnStartCallCount++\n\t\t\t\terr := ProcessState[*resumeTestState](ctx, func(ctx context.Context, s *resumeTestState) error {\n\t\t\t\t\ts.Counter++\n\t\t\t\t\treturn nil\n\t\t\t\t})\n\t\t\t\tif graphOnStartCallCount > 1 {\n\t\t\t\t\tprocessStateErrorOnResume = err\n\t\t\t\t}\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tBuild()\n\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) *resumeTestState {\n\t\treturn &resumeTestState{}\n\t}))\n\n\tlambda := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\twasInterrupted, _, _ := GetInterruptState[*myInterruptState](ctx)\n\t\tif !wasInterrupted {\n\t\t\treturn \"\", StatefulInterrupt(ctx,\n\t\t\t\tmap[string]any{\"reason\": \"test interrupt\"},\n\t\t\t\t&myInterruptState{OriginalInput: input},\n\t\t\t)\n\t\t}\n\n\t\tvar stateCounter int\n\t\terr := ProcessState[*resumeTestState](ctx, func(ctx context.Context, s *resumeTestState) error {\n\t\t\tstateCounter = s.Counter\n\t\t\treturn nil\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, 2, stateCounter, \"Counter should be 2 (first run OnStart + resume OnStart)\")\n\n\t\treturn \"success\", nil\n\t})\n\n\t_ = g.AddLambdaNode(\"lambda\", lambda)\n\t_ = g.AddEdge(START, \"lambda\")\n\t_ = g.AddEdge(\"lambda\", END)\n\n\tgraph, err := g.Compile(context.Background(),\n\t\tWithCheckPointStore(newInMemoryStore()),\n\t\tWithGraphName(\"test-process-state-onstart\"),\n\t)\n\tassert.NoError(t, err)\n\n\tcheckPointID := \"test-checkpoint-process-state\"\n\t_, err = graph.Invoke(context.Background(), \"test input\", WithCheckPointID(checkPointID), WithCallbacks(cb))\n\n\tassert.Error(t, err, \"First invocation should return an error\")\n\tinterruptInfo, isInterrupt := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt, \"Should be an interrupt error\")\n\tassert.NotNil(t, interruptInfo)\n\tassert.Equal(t, 1, graphOnStartCallCount, \"Graph OnStart should be called once on first run\")\n\n\tctx := ResumeWithData(context.Background(), interruptInfo.InterruptContexts[0].ID, &myResumeData{})\n\n\toutput, err := graph.Invoke(ctx, \"\", WithCheckPointID(checkPointID), WithCallbacks(cb))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"success\", output)\n\tassert.Equal(t, 2, graphOnStartCallCount, \"Graph OnStart should be called twice (first run + resume)\")\n\tassert.NoError(t, processStateErrorOnResume, \"ProcessState should work in OnStart during resume\")\n}\n\nfunc TestInterruptStateAndResumeForSubGraph(t *testing.T) {\n\t// create a graph\n\t// create a another graph with a lambda node, as this graph as a sub-graph of the previous graph\n\t// this lambda node will interrupt with a typed state and an info for end-user\n\t// verify the info thrown by the lambda node\n\t// resume with a structured resume data\n\t// within the lambda node, getRunCtx and verify the state and resume data\n\tsubGraph := NewGraph[string, string]()\n\n\tlambda := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\twasInterrupted, hasState, state := GetInterruptState[*myInterruptState](ctx)\n\t\tif !wasInterrupted {\n\t\t\t// First run: interrupt with state\n\t\t\treturn \"\", StatefulInterrupt(ctx,\n\t\t\t\tmap[string]any{\"reason\": \"sub-graph maintenance\"},\n\t\t\t\t&myInterruptState{OriginalInput: input},\n\t\t\t)\n\t\t}\n\n\t\t// Second (resumed) run\n\t\tassert.True(t, hasState)\n\t\tassert.Equal(t, \"main input\", state.OriginalInput)\n\n\t\tisResume, hasData, data := GetResumeContext[*myResumeData](ctx)\n\t\tassert.True(t, isResume)\n\t\tassert.True(t, hasData)\n\t\tassert.Equal(t, \"let's continue sub-graph\", data.Message)\n\n\t\treturn \"Sub-graph resumed successfully\", nil\n\t})\n\n\t_ = subGraph.AddLambdaNode(\"inner_lambda\", lambda)\n\t_ = subGraph.AddEdge(START, \"inner_lambda\")\n\t_ = subGraph.AddEdge(\"inner_lambda\", END)\n\n\t// Create the main graph\n\tmainGraph := NewGraph[string, string]()\n\t_ = mainGraph.AddGraphNode(\"sub_graph_node\", subGraph)\n\t_ = mainGraph.AddEdge(START, \"sub_graph_node\")\n\t_ = mainGraph.AddEdge(\"sub_graph_node\", END)\n\n\tcompiledMainGraph, err := mainGraph.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t// First invocation, which should be interrupted\n\tcheckPointID := \"test-subgraph-checkpoint-1\"\n\t_, err = compiledMainGraph.Invoke(context.Background(), \"main input\", WithCheckPointID(checkPointID))\n\n\t// Verify the interrupt error and extracted info\n\tassert.Error(t, err)\n\tinterruptInfo, isInterrupt := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt)\n\tassert.NotNil(t, interruptInfo)\n\n\tinterruptContexts := interruptInfo.InterruptContexts\n\tassert.Equal(t, 1, len(interruptContexts))\n\tassert.Equal(t, \"runnable:;node:sub_graph_node;node:inner_lambda\", interruptContexts[0].Address.String())\n\tassert.Equal(t, map[string]any{\"reason\": \"sub-graph maintenance\"}, interruptContexts[0].Info)\n\n\t// Prepare resume data\n\tctx := ResumeWithData(context.Background(), interruptContexts[0].ID,\n\t\t&myResumeData{Message: \"let's continue sub-graph\"})\n\n\t// Resume execution\n\toutput, err := compiledMainGraph.Invoke(ctx, \"\", WithCheckPointID(checkPointID))\n\n\t// Verify the final result\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"Sub-graph resumed successfully\", output)\n}\n\nfunc TestInterruptStateAndResumeForToolInNestedSubGraph(t *testing.T) {\n\t// create a ROOT graph.\n\t// create a sub graph A, add A to ROOT graph using AddGraphNode.\n\t// create a sub-sub graph B, add B to A using AddGraphNode.\n\t// within sub-sub graph B, add a ChatModelNode, which is a Mock chat model that implements the ToolCallingChatModel\n\t// interface.\n\t// add a Mock InvokableTool to this mock chat model.\n\t// within sub-sub graph B, also add a ToolsNode that will execute this Mock InvokableTool.\n\t// this tool will interrupt with a typed state and an info for end-user\n\t// verify the info thrown by the tool.\n\t// resume with a structured resume data.\n\t// within the Tool, getRunCtx and verify the state and resume data\n\tctrl := gomock.NewController(t)\n\n\t// 1. Define the interrupting tool\n\tmockTool := &mockInterruptingTool{tt: t}\n\n\t// 2. Define the sub-sub-graph (B)\n\tsubSubGraphB := NewGraph[[]*schema.Message, []*schema.Message]()\n\n\t// Mock Chat Model that calls the tool\n\tmockChatModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmockChatModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).Return(&schema.Message{\n\t\tRole: schema.Assistant,\n\t\tToolCalls: []schema.ToolCall{\n\t\t\t{ID: \"tool_call_123\", Function: schema.FunctionCall{Name: \"interrupt_tool\", Arguments: `{\"input\": \"test\"}`}},\n\t\t},\n\t}, nil).AnyTimes()\n\tmockChatModel.EXPECT().WithTools(gomock.Any()).Return(mockChatModel, nil).AnyTimes()\n\n\ttoolsNode, err := NewToolNode(context.Background(), &ToolsNodeConfig{Tools: []tool.BaseTool{mockTool}})\n\tassert.NoError(t, err)\n\n\t_ = subSubGraphB.AddChatModelNode(\"model\", mockChatModel)\n\t_ = subSubGraphB.AddToolsNode(\"tools\", toolsNode)\n\t_ = subSubGraphB.AddEdge(START, \"model\")\n\t_ = subSubGraphB.AddEdge(\"model\", \"tools\")\n\t_ = subSubGraphB.AddEdge(\"tools\", END)\n\n\t// 3. Define sub-graph (A)\n\tsubGraphA := NewGraph[[]*schema.Message, []*schema.Message]()\n\t_ = subGraphA.AddGraphNode(\"sub_graph_b\", subSubGraphB)\n\t_ = subGraphA.AddEdge(START, \"sub_graph_b\")\n\t_ = subGraphA.AddEdge(\"sub_graph_b\", END)\n\n\t// 4. Define root graph\n\trootGraph := NewGraph[[]*schema.Message, []*schema.Message]()\n\t_ = rootGraph.AddGraphNode(\"sub_graph_a\", subGraphA)\n\t_ = rootGraph.AddEdge(START, \"sub_graph_a\")\n\t_ = rootGraph.AddEdge(\"sub_graph_a\", END)\n\n\t// 5. Compile and run\n\tcompiledRootGraph, err := rootGraph.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()),\n\t\tWithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\n\t// First invocation - should interrupt\n\tcheckPointID := \"test-nested-tool-interrupt\"\n\tinitialInput := []*schema.Message{schema.UserMessage(\"hello\")}\n\t_, err = compiledRootGraph.Invoke(context.Background(), initialInput, WithCheckPointID(checkPointID))\n\n\t// 6. Verify the interrupt\n\tassert.Error(t, err)\n\tinterruptInfo, isInterrupt := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt)\n\tassert.NotNil(t, interruptInfo)\n\n\tinterruptContexts := interruptInfo.InterruptContexts\n\tassert.Len(t, interruptContexts, 1) // Only the root cause is returned\n\n\t// Verify the root cause context\n\trootCause := interruptContexts[0]\n\texpectedPath := \"runnable:root;node:sub_graph_a;node:sub_graph_b;node:tools;tool:interrupt_tool:tool_call_123\"\n\tassert.Equal(t, expectedPath, rootCause.Address.String())\n\tassert.True(t, rootCause.IsRootCause)\n\tassert.Equal(t, map[string]any{\"reason\": \"tool maintenance\"}, rootCause.Info)\n\n\t// Verify the parent via the Parent field\n\tassert.NotNil(t, rootCause.Parent)\n\tassert.Equal(t, \"runnable:root;node:sub_graph_a;node:sub_graph_b;node:tools\", rootCause.Parent.Address.String())\n\tassert.False(t, rootCause.Parent.IsRootCause)\n\n\t// 7. Resume execution\n\tctx := ResumeWithData(context.Background(), rootCause.ID, &myResumeData{Message: \"let's continue tool\"})\n\toutput, err := compiledRootGraph.Invoke(ctx, initialInput, WithCheckPointID(checkPointID))\n\n\t// 8. Verify final result\n\tassert.NoError(t, err)\n\tassert.NotNil(t, output)\n\tassert.Len(t, output, 1)\n\tassert.Equal(t, \"Tool resumed successfully\", output[0].Content)\n}\n\nconst PathSegmentTypeProcess AddressSegmentType = \"process\"\n\n// processState is the state for a single sub-process in the batch test.\ntype processState struct {\n\tStep int\n}\n\n// batchState is the composite state for the whole batch lambda.\ntype batchState struct {\n\tProcessStates map[string]*processState\n\tResults       map[string]string\n}\n\ntype processResumeData struct {\n\tInstruction string\n}\n\nfunc init() {\n\tschema.RegisterName[*myInterruptState](\"my_interrupt_state\")\n\tschema.RegisterName[*batchState](\"batch_state\")\n\tschema.RegisterName[*processState](\"process_state\")\n}\n\nfunc TestMultipleInterruptsAndResumes(t *testing.T) {\n\t// define a new lambda node that act as a 'batch' node\n\t// it kick starts 3 parallel processes, each will interrupt on first run, while preserving their own state.\n\t// each of the process should have their own user-facing interrupt info.\n\t// define a new AddressSegmentType for these sub processes.\n\t// the lambda should use StatefulInterrupt to interrupt and preserve the state,\n\t// which is a specific struct type that implements the CompositeInterruptState interface.\n\t// there should also be a specific struct that that implements the CompositeInterruptInfo interface,\n\t// which helps the end-user to fetch the nested interrupt info.\n\t// put this lambda node within a graph and invoke the graph.\n\t// simulate the user getting the flat list of 3 interrupt points using GetInterruptContexts\n\t// the user then decides to resume two of the three interrupt points\n\t// the first resume has resume data, while the second resume does not.(ResumeWithData vs. Resume)\n\t// verify the resume data and state for the resumed interrupt points.\n\tprocessIDs := []string{\"p0\", \"p1\", \"p2\"}\n\n\t// This is the logic for a single \"process\"\n\trunProcess := func(ctx context.Context, id string) (string, error) {\n\t\t// Check if this specific process was interrupted before\n\t\twasInterrupted, hasState, pState := GetInterruptState[*processState](ctx)\n\t\tif !wasInterrupted {\n\t\t\t// First run for this process, interrupt it.\n\t\t\treturn \"\", StatefulInterrupt(ctx,\n\t\t\t\tmap[string]any{\"reason\": \"process \" + id + \" needs input\"},\n\t\t\t\t&processState{Step: 1},\n\t\t\t)\n\t\t}\n\n\t\tassert.True(t, hasState)\n\t\tassert.Equal(t, 1, pState.Step)\n\n\t\t// Check if we are being resumed\n\t\tisResume, hasData, pData := GetResumeContext[*processResumeData](ctx)\n\t\tif !isResume {\n\t\t\t// Not being resumed, so interrupt again.\n\t\t\treturn \"\", StatefulInterrupt(ctx,\n\t\t\t\tmap[string]any{\"reason\": \"process \" + id + \" still needs input\"},\n\t\t\t\tpState,\n\t\t\t)\n\t\t}\n\n\t\t// We are being resumed.\n\t\tif hasData {\n\t\t\t// Resumed with data\n\t\t\treturn \"process \" + id + \" done with instruction: \" + pData.Instruction, nil\n\t\t}\n\t\t// Resumed without data\n\t\treturn \"process \" + id + \" done\", nil\n\t}\n\n\t// This is the main \"batch\" lambda that orchestrates the processes\n\tbatchLambda := InvokableLambda(func(ctx context.Context, _ string) (map[string]string, error) {\n\t\t// Restore the state of the batch node itself\n\t\t_, _, persistedBatchState := GetInterruptState[*batchState](ctx)\n\t\tif persistedBatchState == nil {\n\t\t\tpersistedBatchState = &batchState{\n\t\t\t\tResults: make(map[string]string),\n\t\t\t}\n\t\t}\n\n\t\tvar errs []error\n\n\t\tfor _, id := range processIDs {\n\t\t\t// If this process already completed in a previous run, skip it.\n\t\t\tif _, done := persistedBatchState.Results[id]; done {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\t// Create a sub-context for each process\n\t\t\tsubCtx := AppendAddressSegment(ctx, PathSegmentTypeProcess, id)\n\t\t\tres, err := runProcess(subCtx, id)\n\n\t\t\tif err != nil {\n\t\t\t\t_, ok := IsInterruptRerunError(err)\n\t\t\t\tassert.True(t, ok)\n\t\t\t\terrs = append(errs, err)\n\t\t\t} else {\n\t\t\t\t// Process completed, save its result to the state for the next run.\n\t\t\t\tpersistedBatchState.Results[id] = res\n\t\t\t}\n\t\t}\n\n\t\tif len(errs) > 0 {\n\t\t\treturn nil, CompositeInterrupt(ctx, nil, persistedBatchState, errs...)\n\t\t}\n\n\t\treturn persistedBatchState.Results, nil\n\t})\n\n\tg := NewGraph[string, map[string]string]()\n\t_ = g.AddLambdaNode(\"batch\", batchLambda)\n\t_ = g.AddEdge(START, \"batch\")\n\t_ = g.AddEdge(\"batch\", END)\n\n\tgraph, err := g.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()),\n\t\tWithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\n\t// --- 1. First invocation, all 3 processes should interrupt ---\n\tcheckPointID := \"multi-interrupt-test\"\n\t_, err = graph.Invoke(context.Background(), \"\", WithCheckPointID(checkPointID))\n\n\tassert.Error(t, err)\n\tinterruptInfo, isInterrupt := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt)\n\tinterruptContexts := interruptInfo.InterruptContexts\n\tassert.Len(t, interruptContexts, 3) // Only the 3 root causes\n\n\tfound := make(map[string]bool)\n\taddrToID := make(map[string]string)\n\tvar parentCtx *InterruptCtx\n\tfor _, iCtx := range interruptContexts {\n\t\taddrStr := iCtx.Address.String()\n\t\tfound[addrStr] = true\n\t\taddrToID[addrStr] = iCtx.ID\n\t\tassert.True(t, iCtx.IsRootCause)\n\t\tassert.Equal(t, map[string]any{\"reason\": \"process \" + iCtx.Address[2].ID + \" needs input\"}, iCtx.Info)\n\t\t// Check that all share the same parent\n\t\tassert.NotNil(t, iCtx.Parent)\n\t\tif parentCtx == nil {\n\t\t\tparentCtx = iCtx.Parent\n\t\t\tassert.Equal(t, \"runnable:root;node:batch\", parentCtx.Address.String())\n\t\t\tassert.False(t, parentCtx.IsRootCause)\n\t\t} else {\n\t\t\tassert.Same(t, parentCtx, iCtx.Parent)\n\t\t}\n\t}\n\tassert.True(t, found[\"runnable:root;node:batch;process:p0\"])\n\tassert.True(t, found[\"runnable:root;node:batch;process:p1\"])\n\tassert.True(t, found[\"runnable:root;node:batch;process:p2\"])\n\n\t// --- 2. Second invocation, resume 2 of 3 processes ---\n\t// Resume p0 with data, and p2 without data. p1 remains interrupted.\n\tresumeCtx := ResumeWithData(context.Background(), addrToID[\"runnable:root;node:batch;process:p0\"], &processResumeData{Instruction: \"do it\"})\n\tresumeCtx = Resume(resumeCtx, addrToID[\"runnable:root;node:batch;process:p2\"])\n\n\t_, err = graph.Invoke(resumeCtx, \"\", WithCheckPointID(checkPointID))\n\n\t// Expect an interrupt again, but only for p1\n\tassert.Error(t, err)\n\tinterruptInfo2, isInterrupt2 := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt2)\n\tinterruptContexts2 := interruptInfo2.InterruptContexts\n\tassert.Len(t, interruptContexts2, 1) // Only p1 is left\n\trootCause2 := interruptContexts2[0]\n\tassert.Equal(t, \"runnable:root;node:batch;process:p1\", rootCause2.Address.String())\n\tassert.NotNil(t, rootCause2.Parent)\n\tassert.Equal(t, \"runnable:root;node:batch\", rootCause2.Parent.Address.String())\n\n\t// --- 3. Third invocation, resume the last process ---\n\tfinalResumeCtx := Resume(context.Background(), rootCause2.ID)\n\tfinalOutput, err := graph.Invoke(finalResumeCtx, \"\", WithCheckPointID(checkPointID))\n\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"process p0 done with instruction: do it\", finalOutput[\"p0\"])\n\tassert.Equal(t, \"process p1 done\", finalOutput[\"p1\"])\n\tassert.Equal(t, \"process p2 done\", finalOutput[\"p2\"])\n}\n\n// toolsNodeResumeTargetCallback captures isResumeTarget for ToolsNode during OnStart\ntype toolsNodeResumeTargetCallback struct {\n\tmu                sync.Mutex\n\tisResumeTargetLog []bool\n}\n\nfunc (c *toolsNodeResumeTargetCallback) OnStart(ctx context.Context, info *callbacks.RunInfo, _ callbacks.CallbackInput) context.Context {\n\tif info.Component == ComponentOfToolsNode {\n\t\tisResumeTarget, _, _ := GetResumeContext[any](ctx)\n\t\tc.mu.Lock()\n\t\tc.isResumeTargetLog = append(c.isResumeTargetLog, isResumeTarget)\n\t\tc.mu.Unlock()\n\t}\n\treturn ctx\n}\n\nfunc (c *toolsNodeResumeTargetCallback) OnEnd(ctx context.Context, _ *callbacks.RunInfo, _ callbacks.CallbackOutput) context.Context {\n\treturn ctx\n}\n\nfunc (c *toolsNodeResumeTargetCallback) OnError(ctx context.Context, _ *callbacks.RunInfo, _ error) context.Context {\n\treturn ctx\n}\n\nfunc (c *toolsNodeResumeTargetCallback) OnStartWithStreamInput(ctx context.Context, _ *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {\n\tinput.Close()\n\treturn ctx\n}\n\nfunc (c *toolsNodeResumeTargetCallback) OnEndWithStreamOutput(ctx context.Context, _ *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context {\n\toutput.Close()\n\treturn ctx\n}\n\n// mockReentryTool is a helper for the reentry test\ntype mockReentryTool struct {\n\tt                     *testing.T\n\tmu                    sync.Mutex\n\tisResumeTargetByRunID map[string]bool\n}\n\nfunc (t *mockReentryTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName:        \"reentry_tool\",\n\t\tDesc:        \"A tool that can be re-entered in a resumed graph.\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\"input\": {Type: schema.String}}),\n\t}, nil\n}\n\nfunc (t *mockReentryTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {\n\twasInterrupted, hasState, _ := tool.GetInterruptState[any](ctx)\n\tisResume, hasData, data := tool.GetResumeContext[*myResumeData](ctx)\n\n\tcallID := GetToolCallID(ctx)\n\n\tt.mu.Lock()\n\tif t.isResumeTargetByRunID != nil {\n\t\tt.isResumeTargetByRunID[callID] = isResume\n\t}\n\tt.mu.Unlock()\n\n\t// Special handling for the re-entrant call to make assertions explicit.\n\tif callID == \"call_3\" {\n\t\tif !isResume {\n\t\t\t// This is the first run of the re-entrant call. Its context must be clean.\n\t\t\t// This is the core assertion for this test.\n\t\t\tassert.False(t.t, wasInterrupted, \"re-entrant call 'call_3' should not have been interrupted on its first run\")\n\t\t\tassert.False(t.t, hasState, \"re-entrant call 'call_3' should not have state on its first run\")\n\t\t\t// Now, interrupt it as part of the test flow.\n\t\t\treturn \"\", tool.StatefulInterrupt(ctx, nil, \"some state for \"+callID)\n\t\t}\n\t\t// This is the resumed run of the re-entrant call.\n\t\tassert.True(t.t, wasInterrupted, \"resumed call 'call_3' must have been interrupted\")\n\t\tassert.True(t.t, hasData, \"resumed call 'call_3' should have data\")\n\t\treturn \"Resumed \" + data.Message, nil\n\t}\n\n\t// Standard logic for the initial calls (call_1, call_2)\n\tif !wasInterrupted {\n\t\t// First run for call_1 and call_2, should interrupt.\n\t\treturn \"\", tool.StatefulInterrupt(ctx, nil, \"some state for \"+callID)\n\t}\n\n\t// From here, wasInterrupted is true for call_1 and call_2.\n\tif isResume {\n\t\t// The user is explicitly resuming this call.\n\t\tassert.True(t.t, hasData, \"call %s should have resume data\", callID)\n\t\treturn \"Resumed \" + data.Message, nil\n\t}\n\n\t// The tool was interrupted before, but is not being resumed now. Re-interrupt.\n\treturn \"\", tool.StatefulInterrupt(ctx, nil, \"some state for \"+callID)\n}\n\nfunc TestReentryForResumedTools(t *testing.T) {\n\t// create a 'ReAct' style graph with a ChatModel node and a ToolsNode.\n\t// within the ToolsNode there is an interruptible tool that will emit interrupt on first run.\n\t// During the first invocation of the graph, there should be two tool calls (of the same tool) that interrupt.\n\t// The user chooses to resume one of the interrupted tool call in second invocation,\n\t// and this time, the resumed tool call should be successful, while the other should interrupt immediately again.\n\t// The user then chooses to resume the other interrupted tool call in third invocation,\n\t// and this time, the ChatModel decides to call the tool again,\n\t// and this time the tool's runCtx should think it was not interrupted nor resumed.\n\tctrl := gomock.NewController(t)\n\n\t// 1. Define the interrupting tool and callback\n\treentryTool := &mockReentryTool{t: t, isResumeTargetByRunID: make(map[string]bool)}\n\ttoolsNodeCB := &toolsNodeResumeTargetCallback{}\n\n\t// 2. Define the graph\n\tg := NewGraph[[]*schema.Message, *schema.Message]()\n\n\t// Mock Chat Model that drives the ReAct loop\n\tmockChatModel := mockModel.NewMockToolCallingChatModel(ctrl)\n\ttoolsNode, err := NewToolNode(context.Background(), &ToolsNodeConfig{Tools: []tool.BaseTool{reentryTool}})\n\tassert.NoError(t, err)\n\n\t// Expectation for the 1st invocation: model returns two tool calls\n\tmockChatModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).Return(&schema.Message{\n\t\tRole: schema.Assistant,\n\t\tToolCalls: []schema.ToolCall{\n\t\t\t{ID: \"call_1\", Function: schema.FunctionCall{Name: \"reentry_tool\", Arguments: `{\"input\": \"a\"}`}},\n\t\t\t{ID: \"call_2\", Function: schema.FunctionCall{Name: \"reentry_tool\", Arguments: `{\"input\": \"b\"}`}},\n\t\t},\n\t}, nil).Times(1)\n\n\t// Expectation for the 2nd invocation (after resuming call_1): model does nothing, graph continues\n\t// Expectation for the 3rd invocation (after resuming call_2): model calls the tool again\n\tmockChatModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\treturn &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{ID: \"call_3\", Function: schema.FunctionCall{Name: \"reentry_tool\", Arguments: `{\"input\": \"c\"}`}},\n\t\t\t},\n\t\t}, nil\n\t}).Times(1)\n\n\t// Expectation for the final invocation: model returns final answer\n\tmockChatModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).Return(&schema.Message{\n\t\tRole:    schema.Assistant,\n\t\tContent: \"all done\",\n\t}, nil).Times(1)\n\n\t_ = g.AddChatModelNode(\"model\", mockChatModel)\n\t_ = g.AddToolsNode(\"tools\", toolsNode)\n\t_ = g.AddEdge(START, \"model\")\n\n\t// Add the crucial branch to decide whether to call tools or end.\n\tmodelBranch := func(ctx context.Context, msg *schema.Message) (string, error) {\n\t\tif len(msg.ToolCalls) > 0 {\n\t\t\treturn \"tools\", nil\n\t\t}\n\t\treturn END, nil\n\t}\n\terr = g.AddBranch(\"model\", NewGraphBranch(modelBranch, map[string]bool{\"tools\": true, END: true}))\n\tassert.NoError(t, err)\n\n\t_ = g.AddEdge(\"tools\", \"model\") // Loop back for ReAct style\n\n\t// 3. Compile and run\n\tgraph, err := g.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()),\n\t\tWithGraphName(\"root\"))\n\tassert.NoError(t, err)\n\tcheckPointID := \"reentry-test\"\n\n\t// --- 1. First invocation: call_1 and call_2 should interrupt ---\n\t_, err = graph.Invoke(context.Background(), []*schema.Message{schema.UserMessage(\"start\")}, WithCheckPointID(checkPointID), WithCallbacks(toolsNodeCB))\n\tassert.Error(t, err)\n\tinterruptInfo1, _ := ExtractInterruptInfo(err)\n\tinterrupts1 := interruptInfo1.InterruptContexts\n\tassert.Len(t, interrupts1, 2) // Only the two tool calls\n\tfound1 := make(map[string]bool)\n\taddrToID1 := make(map[string]string)\n\tfor _, iCtx := range interrupts1 {\n\t\taddrStr := iCtx.Address.String()\n\t\tfound1[addrStr] = true\n\t\taddrToID1[addrStr] = iCtx.ID\n\t\tassert.True(t, iCtx.IsRootCause)\n\t\tassert.NotNil(t, iCtx.Parent)\n\t\tassert.Equal(t, \"runnable:root;node:tools\", iCtx.Parent.Address.String())\n\t}\n\tassert.True(t, found1[\"runnable:root;node:tools;tool:reentry_tool:call_1\"])\n\tassert.True(t, found1[\"runnable:root;node:tools;tool:reentry_tool:call_2\"])\n\n\t// First invocation: neither call_1 nor call_2 should be resume targets\n\tassert.False(t, reentryTool.isResumeTargetByRunID[\"call_1\"], \"first run: call_1 should not be resume target\")\n\tassert.False(t, reentryTool.isResumeTargetByRunID[\"call_2\"], \"first run: call_2 should not be resume target\")\n\n\t// First invocation: ToolsNode should NOT be a resume target\n\tassert.Len(t, toolsNodeCB.isResumeTargetLog, 1, \"ToolsNode OnStart should be called once in first invocation\")\n\tassert.False(t, toolsNodeCB.isResumeTargetLog[0], \"first run: ToolsNode should NOT be resume target\")\n\n\t// Clear for next invocation\n\treentryTool.isResumeTargetByRunID = make(map[string]bool)\n\ttoolsNodeCB.isResumeTargetLog = nil\n\n\t// --- 2. Second invocation: resume call_1, expect call_2 to interrupt again ---\n\tresumeCtx2 := ResumeWithData(context.Background(), addrToID1[\"runnable:root;node:tools;tool:reentry_tool:call_1\"],\n\t\t&myResumeData{Message: \"resume call 1\"})\n\t_, err = graph.Invoke(resumeCtx2, []*schema.Message{schema.UserMessage(\"start\")}, WithCheckPointID(checkPointID), WithCallbacks(toolsNodeCB))\n\tassert.Error(t, err)\n\tinterruptInfo2, _ := ExtractInterruptInfo(err)\n\tinterrupts2 := interruptInfo2.InterruptContexts\n\tassert.Len(t, interrupts2, 1) // Only call_2\n\trootCause2 := interrupts2[0]\n\tassert.Equal(t, \"runnable:root;node:tools;tool:reentry_tool:call_2\", rootCause2.Address.String())\n\tassert.NotNil(t, rootCause2.Parent)\n\tassert.Equal(t, \"runnable:root;node:tools\", rootCause2.Parent.Address.String())\n\n\t// Second invocation: call_1 is resumed, call_2 is NOT resumed (re-interrupts)\n\tassert.True(t, reentryTool.isResumeTargetByRunID[\"call_1\"], \"second run: call_1 should be resume target\")\n\tassert.False(t, reentryTool.isResumeTargetByRunID[\"call_2\"], \"second run: call_2 should NOT be resume target (it re-interrupts)\")\n\n\t// Second invocation: ToolsNode SHOULD be a resume target (because call_1 child is being resumed)\n\tassert.Len(t, toolsNodeCB.isResumeTargetLog, 1, \"ToolsNode OnStart should be called once in second invocation\")\n\tassert.True(t, toolsNodeCB.isResumeTargetLog[0], \"second run: ToolsNode SHOULD be resume target (child call_1 is being resumed)\")\n\n\t// Clear for next invocation\n\treentryTool.isResumeTargetByRunID = make(map[string]bool)\n\ttoolsNodeCB.isResumeTargetLog = nil\n\n\t// --- 3. Third invocation: resume call_2, model makes a new call (call_3) which should interrupt ---\n\tresumeCtx3 := ResumeWithData(context.Background(), rootCause2.ID, &myResumeData{Message: \"resume call 2\"})\n\t_, err = graph.Invoke(resumeCtx3, []*schema.Message{schema.UserMessage(\"start\")}, WithCheckPointID(checkPointID), WithCallbacks(toolsNodeCB))\n\tassert.Error(t, err)\n\tinterruptInfo3, _ := ExtractInterruptInfo(err)\n\tinterrupts3 := interruptInfo3.InterruptContexts\n\tassert.Len(t, interrupts3, 1) // Only call_3\n\trootCause3 := interrupts3[0]\n\tassert.Equal(t, \"runnable:root;node:tools;tool:reentry_tool:call_3\", rootCause3.Address.String()) // Note: this is the new call_3\n\tassert.NotNil(t, rootCause3.Parent)\n\tassert.Equal(t, \"runnable:root;node:tools\", rootCause3.Parent.Address.String())\n\n\t// Third invocation: call_2 is resumed, call_3 is new (not resumed)\n\tassert.True(t, reentryTool.isResumeTargetByRunID[\"call_2\"], \"third run: call_2 should be resume target\")\n\tassert.False(t, reentryTool.isResumeTargetByRunID[\"call_3\"], \"third run: call_3 should NOT be resume target (it's new)\")\n\n\t// Third invocation: ToolsNode is called twice (once for call_2 resume, once for call_3 new)\n\t// First call: ToolsNode SHOULD be resume target (call_2 is being resumed)\n\t// Second call: ToolsNode should NOT be resume target (call_3 is new, no children to resume)\n\tassert.Len(t, toolsNodeCB.isResumeTargetLog, 2, \"ToolsNode OnStart should be called twice in third invocation\")\n\tassert.True(t, toolsNodeCB.isResumeTargetLog[0], \"third run first ToolsNode call: SHOULD be resume target (child call_2 is being resumed)\")\n\tassert.False(t, toolsNodeCB.isResumeTargetLog[1], \"third run second ToolsNode call: should NOT be resume target (call_3 is new)\")\n\n\t// Clear for next invocation\n\treentryTool.isResumeTargetByRunID = make(map[string]bool)\n\ttoolsNodeCB.isResumeTargetLog = nil\n\n\t// --- 4. Final invocation: resume call_3, expect final answer ---\n\tresumeCtx4 := ResumeWithData(context.Background(), rootCause3.ID,\n\t\t&myResumeData{Message: \"resume call 3\"})\n\toutput, err := graph.Invoke(resumeCtx4, []*schema.Message{schema.UserMessage(\"start\")}, WithCheckPointID(checkPointID), WithCallbacks(toolsNodeCB))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"all done\", output.Content)\n\n\t// Fourth invocation: call_3 is resumed\n\tassert.True(t, reentryTool.isResumeTargetByRunID[\"call_3\"], \"fourth run: call_3 should be resume target\")\n\n\t// Fourth invocation: ToolsNode SHOULD be resume target (call_3 is being resumed)\n\tassert.Len(t, toolsNodeCB.isResumeTargetLog, 1, \"ToolsNode OnStart should be called once in fourth invocation\")\n\tassert.True(t, toolsNodeCB.isResumeTargetLog[0], \"fourth run: ToolsNode SHOULD be resume target (child call_3 is being resumed)\")\n}\n\n// mockInterruptingTool is a helper for the nested tool interrupt test\ntype mockInterruptingTool struct {\n\ttt *testing.T\n}\n\nfunc (t *mockInterruptingTool) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"interrupt_tool\",\n\t\tDesc: \"A tool that interrupts execution.\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\"input\": {Type: schema.String, Desc: \"Some input\", Required: true},\n\t\t}),\n\t}, nil\n}\n\nfunc (t *mockInterruptingTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\tvar args map[string]string\n\t_ = json.Unmarshal([]byte(argumentsInJSON), &args)\n\n\twasInterrupted, hasState, state := tool.GetInterruptState[*myInterruptState](ctx)\n\tif !wasInterrupted {\n\t\t// First run: interrupt\n\t\treturn \"\", tool.StatefulInterrupt(ctx,\n\t\t\tmap[string]any{\"reason\": \"tool maintenance\"},\n\t\t\t&myInterruptState{OriginalInput: args[\"input\"]},\n\t\t)\n\t}\n\n\t// Second (resumed) run\n\tassert.True(t.tt, hasState)\n\tassert.Equal(t.tt, \"test\", state.OriginalInput)\n\n\tisResume, hasData, data := tool.GetResumeContext[*myResumeData](ctx)\n\tassert.True(t.tt, isResume)\n\tassert.True(t.tt, hasData)\n\tassert.Equal(t.tt, \"let's continue tool\", data.Message)\n\n\treturn \"Tool resumed successfully\", nil\n}\n\nfunc TestGraphInterruptWithinLambda(t *testing.T) {\n\t// this test case aims to verify behaviors when a standalone graph is within a lambda,\n\t// which in turn is within the root graph.\n\t// the expected behavior is:\n\t// - internal graph will naturally append to the Address\n\t// - internal graph interrupts, where the Address includes steps for both the root graph and the internal graph\n\t// - lambda extracts InterruptInfo, then GetInterruptContexts\n\t// - lambda then acts as a composite node, uses CompositeInterrupt to pass up the\n\t//   internal interrupt points\n\t// - the root graph interrupts\n\t// - end-user extracts the interrupt ID and related info\n\t// - end-user uses ResumeWithData to resume the ID\n\t// - lambda node resumes, invokes the inner graph as usual\n\t// - the internal graph resumes the interrupted node\n\t// To implement this test, within the internal graph you can define another lambda node that can interrupt resume.\n\n\t// 1. Define the innermost lambda that actually interrupts\n\tinterruptingLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\twasInterrupted, hasState, state := GetInterruptState[*myInterruptState](ctx)\n\t\tif !wasInterrupted {\n\t\t\treturn \"\", StatefulInterrupt(ctx, \"inner interrupt\", &myInterruptState{OriginalInput: input})\n\t\t}\n\n\t\tassert.True(t, hasState)\n\t\tassert.Equal(t, \"top level input\", state.OriginalInput)\n\n\t\tisResume, hasData, data := GetResumeContext[*myResumeData](ctx)\n\t\tassert.True(t, isResume)\n\t\tassert.True(t, hasData)\n\t\tassert.Equal(t, \"resume inner\", data.Message)\n\n\t\treturn \"inner lambda resumed successfully\", nil\n\t})\n\n\t// 2. Define the internal graph that contains the interrupting lambda\n\tinnerGraph := NewGraph[string, string]()\n\t_ = innerGraph.AddLambdaNode(\"inner_lambda\", interruptingLambda)\n\t_ = innerGraph.AddEdge(START, \"inner_lambda\")\n\t_ = innerGraph.AddEdge(\"inner_lambda\", END)\n\t// Give the inner graph a name so it can create its \"runnable\" addr step.\n\tcompiledInnerGraph, err := innerGraph.Compile(context.Background(), WithGraphName(\"inner\"), WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t// 3. Define the outer lambda that acts as a composite node\n\tcompositeLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\t// The lambda invokes the inner graph. If the inner graph interrupts, this lambda\n\t\t// must act as a proper composite node and wrap the error.\n\t\toutput, err := compiledInnerGraph.Invoke(ctx, input, WithCheckPointID(\"inner-cp\"))\n\t\tif err != nil {\n\t\t\t_, isInterrupt := ExtractInterruptInfo(err)\n\t\t\tif !isInterrupt {\n\t\t\t\treturn \"\", err // Not an interrupt, just fail\n\t\t\t}\n\n\t\t\t// The composite interrupt itself can be stateless, as it's just a wrapper.\n\t\t\t// It signals to the framework to look inside the subErrs and correctly\n\t\t\t// prepend the current addr to the paths of the inner interrupts.\n\t\t\treturn \"\", CompositeInterrupt(ctx, \"composite interrupt from lambda\", nil, err)\n\t\t}\n\t\treturn output, nil\n\t})\n\n\t// 4. Define the root graph\n\trootGraph := NewGraph[string, string]()\n\t_ = rootGraph.AddLambdaNode(\"composite_lambda\", compositeLambda)\n\t_ = rootGraph.AddEdge(START, \"composite_lambda\")\n\t_ = rootGraph.AddEdge(\"composite_lambda\", END)\n\t// Give the root graph a name for its \"runnable\" addr step.\n\tcompiledRootGraph, err := rootGraph.Compile(context.Background(), WithGraphName(\"root\"), WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t// 5. First invocation - should interrupt\n\tcheckPointID := \"graph-in-lambda-test\"\n\t_, err = compiledRootGraph.Invoke(context.Background(), \"top level input\", WithCheckPointID(checkPointID))\n\n\t// 6. Verify the interrupt\n\tassert.Error(t, err)\n\tinterruptInfo, isInterrupt := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt)\n\tinterruptContexts := interruptInfo.InterruptContexts\n\tassert.Len(t, interruptContexts, 1) // Only the root cause is returned\n\n\t// The addr is now fully qualified, including the runnable steps from both graphs.\n\trootCause := interruptContexts[0]\n\texpectedPath := \"runnable:root;node:composite_lambda;runnable:inner;node:inner_lambda\"\n\tassert.Equal(t, expectedPath, rootCause.Address.String())\n\tassert.Equal(t, \"inner interrupt\", rootCause.Info)\n\tassert.True(t, rootCause.IsRootCause)\n\n\t// Check parent hierarchy\n\tassert.NotNil(t, rootCause.Parent)\n\tassert.Equal(t, \"runnable:root;node:composite_lambda;runnable:inner\", rootCause.Parent.Address.String())\n\tassert.Nil(t, rootCause.Parent.Info) // The inner runnable doesn't have its own info\n\tassert.False(t, rootCause.Parent.IsRootCause)\n\n\t// Check grandparent\n\tassert.NotNil(t, rootCause.Parent.Parent)\n\tassert.Equal(t, \"runnable:root;node:composite_lambda\", rootCause.Parent.Parent.Address.String())\n\tassert.Equal(t, \"composite interrupt from lambda\", rootCause.Parent.Parent.Info)\n\tassert.False(t, rootCause.Parent.Parent.IsRootCause)\n\n\t// 7. Resume execution using the complete, fully-qualified ID\n\tresumeCtx := ResumeWithData(context.Background(), rootCause.ID, &myResumeData{Message: \"resume inner\"})\n\tfinalOutput, err := compiledRootGraph.Invoke(resumeCtx, \"top level input\", WithCheckPointID(checkPointID))\n\n\t// 8. Verify final result\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"inner lambda resumed successfully\", finalOutput)\n}\n\nfunc TestLegacyInterrupt(t *testing.T) {\n\t// this test case aims to test the behavior of the deprecated InterruptAndRerun,\n\t// NewInterruptAndRerunErr within CompositeInterrupt.\n\t// Define two sub-processes(functions), one interrupts with InterruptAndRerun,\n\t// the other interrupts with NewInterruptAndRerunErr.\n\t// create a lambda as a composite node, within the lambda invokes the two sub-processes.\n\t// create the graph, add lambda node and invoke it.\n\t// after verifying the interrupt points, just invokes again without explicit resume.\n\t// verify the same interrupt IDs again.\n\t// then finally Resume() the graph.\n\n\t// 1. Define the sub-processes that use legacy and modern interrupts\n\tsubProcess1 := func(ctx context.Context) (string, error) {\n\t\tisResume, _, data := GetResumeContext[string](ctx)\n\t\tif isResume {\n\t\t\treturn data, nil\n\t\t}\n\t\treturn \"\", deprecatedInterruptAndRerun\n\t}\n\tsubProcess2 := func(ctx context.Context) (string, error) {\n\t\tisResume, _, data := GetResumeContext[string](ctx)\n\t\tif isResume {\n\t\t\treturn data, nil\n\t\t}\n\t\treturn \"\", deprecatedInterruptAndRerunErr(\"legacy info\")\n\t}\n\tsubProcess3 := func(ctx context.Context) (string, error) {\n\t\tisResume, _, data := GetResumeContext[string](ctx)\n\t\tif isResume {\n\t\t\treturn data, nil\n\t\t}\n\t\t// Use the modern, addr-aware interrupt function\n\t\treturn \"\", Interrupt(ctx, \"modern info\")\n\t}\n\n\t// 2. Define the composite lambda\n\tcompositeLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\t// If the lambda itself is being resumed, it means the whole process is done.\n\t\tisResume, _, data := GetResumeContext[string](ctx)\n\n\t\t// Run sub-processes and collect their errors\n\t\tvar (\n\t\t\terrs   []error\n\t\t\toutStr string\n\t\t)\n\n\t\tconst PathStepCustom AddressSegmentType = \"custom\"\n\t\tsubCtx1 := AppendAddressSegment(ctx, PathStepCustom, \"1\")\n\t\tout1, err1 := subProcess1(subCtx1)\n\t\tif err1 != nil {\n\t\t\t// Wrap the legacy error to give it a addr\n\t\t\twrappedErr := WrapInterruptAndRerunIfNeeded(ctx, AddressSegment{Type: PathStepCustom, ID: \"1\"}, err1)\n\t\t\terrs = append(errs, wrappedErr)\n\t\t} else {\n\t\t\toutStr += out1\n\t\t}\n\t\tsubCtx2 := AppendAddressSegment(ctx, PathStepCustom, \"2\")\n\t\tout2, err2 := subProcess2(subCtx2)\n\t\tif err2 != nil {\n\t\t\t// Wrap the legacy error to give it a addr\n\t\t\twrappedErr := WrapInterruptAndRerunIfNeeded(ctx, AddressSegment{Type: PathStepCustom, ID: \"2\"}, err2)\n\t\t\terrs = append(errs, wrappedErr)\n\t\t} else {\n\t\t\toutStr += out2\n\t\t}\n\t\tsubCtx3 := AppendAddressSegment(ctx, PathStepCustom, \"3\")\n\t\tout3, err3 := subProcess3(subCtx3)\n\t\tif err3 != nil {\n\t\t\t// The error from Interrupt() is already addr-aware. WrapInterruptAndRerunIfNeeded\n\t\t\t// should handle this gracefully and return the error as-is.\n\t\t\twrappedErr := WrapInterruptAndRerunIfNeeded(ctx, AddressSegment{Type: PathStepCustom, ID: \"3\"}, err3)\n\t\t\terrs = append(errs, wrappedErr)\n\t\t} else {\n\t\t\toutStr += out3\n\t\t}\n\n\t\tif len(errs) > 0 {\n\t\t\t// Return a composite interrupt containing the wrapped legacy errors\n\t\t\treturn \"\", CompositeInterrupt(ctx, \"legacy composite\", nil, errs...)\n\t\t}\n\n\t\tif isResume {\n\t\t\toutStr = outStr + \" \" + data\n\t\t}\n\n\t\treturn outStr, nil\n\t})\n\n\t// 3. Create and compile the graph\n\trootGraph := NewGraph[string, string]()\n\t_ = rootGraph.AddLambdaNode(\"legacy_composite\", compositeLambda)\n\t_ = rootGraph.AddEdge(START, \"legacy_composite\")\n\t_ = rootGraph.AddEdge(\"legacy_composite\", END)\n\tcompiledGraph, err := rootGraph.Compile(context.Background(), WithGraphName(\"root\"), WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t// 4. First invocation - should interrupt\n\tcheckPointID := \"legacy-interrupt-test\"\n\t_, err = compiledGraph.Invoke(context.Background(), \"input\", WithCheckPointID(checkPointID))\n\n\t// 5. Verify the three interrupt points\n\tassert.Error(t, err)\n\tinfo, isInterrupt := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt)\n\tassert.Len(t, info.InterruptContexts, 3) // Only the 3 root causes\n\n\tfound := make(map[string]any)\n\taddrToID := make(map[string]string)\n\tvar parentCtx *InterruptCtx\n\tfor _, iCtx := range info.InterruptContexts {\n\t\taddrStr := iCtx.Address.String()\n\t\tfound[addrStr] = iCtx.Info\n\t\taddrToID[addrStr] = iCtx.ID\n\t\tassert.True(t, iCtx.IsRootCause)\n\t\t// Check parent\n\t\tassert.NotNil(t, iCtx.Parent)\n\t\tif parentCtx == nil {\n\t\t\tparentCtx = iCtx.Parent\n\t\t\tassert.Equal(t, \"runnable:root;node:legacy_composite\", parentCtx.Address.String())\n\t\t\tassert.Equal(t, \"legacy composite\", parentCtx.Info)\n\t\t\tassert.False(t, parentCtx.IsRootCause)\n\t\t} else {\n\t\t\tassert.Same(t, parentCtx, iCtx.Parent)\n\t\t}\n\t}\n\texpectedID1 := \"runnable:root;node:legacy_composite;custom:1\"\n\texpectedID2 := \"runnable:root;node:legacy_composite;custom:2\"\n\texpectedID3 := \"runnable:root;node:legacy_composite;custom:3\"\n\tassert.Contains(t, found, expectedID1)\n\tassert.Nil(t, found[expectedID1]) // From InterruptAndRerun\n\tassert.Contains(t, found, expectedID2)\n\tassert.Equal(t, \"legacy info\", found[expectedID2]) // From NewInterruptAndRerunErr\n\tassert.Contains(t, found, expectedID3)\n\tassert.Equal(t, \"modern info\", found[expectedID3]) // From Interrupt\n\n\t// 6. Second invocation (re-run without resume) - should yield the same interrupts\n\t_, err = compiledGraph.Invoke(context.Background(), \"input\", WithCheckPointID(checkPointID))\n\tassert.Error(t, err)\n\tinfo2, isInterrupt2 := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt2)\n\tassert.Len(t, info2.InterruptContexts, 3, \"Should have the same number of interrupts on re-run\")\n\n\t// 7. Third invocation - Resume all three interrupt points with specific data\n\tresumeData := map[string]any{\n\t\taddrToID[expectedID1]: \"output1\",\n\t\taddrToID[expectedID2]: \"output2\",\n\t\taddrToID[expectedID3]: \"output3\",\n\t}\n\tresumeCtx := BatchResumeWithData(context.Background(), resumeData)\n\t// TODO: The legacy interrupt wrapping does not currently work correctly with BatchResumeWithData.\n\t// The graph re-interrupts instead of completing. This should be fixed in the core framework.\n\t_, err = compiledGraph.Invoke(resumeCtx, \"input\", WithCheckPointID(checkPointID))\n\tassert.Error(t, err)\n}\n\ntype wrapperToolForTest struct {\n\tcompiledGraph     Runnable[string, string]\n\tisResumeTargetLog []bool\n}\n\nfunc (w *wrapperToolForTest) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"wrapperTool\",\n\t\tDesc: \"A tool that wraps a nested graph\",\n\t}, nil\n}\n\nfunc (w *wrapperToolForTest) InvokableRun(ctx context.Context, input string, opts ...tool.Option) (string, error) {\n\tisResumeTarget, _, _ := tool.GetResumeContext[any](ctx)\n\tw.isResumeTargetLog = append(w.isResumeTargetLog, isResumeTarget)\n\n\tresult, err := w.compiledGraph.Invoke(ctx, input)\n\tif err != nil {\n\t\tif _, ok := ExtractInterruptInfo(err); ok {\n\t\t\treturn \"\", tool.CompositeInterrupt(ctx, \"wrapper tool interrupt\", nil, err)\n\t\t}\n\t\treturn \"\", err\n\t}\n\treturn result, nil\n}\n\nfunc TestToolCompositeInterruptWithNestedGraphInterrupt(t *testing.T) {\n\tctx := context.Background()\n\n\tvar innerNodeIsResumeTarget bool\n\tsubSubGraph := NewGraph[string, string]()\n\terr := subSubGraph.AddLambdaNode(\"interruptNode\", InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\twasInterrupted, _, _ := GetInterruptState[any](ctx)\n\t\tif !wasInterrupted {\n\t\t\treturn \"\", Interrupt(ctx, \"sub-sub graph interrupt info\")\n\t\t}\n\t\tisResumeTarget, _, _ := GetResumeContext[any](ctx)\n\t\tinnerNodeIsResumeTarget = isResumeTarget\n\t\treturn \"resumed successfully\", nil\n\t}))\n\tassert.NoError(t, err)\n\tassert.NoError(t, subSubGraph.AddEdge(START, \"interruptNode\"))\n\tassert.NoError(t, subSubGraph.AddEdge(\"interruptNode\", END))\n\n\tnestedGraph := NewGraph[string, string]()\n\terr = nestedGraph.AddGraphNode(\"subSubGraph\", subSubGraph)\n\tassert.NoError(t, err)\n\tassert.NoError(t, nestedGraph.AddEdge(START, \"subSubGraph\"))\n\tassert.NoError(t, nestedGraph.AddEdge(\"subSubGraph\", END))\n\tcompiledNestedGraph, err := nestedGraph.Compile(ctx)\n\tassert.NoError(t, err)\n\n\twrapperTool := &wrapperToolForTest{compiledGraph: compiledNestedGraph.(Runnable[string, string])}\n\n\ttoolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{Tools: []tool.BaseTool{wrapperTool}})\n\tassert.NoError(t, err)\n\n\touterGraph := NewGraph[*schema.Message, []*schema.Message]()\n\terr = outerGraph.AddToolsNode(\"tools\", toolsNode)\n\tassert.NoError(t, err)\n\tassert.NoError(t, outerGraph.AddEdge(START, \"tools\"))\n\tassert.NoError(t, outerGraph.AddEdge(\"tools\", END))\n\n\tcompiledOuterGraph, err := outerGraph.Compile(ctx, WithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\tcheckpointID := \"test-wrapper-tool-resume\"\n\tinputMsg := &schema.Message{\n\t\tRole: schema.Assistant,\n\t\tToolCalls: []schema.ToolCall{\n\t\t\t{ID: \"call_1\", Function: schema.FunctionCall{Name: \"wrapperTool\", Arguments: `\"test\"`}},\n\t\t},\n\t}\n\n\t_, err = compiledOuterGraph.Invoke(ctx, inputMsg, WithCheckPointID(checkpointID))\n\tassert.Error(t, err)\n\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok, \"should be an interrupt error\")\n\tassert.NotNil(t, info)\n\tassert.NotEmpty(t, info.InterruptContexts)\n\n\trootCause := info.InterruptContexts[0]\n\tassert.Equal(t, \"sub-sub graph interrupt info\", rootCause.Info)\n\tassert.True(t, rootCause.IsRootCause)\n\n\tvar wrapperToolParent *InterruptCtx\n\tfor p := rootCause.Parent; p != nil; p = p.Parent {\n\t\tif p.Info == \"wrapper tool interrupt\" {\n\t\t\twrapperToolParent = p\n\t\t\tbreak\n\t\t}\n\t}\n\tassert.NotNil(t, wrapperToolParent, \"should have parent from wrapper tool with info 'wrapper tool interrupt'\")\n\n\tassert.Len(t, wrapperTool.isResumeTargetLog, 1)\n\tassert.False(t, wrapperTool.isResumeTargetLog[0], \"first invocation: wrapper tool should not be resume target\")\n\n\tresumeCtx := Resume(ctx, rootCause.ID)\n\t_, err = compiledOuterGraph.Invoke(resumeCtx, inputMsg, WithCheckPointID(checkpointID))\n\tassert.NoError(t, err)\n\n\tassert.True(t, innerNodeIsResumeTarget, \"inner node should be resume target\")\n\n\tassert.Len(t, wrapperTool.isResumeTargetLog, 2)\n\tassert.True(t, wrapperTool.isResumeTargetLog[1], \"second invocation: wrapper tool should be resume target because its child is targeted\")\n}\n"
  },
  {
    "path": "compose/runnable.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Runnable is the interface for an executable object. Graph, Chain can be compiled into Runnable.\n// runnable is the core conception of eino, we do downgrade compatibility for four data flow patterns,\n// and can automatically connect components that only implement one or more methods.\n// eg, if a component only implements Stream() method, you can still call Invoke() to convert stream output to invoke output.\ntype Runnable[I, O any] interface {\n\tInvoke(ctx context.Context, input I, opts ...Option) (output O, err error)\n\tStream(ctx context.Context, input I, opts ...Option) (output *schema.StreamReader[O], err error)\n\tCollect(ctx context.Context, input *schema.StreamReader[I], opts ...Option) (output O, err error)\n\tTransform(ctx context.Context, input *schema.StreamReader[I], opts ...Option) (output *schema.StreamReader[O], err error)\n}\n\ntype invoke func(ctx context.Context, input any, opts ...any) (output any, err error)\ntype transform func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error)\n\n// composableRunnable the wrapper for all executable object directly provided by the user.\n// one instance corresponds to one instance of the executable object.\n// all information comes from executable object without any other dimensions of information.\n// for the graphNode, ChainBranch, StatePreHandler, StatePostHandler etc.\ntype composableRunnable struct {\n\ti invoke\n\tt transform\n\n\tinputType  reflect.Type\n\toutputType reflect.Type\n\toptionType reflect.Type\n\n\t*genericHelper\n\n\tisPassthrough bool\n\n\tmeta *executorMeta\n\n\t// only available when in Graph node\n\t// if composableRunnable not in Graph node, this field would be nil\n\tnodeInfo *nodeInfo\n}\n\nfunc runnableLambda[I, O, TOption any](i Invoke[I, O, TOption], s Stream[I, O, TOption], c Collect[I, O, TOption],\n\tt Transform[I, O, TOption], enableCallback bool) *composableRunnable {\n\trp := newRunnablePacker(i, s, c, t, enableCallback)\n\n\treturn rp.toComposableRunnable()\n}\n\ntype runnablePacker[I, O, TOption any] struct {\n\ti Invoke[I, O, TOption]\n\ts Stream[I, O, TOption]\n\tc Collect[I, O, TOption]\n\tt Transform[I, O, TOption]\n}\n\nfunc (rp *runnablePacker[I, O, TOption]) wrapRunnableCtx(ctxWrapper func(ctx context.Context, opts ...TOption) context.Context) {\n\ti, s, c, t := rp.i, rp.s, rp.c, rp.t\n\trp.i = func(ctx context.Context, input I, opts ...TOption) (output O, err error) {\n\t\tctx = ctxWrapper(ctx, opts...)\n\t\treturn i(ctx, input, opts...)\n\t}\n\trp.s = func(ctx context.Context, input I, opts ...TOption) (output *schema.StreamReader[O], err error) {\n\t\tctx = ctxWrapper(ctx, opts...)\n\t\treturn s(ctx, input, opts...)\n\t}\n\trp.c = func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output O, err error) {\n\t\tctx = ctxWrapper(ctx, opts...)\n\t\treturn c(ctx, input, opts...)\n\t}\n\n\trp.t = func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output *schema.StreamReader[O], err error) {\n\t\tctx = ctxWrapper(ctx, opts...)\n\t\treturn t(ctx, input, opts...)\n\t}\n}\n\nfunc (rp *runnablePacker[I, O, TOption]) toComposableRunnable() *composableRunnable {\n\tinputType := generic.TypeOf[I]()\n\toutputType := generic.TypeOf[O]()\n\toptionType := generic.TypeOf[TOption]()\n\tc := &composableRunnable{\n\t\tgenericHelper: newGenericHelper[I, O](),\n\t\tinputType:     inputType,\n\t\toutputType:    outputType,\n\t\toptionType:    optionType,\n\t}\n\n\ti := func(ctx context.Context, input any, opts ...any) (output any, err error) {\n\t\tin, ok := input.(I)\n\t\tif !ok {\n\t\t\t// When a nil is passed as an 'any' type, its original type information is lost,\n\t\t\t// becoming an untyped nil. This would cause type assertions to fail.\n\t\t\t// So if the input is nil and the target type I is an interface, we need to explicitly create a nil of type I.\n\t\t\tif input == nil && reflect.TypeOf((*I)(nil)).Elem().Kind() == reflect.Interface {\n\t\t\t\tvar i I\n\t\t\t\tin = i\n\t\t\t} else {\n\t\t\t\tpanic(newUnexpectedInputTypeErr(inputType, reflect.TypeOf(input)))\n\t\t\t}\n\t\t}\n\n\t\ttos, err := convertOption[TOption](opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn rp.Invoke(ctx, in, tos...)\n\t}\n\n\tt := func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) {\n\t\tin, ok := unpackStreamReader[I](input)\n\t\tif !ok {\n\t\t\tpanic(newUnexpectedInputTypeErr(reflect.TypeOf(in), input.getType()))\n\t\t}\n\n\t\ttos, err := convertOption[TOption](opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tout, err := rp.Transform(ctx, in, tos...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn packStreamReader(out), nil\n\t}\n\n\tc.i = i\n\tc.t = t\n\n\treturn c\n}\n\n// Invoke works like `ping => pong`.\nfunc (rp *runnablePacker[I, O, TOption]) Invoke(ctx context.Context,\n\tinput I, opts ...TOption) (output O, err error) {\n\treturn rp.i(ctx, input, opts...)\n}\n\n// Stream works like `ping => stream output`.\nfunc (rp *runnablePacker[I, O, TOption]) Stream(ctx context.Context,\n\tinput I, opts ...TOption) (output *schema.StreamReader[O], err error) {\n\n\treturn rp.s(ctx, input, opts...)\n}\n\n// Collect works like `stream input => pong`.\nfunc (rp *runnablePacker[I, O, TOption]) Collect(ctx context.Context,\n\tinput *schema.StreamReader[I], opts ...TOption) (output O, err error) {\n\treturn rp.c(ctx, input, opts...)\n}\n\n// Transform works like `stream input => stream output`.\nfunc (rp *runnablePacker[I, O, TOption]) Transform(ctx context.Context,\n\tinput *schema.StreamReader[I], opts ...TOption) (output *schema.StreamReader[O], err error) {\n\treturn rp.t(ctx, input, opts...)\n}\n\nfunc defaultImplConcatStreamReader[T any](\n\tsr *schema.StreamReader[T]) (T, error) {\n\n\tc, err := concatStreamReader(sr)\n\tif err != nil {\n\t\tvar t T\n\t\treturn t, fmt.Errorf(\"concat stream reader fail: %w\", err)\n\t}\n\n\treturn c, nil\n}\n\nfunc invokeByStream[I, O, TOption any](s Stream[I, O, TOption]) Invoke[I, O, TOption] {\n\treturn func(ctx context.Context, input I, opts ...TOption) (output O, err error) {\n\t\tsr, err := s(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\treturn defaultImplConcatStreamReader(sr)\n\t}\n}\n\nfunc invokeByCollect[I, O, TOption any](c Collect[I, O, TOption]) Invoke[I, O, TOption] {\n\treturn func(ctx context.Context, input I, opts ...TOption) (output O, err error) {\n\t\tsr := schema.StreamReaderFromArray([]I{input})\n\n\t\treturn c(ctx, sr, opts...)\n\t}\n}\n\nfunc invokeByTransform[I, O, TOption any](t Transform[I, O, TOption]) Invoke[I, O, TOption] {\n\treturn func(ctx context.Context, input I, opts ...TOption) (output O, err error) {\n\t\tsrInput := schema.StreamReaderFromArray([]I{input})\n\n\t\tsrOutput, err := t(ctx, srInput, opts...)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\treturn defaultImplConcatStreamReader(srOutput)\n\t}\n}\n\nfunc streamByTransform[I, O, TOption any](t Transform[I, O, TOption]) Stream[I, O, TOption] {\n\treturn func(ctx context.Context, input I, opts ...TOption) (output *schema.StreamReader[O], err error) {\n\t\tsrInput := schema.StreamReaderFromArray([]I{input})\n\n\t\treturn t(ctx, srInput, opts...)\n\t}\n}\n\nfunc streamByInvoke[I, O, TOption any](i Invoke[I, O, TOption]) Stream[I, O, TOption] {\n\treturn func(ctx context.Context, input I, opts ...TOption) (output *schema.StreamReader[O], err error) {\n\t\tout, err := i(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn schema.StreamReaderFromArray([]O{out}), nil\n\t}\n}\n\nfunc streamByCollect[I, O, TOption any](c Collect[I, O, TOption]) Stream[I, O, TOption] {\n\treturn func(ctx context.Context, input I, opts ...TOption) (output *schema.StreamReader[O], err error) {\n\t\tsrInput := schema.StreamReaderFromArray([]I{input})\n\t\tout, err := c(ctx, srInput, opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn schema.StreamReaderFromArray([]O{out}), nil\n\t}\n}\n\nfunc collectByTransform[I, O, TOption any](t Transform[I, O, TOption]) Collect[I, O, TOption] {\n\treturn func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output O, err error) {\n\t\tsrOutput, err := t(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\treturn defaultImplConcatStreamReader(srOutput)\n\t}\n}\n\nfunc collectByInvoke[I, O, TOption any](i Invoke[I, O, TOption]) Collect[I, O, TOption] {\n\treturn func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output O, err error) {\n\t\tin, err := defaultImplConcatStreamReader(input)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\treturn i(ctx, in, opts...)\n\t}\n}\n\nfunc collectByStream[I, O, TOption any](s Stream[I, O, TOption]) Collect[I, O, TOption] {\n\treturn func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output O, err error) {\n\t\tin, err := defaultImplConcatStreamReader(input)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\tsrOutput, err := s(ctx, in, opts...)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\treturn defaultImplConcatStreamReader(srOutput)\n\t}\n}\n\nfunc transformByStream[I, O, TOption any](s Stream[I, O, TOption]) Transform[I, O, TOption] {\n\treturn func(ctx context.Context, input *schema.StreamReader[I],\n\t\topts ...TOption) (output *schema.StreamReader[O], err error) {\n\t\tin, err := defaultImplConcatStreamReader(input)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\treturn s(ctx, in, opts...)\n\t}\n}\n\nfunc transformByCollect[I, O, TOption any](c Collect[I, O, TOption]) Transform[I, O, TOption] {\n\treturn func(ctx context.Context, input *schema.StreamReader[I],\n\t\topts ...TOption) (output *schema.StreamReader[O], err error) {\n\t\tout, err := c(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\treturn schema.StreamReaderFromArray([]O{out}), nil\n\t}\n}\n\nfunc transformByInvoke[I, O, TOption any](i Invoke[I, O, TOption]) Transform[I, O, TOption] {\n\treturn func(ctx context.Context, input *schema.StreamReader[I],\n\t\topts ...TOption) (output *schema.StreamReader[O], err error) {\n\t\tin, err := defaultImplConcatStreamReader(input)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\tout, err := i(ctx, in, opts...)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\treturn schema.StreamReaderFromArray([]O{out}), nil\n\t}\n}\n\nfunc newRunnablePacker[I, O, TOption any](i Invoke[I, O, TOption], s Stream[I, O, TOption],\n\tc Collect[I, O, TOption], t Transform[I, O, TOption], enableCallback bool) *runnablePacker[I, O, TOption] {\n\n\tr := &runnablePacker[I, O, TOption]{}\n\n\tif enableCallback {\n\t\tif i != nil {\n\t\t\ti = invokeWithCallbacks(i)\n\t\t}\n\n\t\tif s != nil {\n\t\t\ts = streamWithCallbacks(s)\n\t\t}\n\n\t\tif c != nil {\n\t\t\tc = collectWithCallbacks(c)\n\t\t}\n\n\t\tif t != nil {\n\t\t\tt = transformWithCallbacks(t)\n\t\t}\n\t}\n\n\tif i != nil {\n\t\tr.i = i\n\t} else if s != nil {\n\t\tr.i = invokeByStream(s)\n\t} else if c != nil {\n\t\tr.i = invokeByCollect(c)\n\t} else {\n\t\tr.i = invokeByTransform(t)\n\t}\n\n\tif s != nil {\n\t\tr.s = s\n\t} else if t != nil {\n\t\tr.s = streamByTransform(t)\n\t} else if i != nil {\n\t\tr.s = streamByInvoke(i)\n\t} else {\n\t\tr.s = streamByCollect(c)\n\t}\n\n\tif c != nil {\n\t\tr.c = c\n\t} else if t != nil {\n\t\tr.c = collectByTransform(t)\n\t} else if i != nil {\n\t\tr.c = collectByInvoke(i)\n\t} else {\n\t\tr.c = collectByStream(s)\n\t}\n\n\tif t != nil {\n\t\tr.t = t\n\t} else if s != nil {\n\t\tr.t = transformByStream(s)\n\t} else if c != nil {\n\t\tr.t = transformByCollect(c)\n\t} else {\n\t\tr.t = transformByInvoke(i)\n\t}\n\n\treturn r\n}\n\nfunc toGenericRunnable[I, O any](cr *composableRunnable, ctxWrapper func(ctx context.Context, opts ...Option) context.Context) (\n\t*runnablePacker[I, O, Option], error) {\n\ti := func(ctx context.Context, input I, opts ...Option) (output O, err error) {\n\t\tout, err := cr.i(ctx, input, toAnyList(opts)...)\n\t\tif err != nil {\n\t\t\treturn output, err\n\t\t}\n\n\t\tto, ok := out.(O)\n\t\tif !ok {\n\t\t\t// When a nil is passed as an 'any' type, its original type information is lost,\n\t\t\t// becoming an untyped nil. This would cause type assertions to fail.\n\t\t\t// So if the output is nil and the target type O is an interface, we need to explicitly create a nil of type O.\n\t\t\tif out == nil && generic.TypeOf[O]().Kind() == reflect.Interface {\n\t\t\t\tvar o O\n\t\t\t\tto = o\n\t\t\t} else {\n\t\t\t\tpanic(newUnexpectedInputTypeErr(generic.TypeOf[O](), reflect.TypeOf(out)))\n\t\t\t}\n\t\t}\n\t\treturn to, nil\n\t}\n\n\tt := func(ctx context.Context, input *schema.StreamReader[I],\n\t\topts ...Option) (output *schema.StreamReader[O], err error) {\n\t\tin := packStreamReader(input)\n\t\tout, err := cr.t(ctx, in, toAnyList(opts)...)\n\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\toutput, ok := unpackStreamReader[O](out)\n\t\tif !ok {\n\t\t\tpanic(\"impossible\")\n\t\t}\n\n\t\treturn output, nil\n\t}\n\n\tr := newRunnablePacker(i, nil, nil, t, false)\n\tr.wrapRunnableCtx(ctxWrapper)\n\n\treturn r, nil\n}\n\nfunc inputKeyedComposableRunnable(key string, r *composableRunnable) *composableRunnable {\n\twrapper := *r\n\twrapper.genericHelper = wrapper.genericHelper.forMapInput()\n\ti := r.i\n\twrapper.i = func(ctx context.Context, input any, opts ...any) (output any, err error) {\n\t\tv, ok := input.(map[string]any)[key]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"cannot find input key: %s\", key)\n\t\t}\n\t\tout, err := i(ctx, v, opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn out, nil\n\t}\n\n\tt := r.t\n\twrapper.t = func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) {\n\t\tnInput, ok := r.inputStreamFilter(key, input)\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"inputStreamFilter failed, key= %s, node name= %s, err= %w\", key, r.nodeInfo.name, err)\n\t\t}\n\t\tout, err := t(ctx, nInput, opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn out, nil\n\t}\n\n\twrapper.inputType = generic.TypeOf[map[string]any]()\n\treturn &wrapper\n}\n\nfunc outputKeyedComposableRunnable(key string, r *composableRunnable) *composableRunnable {\n\twrapper := *r\n\twrapper.genericHelper = wrapper.genericHelper.forMapOutput()\n\ti := r.i\n\twrapper.i = func(ctx context.Context, input any, opts ...any) (output any, err error) {\n\t\tout, err := i(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn map[string]any{key: out}, nil\n\t}\n\n\tt := r.t\n\twrapper.t = func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) {\n\t\tout, err := t(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn out.withKey(key), nil\n\t}\n\n\twrapper.outputType = generic.TypeOf[map[string]any]()\n\n\treturn &wrapper\n}\n\n// composablePassthrough special runnable that passthrough input to output\nfunc composablePassthrough() *composableRunnable {\n\tr := &composableRunnable{isPassthrough: true, nodeInfo: &nodeInfo{}}\n\n\tr.i = func(ctx context.Context, input any, opts ...any) (output any, err error) {\n\t\treturn input, nil\n\t}\n\n\tr.t = func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) {\n\t\treturn input, nil\n\t}\n\n\tr.meta = &executorMeta{\n\t\tcomponent:                  ComponentOfPassthrough,\n\t\tisComponentCallbackEnabled: false,\n\t\tcomponentImplType:          \"Passthrough\",\n\t}\n\n\treturn r\n}\n"
  },
  {
    "path": "compose/runnable_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"strconv\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestRunnableLambda(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"invoke_to_runnable\", func(t *testing.T) {\n\t\trl := runnableLambda(\n\t\t\tfunc(ctx context.Context, input int, opts ...Option) (output string, err error) {\n\t\t\t\treturn strconv.Itoa(input) + \"+\" + opts[0].options[0].(string), nil\n\t\t\t},\n\t\t\tnil, nil, nil, false)\n\n\t\tctxWrapper := func(ctx context.Context, opts ...Option) context.Context {\n\t\t\treturn ctx\n\t\t}\n\t\tgr, err := toGenericRunnable[int, string](rl, ctxWrapper)\n\t\tassert.NoError(t, err)\n\t\tout, err := gr.Invoke(ctx, 10, WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsr, err := gr.Stream(ctx, 10, WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tout, err = concatStreamReader(sr)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsri, swi := schema.Pipe[int](1)\n\t\t_ = swi.Send(10, nil)\n\t\tswi.Close()\n\t\tsriArr := sri.Copy(2)\n\n\t\tout, err = gr.Collect(ctx, sriArr[0], WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsr, err = gr.Transform(ctx, sriArr[1], WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tout, err = concatStreamReader(sr)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\t})\n\n\tt.Run(\"stream_to_runnable\", func(t *testing.T) {\n\t\trl := runnableLambda(nil,\n\t\t\tfunc(ctx context.Context, input int, opts ...Option) (output *schema.StreamReader[string], err error) {\n\t\t\t\tsro, swo := schema.Pipe[string](3)\n\t\t\t\t_ = swo.Send(strconv.Itoa(input), nil)\n\t\t\t\t_ = swo.Send(\"+\", nil)\n\t\t\t\t_ = swo.Send(opts[0].options[0].(string), nil)\n\t\t\t\tswo.Close()\n\t\t\t\treturn sro, nil\n\t\t\t}, nil, nil, false)\n\n\t\tctxWrapper := func(ctx context.Context, opts ...Option) context.Context {\n\t\t\treturn ctx\n\t\t}\n\t\tgr, err := toGenericRunnable[int, string](rl, ctxWrapper)\n\t\tassert.NoError(t, err)\n\t\tout, err := gr.Invoke(ctx, 10, WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsr, err := gr.Stream(ctx, 10, WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tout, err = concatStreamReader(sr)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsri, swi := schema.Pipe[int](1)\n\t\t_ = swi.Send(10, nil)\n\t\tswi.Close()\n\t\tsriArr := sri.Copy(2)\n\n\t\tout, err = gr.Collect(ctx, sriArr[0], WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsr, err = gr.Transform(ctx, sriArr[1], WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tout, err = concatStreamReader(sr)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\t})\n\n\tt.Run(\"transform_to_runnable\", func(t *testing.T) {\n\t\trl := runnableLambda(\n\t\t\tnil, nil, nil,\n\t\t\tfunc(ctx context.Context, input *schema.StreamReader[int], opts ...Option) (output *schema.StreamReader[string], err error) {\n\n\t\t\t\tin, e := input.Recv()\n\t\t\t\tif errors.Is(e, io.EOF) {\n\t\t\t\t\treturn nil, fmt.Errorf(\"unpected EOF\")\n\t\t\t\t}\n\t\t\t\tinput.Close()\n\n\t\t\t\tsro, swo := schema.Pipe[string](3)\n\t\t\t\t_ = swo.Send(strconv.Itoa(in), nil)\n\t\t\t\t_ = swo.Send(\"+\", nil)\n\t\t\t\t_ = swo.Send(opts[0].options[0].(string), nil)\n\t\t\t\tswo.Close()\n\t\t\t\treturn sro, nil\n\t\t\t},\n\t\t\tfalse)\n\n\t\tctxWrapper := func(ctx context.Context, opts ...Option) context.Context {\n\t\t\treturn ctx\n\t\t}\n\t\tgr, err := toGenericRunnable[int, string](rl, ctxWrapper)\n\t\tassert.NoError(t, err)\n\t\tout, err := gr.Invoke(ctx, 10, WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsr, err := gr.Stream(ctx, 10, WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tout, err = concatStreamReader(sr)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsri, swi := schema.Pipe[int](1)\n\t\t_ = swi.Send(10, nil)\n\t\tswi.Close()\n\t\tsriArr := sri.Copy(2)\n\n\t\tout, err = gr.Collect(ctx, sriArr[0], WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsr, err = gr.Transform(ctx, sriArr[1], WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tout, err = concatStreamReader(sr)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\t})\n\n\tt.Run(\"collect_to_runnable\", func(t *testing.T) {\n\t\trl := runnableLambda(nil, nil,\n\t\t\tfunc(ctx context.Context, input *schema.StreamReader[int], opts ...Option) (output string, err error) {\n\t\t\t\tin, e := input.Recv()\n\t\t\t\tif errors.Is(e, io.EOF) {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"unpected EOF\")\n\t\t\t\t}\n\t\t\t\tinput.Close()\n\n\t\t\t\treturn strconv.Itoa(in) + \"+\" + opts[0].options[0].(string), nil\n\t\t\t},\n\t\t\tnil, false)\n\n\t\tctxWrapper := func(ctx context.Context, opts ...Option) context.Context {\n\t\t\treturn ctx\n\t\t}\n\n\t\tgr, err := toGenericRunnable[int, string](rl, ctxWrapper)\n\t\tassert.NoError(t, err)\n\t\tout, err := gr.Invoke(ctx, 10, WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsr, err := gr.Stream(ctx, 10, WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tout, err = concatStreamReader(sr)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsri, swi := schema.Pipe[int](1)\n\t\t_ = swi.Send(10, nil)\n\t\tswi.Close()\n\t\tsriArr := sri.Copy(2)\n\n\t\tout, err = gr.Collect(ctx, sriArr[0], WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\n\t\tsr, err = gr.Transform(ctx, sriArr[1], WithLambdaOption(\"100\"))\n\t\tassert.NoError(t, err)\n\t\tout, err = concatStreamReader(sr)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"10+100\", out)\n\t})\n}\n"
  },
  {
    "path": "compose/state.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"reflect\"\n\t\"sync\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// GenLocalState is a function that generates the state.\ntype GenLocalState[S any] func(ctx context.Context) (state S)\n\ntype stateKey struct{}\n\ntype internalState struct {\n\tstate  any\n\tmu     sync.Mutex\n\tparent *internalState\n}\n\n// StatePreHandler is a function called before the node is executed.\n// Notice: if user called Stream but with StatePreHandler, the StatePreHandler will read all stream chunks and merge them into a single object.\ntype StatePreHandler[I, S any] func(ctx context.Context, in I, state S) (I, error)\n\n// StatePostHandler is a function called after the node is executed.\n// Notice: if user called Stream but with StatePostHandler, the StatePostHandler will read all stream chunks and merge them into a single object.\ntype StatePostHandler[O, S any] func(ctx context.Context, out O, state S) (O, error)\n\n// StreamStatePreHandler is a function that is called before the node is executed with stream input and output.\ntype StreamStatePreHandler[I, S any] func(ctx context.Context, in *schema.StreamReader[I], state S) (*schema.StreamReader[I], error)\n\n// StreamStatePostHandler is a function that is called after the node is executed with stream input and output.\ntype StreamStatePostHandler[O, S any] func(ctx context.Context, out *schema.StreamReader[O], state S) (*schema.StreamReader[O], error)\n\nfunc convertPreHandler[I, S any](handler StatePreHandler[I, S]) *composableRunnable {\n\trf := func(ctx context.Context, in I, opts ...any) (I, error) {\n\t\tcState, pMu, err := getState[S](ctx)\n\t\tif err != nil {\n\t\t\treturn in, err\n\t\t}\n\t\tpMu.Lock()\n\t\tdefer pMu.Unlock()\n\n\t\treturn handler(ctx, in, cState)\n\t}\n\n\treturn runnableLambda[I, I](rf, nil, nil, nil, false)\n}\n\nfunc convertPostHandler[O, S any](handler StatePostHandler[O, S]) *composableRunnable {\n\trf := func(ctx context.Context, out O, opts ...any) (O, error) {\n\t\tcState, pMu, err := getState[S](ctx)\n\t\tif err != nil {\n\t\t\treturn out, err\n\t\t}\n\t\tpMu.Lock()\n\t\tdefer pMu.Unlock()\n\n\t\treturn handler(ctx, out, cState)\n\t}\n\n\treturn runnableLambda[O, O](rf, nil, nil, nil, false)\n}\n\nfunc streamConvertPreHandler[I, S any](handler StreamStatePreHandler[I, S]) *composableRunnable {\n\trf := func(ctx context.Context, in *schema.StreamReader[I], opts ...any) (*schema.StreamReader[I], error) {\n\t\tcState, pMu, err := getState[S](ctx)\n\t\tif err != nil {\n\t\t\treturn in, err\n\t\t}\n\t\tpMu.Lock()\n\t\tdefer pMu.Unlock()\n\n\t\treturn handler(ctx, in, cState)\n\t}\n\n\treturn runnableLambda[I, I](nil, nil, nil, rf, false)\n}\n\nfunc streamConvertPostHandler[O, S any](handler StreamStatePostHandler[O, S]) *composableRunnable {\n\trf := func(ctx context.Context, out *schema.StreamReader[O], opts ...any) (*schema.StreamReader[O], error) {\n\t\tcState, pMu, err := getState[S](ctx)\n\t\tif err != nil {\n\t\t\treturn out, err\n\t\t}\n\t\tpMu.Lock()\n\t\tdefer pMu.Unlock()\n\n\t\treturn handler(ctx, out, cState)\n\t}\n\n\treturn runnableLambda[O, O](nil, nil, nil, rf, false)\n}\n\n// ProcessState processes the state from the context in a concurrency-safe way.\n// This is the recommended way to access and modify state in custom nodes.\n// The provided function handler will be executed with exclusive access to the state (protected by mutex).\n//\n// State Lookup Behavior:\n// - If the requested state type exists in the current graph, it will be returned\n// - If not found in current graph, ProcessState will search in parent graph states (for nested graphs)\n// - This enables nested graphs to access state from their parent graphs\n// - Follows lexical scoping: inner state of the same type shadows outer state\n//\n// Concurrency Safety:\n// - ProcessState automatically locks the mutex of the state being accessed (current or parent level)\n// - Each state level has its own mutex, allowing concurrent access to different levels\n// - The lock is held for the entire duration of the handler function\n//\n// Note: This method will report an error if the state type doesn't match or state is not found in the context chain.\n//\n// Example - Basic usage in a single graph:\n//\n//\tlambdaFunc := func(ctx context.Context, in string, opts ...any) (string, error) {\n//\t\terr := compose.ProcessState[*MyState](ctx, func(ctx context.Context, state *MyState) error {\n//\t\t\t// Safely modify state\n//\t\t\tstate.Count++\n//\t\t\treturn nil\n//\t\t})\n//\t\tif err != nil {\n//\t\t\treturn \"\", err\n//\t\t}\n//\t\treturn in, nil\n//\t}\n//\n// Example - Nested graph accessing parent state:\n//\n//\t// In an inner graph node\n//\tinnerNode := func(ctx context.Context, input string) (string, error) {\n//\t\t// Access parent graph's state\n//\t\terr := compose.ProcessState[*OuterState](ctx, func(ctx context.Context, s *OuterState) error {\n//\t\t\ts.Counter++  // Safely modify parent state\n//\t\t\treturn nil\n//\t\t})\n//\t\tif err != nil {\n//\t\t\treturn \"\", err\n//\t\t}\n//\n//\t\t// Also access inner graph's own state\n//\t\terr = compose.ProcessState[*InnerState](ctx, func(ctx context.Context, s *InnerState) error {\n//\t\t\ts.Data = \"processed\"\n//\t\t\treturn nil\n//\t\t})\n//\t\treturn input, nil\n//\t}\nfunc ProcessState[S any](ctx context.Context, handler func(context.Context, S) error) error {\n\ts, pMu, err := getState[S](ctx)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"get state from context fail: %w\", err)\n\t}\n\tpMu.Lock()\n\tdefer pMu.Unlock()\n\treturn handler(ctx, s)\n}\n\nfunc getState[S any](ctx context.Context) (S, *sync.Mutex, error) {\n\tstate := ctx.Value(stateKey{})\n\n\tif state == nil {\n\t\tvar s S\n\t\treturn s, nil, fmt.Errorf(\"have not set state\")\n\t}\n\n\tinterState := state.(*internalState)\n\n\tfor interState != nil {\n\t\tif cState, ok := interState.state.(S); ok {\n\t\t\treturn cState, &interState.mu, nil\n\t\t}\n\t\tinterState = interState.parent\n\t}\n\n\tvar s S\n\treturn s, nil, fmt.Errorf(\"cannot find state with type: %v in states chain, \"+\n\t\t\"current state type: %v\",\n\t\tgeneric.TypeOf[S](), reflect.TypeOf(state.(*internalState).state))\n}\n"
  },
  {
    "path": "compose/state_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"io\"\n\t\"strings\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n\t\"unicode/utf8\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype midStr string\n\nfunc TestStateGraphWithEdge(t *testing.T) {\n\n\tctx := context.Background()\n\n\tconst (\n\t\tnodeOfL1 = \"invokable\"\n\t\tnodeOfL2 = \"streamable\"\n\t\tnodeOfL3 = \"transformable\"\n\t)\n\n\ttype testState struct {\n\t\tms []string\n\t}\n\n\tgen := func(ctx context.Context) *testState {\n\t\treturn &testState{}\n\t}\n\n\tsg := NewGraph[string, string](WithGenLocalState(gen))\n\n\tl1 := InvokableLambda(func(ctx context.Context, in string) (out midStr, err error) {\n\t\treturn midStr(\"InvokableLambda: \" + in), nil\n\t})\n\n\tl1StateToInput := func(ctx context.Context, in string, state *testState) (string, error) {\n\t\tstate.ms = append(state.ms, in)\n\t\treturn in, nil\n\t}\n\n\tl1StateToOutput := func(ctx context.Context, out midStr, state *testState) (midStr, error) {\n\t\tstate.ms = append(state.ms, string(out))\n\t\treturn out, nil\n\t}\n\n\terr := sg.AddLambdaNode(nodeOfL1, l1,\n\t\tWithStatePreHandler(l1StateToInput), WithStatePostHandler(l1StateToOutput))\n\tassert.NoError(t, err)\n\n\tl2 := StreamableLambda(func(ctx context.Context, input midStr) (output *schema.StreamReader[string], err error) {\n\t\toutStr := \"StreamableLambda: \" + string(input)\n\n\t\tsr, sw := schema.Pipe[string](utf8.RuneCountInString(outStr))\n\n\t\tgo func() {\n\t\t\tfor _, field := range strings.Fields(outStr) {\n\t\t\t\tsw.Send(field+\" \", nil)\n\t\t\t}\n\t\t\tsw.Close()\n\t\t}()\n\n\t\treturn sr, nil\n\t})\n\n\tl2StateToOutput := func(ctx context.Context, out string, state *testState) (string, error) {\n\t\tstate.ms = append(state.ms, out)\n\t\treturn out, nil\n\t}\n\n\terr = sg.AddLambdaNode(nodeOfL2, l2, WithStatePostHandler(l2StateToOutput))\n\tassert.NoError(t, err)\n\n\tl3 := TransformableLambda(func(ctx context.Context, input *schema.StreamReader[string]) (\n\t\toutput *schema.StreamReader[string], err error) {\n\n\t\tprefix := \"TransformableLambda: \"\n\t\tsr, sw := schema.Pipe[string](20)\n\n\t\tgo func() {\n\t\t\tfor _, field := range strings.Fields(prefix) {\n\t\t\t\tsw.Send(field+\" \", nil)\n\t\t\t}\n\t\t\tdefer input.Close()\n\n\t\t\tfor {\n\t\t\t\tchunk, err := input.Recv()\n\t\t\t\tif err != nil {\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t\t// TODO: how to trace this kind of error in the goroutine of processing stream\n\t\t\t\t\tsw.Send(chunk, err)\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tsw.Send(chunk, nil)\n\n\t\t\t}\n\t\t\tsw.Close()\n\t\t}()\n\n\t\treturn sr, nil\n\t})\n\n\tl3StateToOutput := func(ctx context.Context, out string, state *testState) (string, error) {\n\t\tstate.ms = append(state.ms, out)\n\t\tassert.Len(t, state.ms, 4)\n\t\treturn out, nil\n\t}\n\n\terr = sg.AddLambdaNode(nodeOfL3, l3, WithStatePostHandler(l3StateToOutput))\n\tassert.NoError(t, err)\n\n\terr = sg.AddEdge(START, nodeOfL1)\n\tassert.NoError(t, err)\n\n\terr = sg.AddEdge(nodeOfL1, nodeOfL2)\n\tassert.NoError(t, err)\n\n\terr = sg.AddEdge(nodeOfL2, nodeOfL3)\n\tassert.NoError(t, err)\n\n\terr = sg.AddEdge(nodeOfL3, END)\n\tassert.NoError(t, err)\n\n\trun, err := sg.Compile(ctx)\n\tassert.NoError(t, err)\n\n\tout, err := run.Invoke(ctx, \"how are you\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"TransformableLambda: StreamableLambda: InvokableLambda: how are you \", out)\n\n\tstream, err := run.Stream(ctx, \"how are you\")\n\tassert.NoError(t, err)\n\tout, err = concatStreamReader(stream)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"TransformableLambda: StreamableLambda: InvokableLambda: how are you \", out)\n\n\tsr, sw := schema.Pipe[string](1)\n\tsw.Send(\"how are you\", nil)\n\tsw.Close()\n\n\tstream, err = run.Transform(ctx, sr)\n\tassert.NoError(t, err)\n\tout, err = concatStreamReader(stream)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"TransformableLambda: StreamableLambda: InvokableLambda: how are you \", out)\n}\n\nfunc TestStateGraphUtils(t *testing.T) {\n\tt.Run(\"getState_success\", func(t *testing.T) {\n\t\ttype testStruct struct {\n\t\t\tUserID int64\n\t\t}\n\n\t\tctx := context.Background()\n\n\t\tctx = context.WithValue(ctx, stateKey{}, &internalState{\n\t\t\tstate: &testStruct{UserID: 10},\n\t\t})\n\n\t\tvar userID int64\n\t\terr := ProcessState[*testStruct](ctx, func(_ context.Context, state *testStruct) error {\n\t\t\tuserID = state.UserID\n\t\t\treturn nil\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, int64(10), userID)\n\t})\n\n\tt.Run(\"getState_nil\", func(t *testing.T) {\n\t\ttype testStruct struct {\n\t\t\tUserID int64\n\t\t}\n\n\t\tctx := context.Background()\n\t\tctx = context.WithValue(ctx, stateKey{}, &internalState{})\n\n\t\terr := ProcessState[*testStruct](ctx, func(_ context.Context, state *testStruct) error {\n\t\t\treturn nil\n\t\t})\n\t\tassert.ErrorContains(t, err, \"cannot find state with type: *compose.testStruct in states chain, \"+\n\t\t\t\"current state type: <nil>\")\n\t})\n\n\tt.Run(\"getState_type_error\", func(t *testing.T) {\n\t\ttype testStruct struct {\n\t\t\tUserID int64\n\t\t}\n\n\t\tctx := context.Background()\n\t\tctx = context.WithValue(ctx, stateKey{}, &internalState{\n\t\t\tstate: &testStruct{UserID: 10},\n\t\t})\n\n\t\terr := ProcessState[string](ctx, func(_ context.Context, state string) error {\n\t\t\treturn nil\n\t\t})\n\t\tassert.ErrorContains(t, err, \"cannot find state with type: string in states chain, \"+\n\t\t\t\"current state type: *compose.testStruct\")\n\n\t})\n}\n\nfunc TestStateChain(t *testing.T) {\n\tctx := context.Background()\n\ttype testState struct {\n\t\tField1 string\n\t\tField2 string\n\t}\n\tsc := NewChain[string, string](WithGenLocalState(func(ctx context.Context) (state *testState) {\n\t\treturn &testState{}\n\t}))\n\n\tr, err := sc.AppendLambda(InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\terr = ProcessState[*testState](ctx, func(_ context.Context, state *testState) error {\n\t\t\tstate.Field1 = \"node1\"\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn input, nil\n\t}), WithStatePostHandler(func(ctx context.Context, out string, state *testState) (string, error) {\n\t\tstate.Field2 = \"node2\"\n\t\treturn out, nil\n\t})).\n\t\tAppendLambda(InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\treturn input, nil\n\t\t}), WithStatePreHandler(func(ctx context.Context, in string, state *testState) (string, error) {\n\t\t\treturn in + state.Field1 + state.Field2, nil\n\t\t})).Compile(ctx)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tresult, err := r.Invoke(ctx, \"start\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif result != \"startnode1node2\" {\n\t\tt.Fatal(\"result is unexpected\")\n\t}\n}\n\nfunc TestStreamState(t *testing.T) {\n\ttype testState struct {\n\t\tField1 string\n\t}\n\tctx := context.Background()\n\ts := &testState{Field1: \"1\"}\n\tg := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testState) { return s }))\n\terr := g.AddLambdaNode(\"1\", TransformableLambda(func(ctx context.Context, input *schema.StreamReader[string]) (output *schema.StreamReader[string], err error) {\n\t\treturn input, nil\n\t}), WithStreamStatePreHandler(func(ctx context.Context, in *schema.StreamReader[string], state *testState) (*schema.StreamReader[string], error) {\n\t\tsr, sw := schema.Pipe[string](5)\n\t\tfor i := 0; i < 5; i++ {\n\t\t\tsw.Send(state.Field1, nil)\n\t\t}\n\t\tsw.Close()\n\t\treturn sr, nil\n\t}), WithStreamStatePostHandler(func(ctx context.Context, in *schema.StreamReader[string], state *testState) (*schema.StreamReader[string], error) {\n\t\tss := in.Copy(2)\n\t\tfor {\n\t\t\tchunk, err := ss[0].Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\treturn ss[1], nil\n\t\t\t}\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tstate.Field1 += chunk\n\t\t}\n\t}))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(START, \"1\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\terr = g.AddEdge(\"1\", END)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tr, err := g.Compile(ctx)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tsr, _ := schema.Pipe[string](1)\n\tstreamResult, err := r.Transform(ctx, sr)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif s.Field1 != \"111111\" {\n\t\tt.Fatal(\"state is unexpected\")\n\t}\n\tfor i := 0; i < 5; i++ {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tif chunk != \"1\" {\n\t\t\tt.Fatal(\"result is unexpected\")\n\t\t}\n\t}\n\t_, err = streamResult.Recv()\n\tif err != io.EOF {\n\t\tt.Fatal(\"result is unexpected\")\n\t}\n}\n\n// Nested Graph State Tests\n\ntype NestedOuterState struct {\n\tValue   string\n\tCounter int\n}\n\ntype NestedInnerState struct {\n\tValue string\n}\n\nfunc init() {\n\tschema.RegisterName[*NestedOuterState](\"NestedOuterState\")\n\tschema.RegisterName[*NestedInnerState](\"NestedInnerState\")\n}\n\nfunc TestNestedGraphStateAccess(t *testing.T) {\n\t// Test that inner graph can access outer graph's state\n\tgenOuterState := func(ctx context.Context) *NestedOuterState {\n\t\treturn &NestedOuterState{Value: \"outer\", Counter: 0}\n\t}\n\n\tgenInnerState := func(ctx context.Context) *NestedInnerState {\n\t\treturn &NestedInnerState{Value: \"inner\"}\n\t}\n\n\tinnerNode := func(ctx context.Context, input string) (string, error) {\n\t\t// Access both inner and outer state\n\t\tvar outerValue string\n\t\terr := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error {\n\t\t\touterValue = s.Value\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\tvar innerValue string\n\t\terr = ProcessState(ctx, func(ctx context.Context, s *NestedInnerState) error {\n\t\t\tinnerValue = s.Value\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn fmt.Sprintf(\"%s_inner=%s_outer=%s\", input, innerValue, outerValue), nil\n\t}\n\n\tinnerGraph := NewGraph[string, string](WithGenLocalState(genInnerState))\n\t_ = innerGraph.AddLambdaNode(\"inner_node\", InvokableLambda(innerNode))\n\t_ = innerGraph.AddEdge(START, \"inner_node\")\n\t_ = innerGraph.AddEdge(\"inner_node\", END)\n\n\touterGraph := NewGraph[string, string](WithGenLocalState(genOuterState))\n\t_ = outerGraph.AddGraphNode(\"inner_graph\", innerGraph)\n\t_ = outerGraph.AddEdge(START, \"inner_graph\")\n\t_ = outerGraph.AddEdge(\"inner_graph\", END)\n\n\tr, err := outerGraph.Compile(context.Background())\n\tassert.NoError(t, err)\n\n\tout, err := r.Invoke(context.Background(), \"start\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"start_inner=inner_outer=outer\", out)\n}\n\nfunc TestNestedGraphStateShadowing(t *testing.T) {\n\t// Test that inner state shadows outer state of the same type (lexical scoping)\n\ttype CommonState struct {\n\t\tValue string\n\t}\n\n\tgenOuterState := func(ctx context.Context) *CommonState {\n\t\treturn &CommonState{Value: \"outer\"}\n\t}\n\n\tgenInnerState := func(ctx context.Context) *CommonState {\n\t\treturn &CommonState{Value: \"inner\"}\n\t}\n\n\tinnerNode := func(ctx context.Context, input string) (string, error) {\n\t\tvar value string\n\t\terr := ProcessState(ctx, func(ctx context.Context, s *CommonState) error {\n\t\t\t// Should see \"inner\" because inner state shadows outer state\n\t\t\tvalue = s.Value\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn input + \"_\" + value, nil\n\t}\n\n\tinnerGraph := NewGraph[string, string](WithGenLocalState(genInnerState))\n\t_ = innerGraph.AddLambdaNode(\"inner_node\", InvokableLambda(innerNode))\n\t_ = innerGraph.AddEdge(START, \"inner_node\")\n\t_ = innerGraph.AddEdge(\"inner_node\", END)\n\n\touterGraph := NewGraph[string, string](WithGenLocalState(genOuterState))\n\t_ = outerGraph.AddGraphNode(\"inner_graph\", innerGraph)\n\t_ = outerGraph.AddEdge(START, \"inner_graph\")\n\t_ = outerGraph.AddEdge(\"inner_graph\", END)\n\n\tr, err := outerGraph.Compile(context.Background())\n\tassert.NoError(t, err)\n\n\tout, err := r.Invoke(context.Background(), \"start\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"start_inner\", out)\n}\n\nfunc TestNestedGraphStateAfterResume(t *testing.T) {\n\t// Test that state parent linking works correctly after resume\n\t// when the outer state is restored from checkpoint (new instance)\n\tgenOuterState := func(ctx context.Context) *NestedOuterState {\n\t\treturn &NestedOuterState{Value: \"outer\", Counter: 0}\n\t}\n\n\tgenInnerState := func(ctx context.Context) *NestedInnerState {\n\t\treturn &NestedInnerState{Value: \"inner\"}\n\t}\n\n\t// Node that modifies outer state\n\touterNode := func(ctx context.Context, input string) (string, error) {\n\t\terr := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error {\n\t\t\ts.Counter = 42\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn input, nil\n\t}\n\n\t// Inner node that reads outer state\n\tinnerNode := func(ctx context.Context, input string) (string, error) {\n\t\tvar outerCounter int\n\t\tvar outerValue string\n\t\terr := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error {\n\t\t\t// Should see the modified counter value from the restored state\n\t\t\touterCounter = s.Counter\n\t\t\touterValue = s.Value\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn fmt.Sprintf(\"%s_counter=%d_value=%s\", input, outerCounter, outerValue), nil\n\t}\n\n\tinnerGraph := NewGraph[string, string](WithGenLocalState(genInnerState))\n\t_ = innerGraph.AddLambdaNode(\"inner_node\", InvokableLambda(innerNode))\n\t_ = innerGraph.AddEdge(START, \"inner_node\")\n\t_ = innerGraph.AddEdge(\"inner_node\", END)\n\n\touterGraph := NewGraph[string, string](WithGenLocalState(genOuterState))\n\t_ = outerGraph.AddLambdaNode(\"outer_node\", InvokableLambda(outerNode))\n\t_ = outerGraph.AddGraphNode(\"inner_graph\", innerGraph, WithGraphCompileOptions(WithInterruptBeforeNodes([]string{\"inner_node\"})))\n\t_ = outerGraph.AddEdge(START, \"outer_node\")\n\t_ = outerGraph.AddEdge(\"outer_node\", \"inner_graph\")\n\t_ = outerGraph.AddEdge(\"inner_graph\", END)\n\n\tstore := newInMemoryStore()\n\tr, err := outerGraph.Compile(context.Background(), WithCheckPointStore(store))\n\tassert.NoError(t, err)\n\n\t// First run - should interrupt after modifying outer state\n\t_, err = r.Invoke(context.Background(), \"start\", WithCheckPointID(\"state_resume_test\"))\n\tassert.Error(t, err)\n\n\t// Resume - outer state should be restored with Counter=42\n\t// Inner graph should link to this restored outer state\n\tout, err := r.Invoke(context.Background(), \"start\", WithCheckPointID(\"state_resume_test\"))\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"start_counter=42_value=outer\", out)\n}\n\nfunc TestLambdaNestedGraphStateAccess(t *testing.T) {\n\t// Test that inner graph invoked from a lambda can access outer graph's state\n\t// This tests the case: outer graph -> lambda node -> inner graph (using CompositeInterrupt)\n\tgenOuterState := func(ctx context.Context) *NestedOuterState {\n\t\treturn &NestedOuterState{Value: \"outer\", Counter: 100}\n\t}\n\n\tgenInnerState := func(ctx context.Context) *NestedInnerState {\n\t\treturn &NestedInnerState{Value: \"inner\"}\n\t}\n\n\t// Inner node that accesses outer state\n\tinnerNode := func(ctx context.Context, input string) (string, error) {\n\t\tvar outerValue string\n\t\tvar outerCounter int\n\t\terr := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error {\n\t\t\touterValue = s.Value\n\t\t\touterCounter = s.Counter\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\tvar innerValue string\n\t\terr = ProcessState(ctx, func(ctx context.Context, s *NestedInnerState) error {\n\t\t\tinnerValue = s.Value\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn fmt.Sprintf(\"%s_inner=%s_outer=%s_%d\", input, innerValue, outerValue, outerCounter), nil\n\t}\n\n\t// Build inner graph\n\tinnerGraph := NewGraph[string, string](WithGenLocalState(genInnerState))\n\t_ = innerGraph.AddLambdaNode(\"inner_node\", InvokableLambda(innerNode))\n\t_ = innerGraph.AddEdge(START, \"inner_node\")\n\t_ = innerGraph.AddEdge(\"inner_node\", END)\n\n\t// Compile inner graph as a standalone runnable\n\tinnerRunnable, err := innerGraph.Compile(context.Background())\n\tassert.NoError(t, err)\n\n\t// Lambda that invokes the inner graph\n\tlambdaNode := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\t// Simply invoke the inner graph - state context is passed through\n\t\treturn innerRunnable.Invoke(ctx, input)\n\t})\n\n\t// Build outer graph\n\touterGraph := NewGraph[string, string](WithGenLocalState(genOuterState))\n\t_ = outerGraph.AddLambdaNode(\"lambda_with_graph\", lambdaNode)\n\t_ = outerGraph.AddEdge(START, \"lambda_with_graph\")\n\t_ = outerGraph.AddEdge(\"lambda_with_graph\", END)\n\n\tr, err := outerGraph.Compile(context.Background())\n\tassert.NoError(t, err)\n\n\tout, err := r.Invoke(context.Background(), \"start\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"start_inner=inner_outer=outer_100\", out)\n}\n\nfunc TestLambdaNestedGraphStateAfterResume(t *testing.T) {\n\t// Test that state parent linking works correctly after resume\n\t// in the lambda-nested case (outer graph -> lambda -> inner graph)\n\tgenOuterState := func(ctx context.Context) *NestedOuterState {\n\t\treturn &NestedOuterState{Value: \"outer\", Counter: 0}\n\t}\n\n\tgenInnerState := func(ctx context.Context) *NestedInnerState {\n\t\treturn &NestedInnerState{Value: \"inner\"}\n\t}\n\n\t// Outer node that modifies state\n\touterNode := func(ctx context.Context, input string) (string, error) {\n\t\terr := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error {\n\t\t\ts.Counter = 99\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn input, nil\n\t}\n\n\t// Inner lambda that interrupts on first run, reads outer state on resume\n\tinnerLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\twasInterrupted, _, _ := GetInterruptState[*NestedInnerState](ctx)\n\t\tif !wasInterrupted {\n\t\t\t// First run: interrupt\n\t\t\treturn \"\", StatefulInterrupt(ctx, \"inner interrupt\", &NestedInnerState{Value: \"inner\"})\n\t\t}\n\n\t\t// Resumed: read outer state\n\t\tvar outerCounter int\n\t\tvar outerValue string\n\t\terr := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error {\n\t\t\t// Should see the modified counter from the restored state\n\t\t\touterCounter = s.Counter\n\t\t\touterValue = s.Value\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn fmt.Sprintf(\"%s_counter=%d_value=%s\", input, outerCounter, outerValue), nil\n\t})\n\n\t// Build inner graph\n\tinnerGraph := NewGraph[string, string](WithGenLocalState(genInnerState))\n\t_ = innerGraph.AddLambdaNode(\"inner_lambda\", innerLambda)\n\t_ = innerGraph.AddEdge(START, \"inner_lambda\")\n\t_ = innerGraph.AddEdge(\"inner_lambda\", END)\n\n\t// Compile inner graph as standalone runnable with checkpoint support\n\tinnerRunnable, err := innerGraph.Compile(context.Background(),\n\t\tWithGraphName(\"inner\"),\n\t\tWithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t// Composite lambda that invokes the inner graph and handles interrupts\n\tcompositeLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) {\n\t\toutput, err := innerRunnable.Invoke(ctx, input, WithCheckPointID(\"inner-cp\"))\n\t\tif err != nil {\n\t\t\t_, isInterrupt := ExtractInterruptInfo(err)\n\t\t\tif !isInterrupt {\n\t\t\t\treturn \"\", err\n\t\t\t}\n\t\t\t// Wrap the interrupt using CompositeInterrupt\n\t\t\treturn \"\", CompositeInterrupt(ctx, \"composite interrupt\", nil, err)\n\t\t}\n\t\treturn output, nil\n\t})\n\n\t// Build outer graph\n\touterGraph := NewGraph[string, string](WithGenLocalState(genOuterState))\n\t_ = outerGraph.AddLambdaNode(\"outer_node\", InvokableLambda(outerNode))\n\t_ = outerGraph.AddLambdaNode(\"composite_lambda\", compositeLambda)\n\t_ = outerGraph.AddEdge(START, \"outer_node\")\n\t_ = outerGraph.AddEdge(\"outer_node\", \"composite_lambda\")\n\t_ = outerGraph.AddEdge(\"composite_lambda\", END)\n\n\t// Compile outer graph\n\touterRunnable, err := outerGraph.Compile(context.Background(),\n\t\tWithGraphName(\"root\"),\n\t\tWithCheckPointStore(newInMemoryStore()))\n\tassert.NoError(t, err)\n\n\t// First run - should interrupt after modifying outer state\n\tcheckPointID := \"lambda_state_resume_test\"\n\t_, err = outerRunnable.Invoke(context.Background(), \"start\", WithCheckPointID(checkPointID))\n\tassert.Error(t, err)\n\n\tinterruptInfo, isInterrupt := ExtractInterruptInfo(err)\n\tassert.True(t, isInterrupt)\n\n\t// Resume - outer state should be restored with Counter=99\n\t// Inner lambda should link to this restored outer state\n\tctx := ResumeWithData(context.Background(), interruptInfo.InterruptContexts[0].ID, nil)\n\tout, err := outerRunnable.Invoke(ctx, \"start\", WithCheckPointID(checkPointID))\n\tassert.NoError(t, err)\n\n\t// Verify the inner lambda saw the modified counter from the restored outer state\n\tassert.Contains(t, out, \"counter=99\")\n\tassert.Contains(t, out, \"value=outer\")\n}\n\nfunc TestNestedGraphStateConcurrency(t *testing.T) {\n\t// Test that concurrent access to parent and child states uses correct locks\n\t// This verifies that ProcessState properly locks the parent state's mutex when accessing it\n\tgenOuterState := func(ctx context.Context) *NestedOuterState {\n\t\treturn &NestedOuterState{Value: \"outer\", Counter: 0}\n\t}\n\n\tgenInnerState := func(ctx context.Context) *NestedInnerState {\n\t\treturn &NestedInnerState{Value: \"inner\"}\n\t}\n\n\t// Inner node that concurrently modifies both outer and inner state\n\tinnerNode := func(ctx context.Context, input string) (string, error) {\n\t\tvar wg sync.WaitGroup\n\t\terrors := make(chan error, 20)\n\n\t\t// Launch 10 goroutines that modify outer state\n\t\t// If locks don't work correctly, we'll see race conditions\n\t\tfor i := 0; i < 10; i++ {\n\t\t\twg.Add(1)\n\t\t\tgo func() {\n\t\t\t\tdefer wg.Done()\n\t\t\t\terr := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error {\n\t\t\t\t\t// ProcessState should hold the parent's lock during this entire function\n\t\t\t\t\tcurrent := s.Counter\n\t\t\t\t\ttime.Sleep(time.Millisecond) // Simulate work\n\t\t\t\t\ts.Counter = current + 1\n\t\t\t\t\treturn nil\n\t\t\t\t})\n\t\t\t\tif err != nil {\n\t\t\t\t\terrors <- err\n\t\t\t\t}\n\t\t\t}()\n\t\t}\n\n\t\t// Launch 10 goroutines that modify inner state\n\t\tfor i := 0; i < 10; i++ {\n\t\t\twg.Add(1)\n\t\t\tgo func() {\n\t\t\t\tdefer wg.Done()\n\t\t\t\terr := ProcessState(ctx, func(ctx context.Context, s *NestedInnerState) error {\n\t\t\t\t\t// This uses the inner state's own lock\n\t\t\t\t\treturn nil\n\t\t\t\t})\n\t\t\t\tif err != nil {\n\t\t\t\t\terrors <- err\n\t\t\t\t}\n\t\t\t}()\n\t\t}\n\n\t\twg.Wait()\n\t\tclose(errors)\n\n\t\t// Check for errors\n\t\tfor err := range errors {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn input, nil\n\t}\n\n\tinnerGraph := NewGraph[string, string](WithGenLocalState(genInnerState))\n\t_ = innerGraph.AddLambdaNode(\"inner_node\", InvokableLambda(innerNode))\n\t_ = innerGraph.AddEdge(START, \"inner_node\")\n\t_ = innerGraph.AddEdge(\"inner_node\", END)\n\n\touterGraph := NewGraph[string, string](WithGenLocalState(genOuterState))\n\t_ = outerGraph.AddGraphNode(\"inner_graph\", innerGraph)\n\t_ = outerGraph.AddEdge(START, \"inner_graph\")\n\t_ = outerGraph.AddEdge(\"inner_graph\", END)\n\n\tr, err := outerGraph.Compile(context.Background())\n\tassert.NoError(t, err)\n\n\t_, err = r.Invoke(context.Background(), \"start\")\n\tassert.NoError(t, err)\n\n\t// Note: This test is primarily validated by running with -race flag\n\t// If locks don't work correctly, the race detector will catch it\n}\n"
  },
  {
    "path": "compose/stream_concat.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"errors\"\n\t\"io\"\n\n\t\"github.com/cloudwego/eino/internal\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// RegisterStreamChunkConcatFunc registers a function to concat stream chunks.\n// It's required when you want to concat stream chunks of a specific type.\n// for example you call Invoke() but node only implements Stream().\n// call at process init\n// not thread safe\n// eg.\n//\n//\ttype testStruct struct {\n//\t\tfield1 string\n//\t\tfield2 int\n//\t}\n//\tcompose.RegisterStreamChunkConcatFunc(func(items []testStruct) (testStruct, error) {\n//\t\treturn testStruct{\n//\t\t\tfield1: items[1].field1, // may implement inplace logic by your scenario\n//\t\t\tfield2: items[0].field2 + items[1].field2,\n//\t\t}, nil\n//\t})\nfunc RegisterStreamChunkConcatFunc[T any](fn func([]T) (T, error)) {\n\tinternal.RegisterStreamChunkConcatFunc(fn)\n}\n\nvar emptyStreamConcatErr = errors.New(\"stream reader is empty, concat fail\")\n\nfunc concatStreamReader[T any](sr *schema.StreamReader[T]) (T, error) {\n\tdefer sr.Close()\n\n\tvar items []T\n\n\tfor {\n\t\tchunk, err := sr.Recv()\n\t\tif err != nil {\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tif _, ok := schema.GetSourceName(err); ok {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tvar t T\n\t\t\treturn t, newStreamReadError(err)\n\t\t}\n\n\t\titems = append(items, chunk)\n\t}\n\n\tif len(items) == 0 {\n\t\tvar t T\n\t\treturn t, emptyStreamConcatErr\n\t}\n\n\tif len(items) == 1 {\n\t\treturn items[0], nil\n\t}\n\n\tres, err := internal.ConcatItems(items)\n\tif err != nil {\n\t\tvar t T\n\t\treturn t, err\n\t}\n\treturn res, nil\n}\n"
  },
  {
    "path": "compose/stream_concat_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"errors\"\n\t\"strconv\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/internal\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype tStreamConcatItemForTest struct {\n\ts string\n}\n\nfunc concatTStreamForTest(items []tStreamConcatItemForTest) (tStreamConcatItemForTest, error) {\n\tvar s string\n\tfor _, item := range items {\n\t\ts += item.s\n\t}\n\n\treturn tStreamConcatItemForTest{s: s}, nil\n}\n\nfunc concatIntForTest(items []int) (int, error) {\n\tvar i int\n\tfor _, item := range items {\n\t\ti += item\n\t}\n\n\treturn i, nil\n}\n\ntype tConcatErrForTest struct{}\n\nfunc concatTStreamError(_ []tConcatErrForTest) (tConcatErrForTest, error) {\n\treturn tConcatErrForTest{}, errors.New(\"test error\")\n}\n\nfunc TestConcatRegistry(t *testing.T) {\n\n\tRegisterStreamChunkConcatFunc(concatTStreamForTest)\n\n\tsr, sw := schema.Pipe[tStreamConcatItemForTest](10)\n\tgo func() {\n\t\tfor i := 0; i < 10; i++ {\n\t\t\tsw.Send(tStreamConcatItemForTest{s: strconv.Itoa(i)}, nil)\n\t\t}\n\t\tsw.Close()\n\t}()\n\n\tlastVal, err := concatStreamReader(sr)\n\tassert.Nil(t, err)\n\n\tassert.Equal(t, \"0123456789\", lastVal.s)\n}\n\nfunc TestStringConcat(t *testing.T) {\n\tsr, sw := schema.Pipe[string](10)\n\tgo func() {\n\t\tfor i := 0; i < 10; i++ {\n\t\t\tsw.Send(strconv.Itoa(i), nil)\n\t\t}\n\n\t\tsw.Close()\n\t}()\n\n\tlastVal, err := concatStreamReader(sr)\n\tassert.Nil(t, err)\n\n\tassert.Equal(t, \"0123456789\", lastVal)\n}\n\nfunc TestMessageConcat(t *testing.T) {\n\tsr, sw := schema.Pipe[*schema.Message](10)\n\tgo func() {\n\t\tfor i := 0; i < 10; i++ {\n\t\t\tcontent := schema.UserMessage(strconv.Itoa(i))\n\t\t\tif i%4 == 0 {\n\t\t\t\tcontent.Extra = map[string]any{\n\t\t\t\t\t\"key_1\":         strconv.Itoa(i),\n\t\t\t\t\tstrconv.Itoa(i): strconv.Itoa(i),\n\t\t\t\t}\n\t\t\t}\n\t\t\tsw.Send(content, nil)\n\t\t}\n\t\tsw.Close()\n\t}()\n\n\tlastVal, err := concatStreamReader(sr)\n\tassert.Nil(t, err)\n\tassert.Equal(t, \"0123456789\", lastVal.Content)\n\tassert.Len(t, lastVal.Extra, 4)\n\tassert.Equal(t, map[string]any{\n\t\t\"key_1\": \"048\",\n\t\t\"0\":     \"0\",\n\t\t\"4\":     \"4\",\n\t\t\"8\":     \"8\",\n\t}, lastVal.Extra)\n\n}\n\nfunc TestMapConcat(t *testing.T) {\n\tRegisterStreamChunkConcatFunc(concatTStreamForTest)\n\tRegisterStreamChunkConcatFunc(concatIntForTest)\n\n\tt.Run(\"simple map\", func(t *testing.T) {\n\t\tsr, sw := schema.Pipe[map[string]any](10)\n\n\t\tgo func() {\n\t\t\tfor i := 0; i < 10; i++ {\n\t\t\t\tsw.Send(map[string]any{\n\t\t\t\t\t\"string\":        strconv.Itoa(i),\n\t\t\t\t\t\"custom_concat\": tStreamConcatItemForTest{s: strconv.Itoa(9 - i)},\n\t\t\t\t\t\"count\":         i,\n\t\t\t\t}, nil)\n\t\t\t}\n\t\t\tsw.Close()\n\t\t}()\n\n\t\tlastVal, err := concatStreamReader(sr)\n\t\tassert.Nil(t, err)\n\n\t\tassert.Equal(t, \"0123456789\", lastVal[\"string\"])\n\t\tassert.Equal(t, \"9876543210\", lastVal[\"custom_concat\"].(tStreamConcatItemForTest).s)\n\t\tassert.Equal(t, 45, lastVal[\"count\"])\n\n\t})\n\n\tt.Run(\"complex map\", func(t *testing.T) {\n\t\tsr, sw := schema.Pipe[map[string]any](10)\n\n\t\tgo func() {\n\t\t\tfor i := 0; i < 10; i++ {\n\t\t\t\t// 嵌套 map, 仅允许第一层做类型合并，第二层直接覆盖\n\t\t\t\tsw.Send(map[string]any{ // 嵌套 map\n\t\t\t\t\t\"string\": strconv.Itoa(i),\n\t\t\t\t\t\"deep_map\": map[string]any{\n\t\t\t\t\t\t\"message\": &schema.Message{\n\t\t\t\t\t\t\tContent: strconv.Itoa(i),\n\t\t\t\t\t\t},\n\t\t\t\t\t\t\"custom_concat_deep\": tStreamConcatItemForTest{s: strconv.Itoa(9 - i)},\n\t\t\t\t\t\t\"count\":              i,\n\t\t\t\t\t},\n\t\t\t\t\t\"custom_concat\": tStreamConcatItemForTest{s: strconv.Itoa(9 - i)},\n\t\t\t\t\t\"count\":         i,\n\t\t\t\t}, nil)\n\t\t\t}\n\t\t\tsw.Close()\n\t\t}()\n\n\t\tlastVal, err := concatStreamReader(sr)\n\t\tassert.Nil(t, err)\n\n\t\tassert.Equal(t, \"0123456789\", lastVal[\"string\"])\n\t\tassert.Equal(t, 45, lastVal[\"count\"])\n\t\tassert.Equal(t, \"0123456789\", lastVal[\"deep_map\"].(map[string]any)[\"message\"].(*schema.Message).Content)\n\t\tassert.Equal(t, \"9876543210\", lastVal[\"deep_map\"].(map[string]any)[\"custom_concat_deep\"].(tStreamConcatItemForTest).s)\n\t\tassert.Equal(t, 45, lastVal[\"deep_map\"].(map[string]any)[\"count\"])\n\t})\n}\n\nfunc TestConcatError(t *testing.T) {\n\tt.Run(\"map type not equal\", func(t *testing.T) {\n\t\ta := map[string]any{\n\t\t\t\"str\": \"string_01\",\n\t\t\t\"x\":   \"string_in_a\",\n\t\t}\n\n\t\tb := map[string]any{\n\t\t\t\"str\": \"string_02\",\n\t\t\t\"x\":   123,\n\t\t}\n\t\t_, err := internal.ConcatItems([]map[string]any{a, b})\n\t\tassert.NotNil(t, err)\n\t})\n\n\tt.Run(\"merge error\", func(t *testing.T) {\n\t\tRegisterStreamChunkConcatFunc(concatTStreamError)\n\n\t\t_, err := internal.ConcatItems([]tConcatErrForTest{{}, {}})\n\t\tassert.NotNil(t, err)\n\t})\n}\n\nfunc TestConcatSliceValue(t *testing.T) {\n\ttype testStruct struct {\n\t\tA string\n\t}\n\n\ts := []testStruct{{}, {A: \"123\"}, {}}\n\tresult, err := internal.ConcatItems(s)\n\tassert.Nil(t, err)\n\tassert.Equal(t, testStruct{A: \"123\"}, result)\n\n\ts = []testStruct{{}, {}, {}}\n\tresult, err = internal.ConcatItems(s)\n\tassert.Nil(t, err)\n\tassert.Equal(t, testStruct{}, result)\n}\n"
  },
  {
    "path": "compose/stream_reader.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype streamReader interface {\n\tcopy(n int) []streamReader\n\tgetType() reflect.Type\n\tgetChunkType() reflect.Type\n\tmerge([]streamReader) streamReader\n\twithKey(string) streamReader\n\tclose()\n\ttoAnyStreamReader() *schema.StreamReader[any]\n\tmergeWithNames([]streamReader, []string) streamReader\n}\n\ntype streamReaderPacker[T any] struct {\n\tsr *schema.StreamReader[T]\n}\n\nfunc (srp streamReaderPacker[T]) close() {\n\tsrp.sr.Close()\n}\n\nfunc (srp streamReaderPacker[T]) copy(n int) []streamReader {\n\tret := make([]streamReader, n)\n\tsrs := srp.sr.Copy(n)\n\n\tfor i := 0; i < n; i++ {\n\t\tret[i] = streamReaderPacker[T]{srs[i]}\n\t}\n\n\treturn ret\n}\n\nfunc (srp streamReaderPacker[T]) getType() reflect.Type {\n\treturn reflect.TypeOf(srp.sr)\n}\n\nfunc (srp streamReaderPacker[T]) getChunkType() reflect.Type {\n\treturn generic.TypeOf[T]()\n}\n\nfunc (srp streamReaderPacker[T]) toStreamReaders(srs []streamReader) []*schema.StreamReader[T] {\n\tret := make([]*schema.StreamReader[T], len(srs)+1)\n\tret[0] = srp.sr\n\tfor i := 1; i < len(ret); i++ {\n\t\tsr, ok := unpackStreamReader[T](srs[i-1])\n\t\tif !ok {\n\t\t\treturn nil\n\t\t}\n\n\t\tret[i] = sr\n\t}\n\n\treturn ret\n}\n\nfunc (srp streamReaderPacker[T]) merge(isrs []streamReader) streamReader {\n\tsrs := srp.toStreamReaders(isrs)\n\n\tsr := schema.MergeStreamReaders(srs)\n\n\treturn packStreamReader(sr)\n}\n\nfunc (srp streamReaderPacker[T]) mergeWithNames(isrs []streamReader, names []string) streamReader {\n\tsrs := srp.toStreamReaders(isrs)\n\n\tsr := schema.InternalMergeNamedStreamReaders(srs, names)\n\n\treturn packStreamReader(sr)\n}\n\nfunc (srp streamReaderPacker[T]) withKey(key string) streamReader {\n\tcvt := func(v T) (map[string]any, error) {\n\t\treturn map[string]any{key: v}, nil\n\t}\n\n\tret := schema.StreamReaderWithConvert[T, map[string]any](srp.sr, cvt)\n\n\treturn packStreamReader(ret)\n}\n\nfunc (srp streamReaderPacker[T]) toAnyStreamReader() *schema.StreamReader[any] {\n\treturn schema.StreamReaderWithConvert(srp.sr, func(t T) (any, error) {\n\t\treturn t, nil\n\t})\n}\n\nfunc packStreamReader[T any](sr *schema.StreamReader[T]) streamReader {\n\treturn streamReaderPacker[T]{sr}\n}\n\nfunc unpackStreamReader[T any](isr streamReader) (*schema.StreamReader[T], bool) {\n\tc, ok := isr.(streamReaderPacker[T])\n\tif ok {\n\t\treturn c.sr, true\n\t}\n\n\ttyp := generic.TypeOf[T]()\n\tif typ.Kind() == reflect.Interface {\n\t\treturn schema.StreamReaderWithConvert(isr.toAnyStreamReader(), func(t any) (T, error) {\n\t\t\treturn t.(T), nil\n\t\t}), true\n\t}\n\n\treturn nil, false\n}\n"
  },
  {
    "path": "compose/stream_reader_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"io\"\n\t\"reflect\"\n\t\"testing\"\n\n\t\"github.com/cloudwego/eino/schema\"\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestArrayStreamMerge(t *testing.T) {\n\n\tt.Run(\"unpack_to_equal_type\", func(t *testing.T) {\n\t\ta1 := []int{1, 2, 3}\n\t\ta2 := []int{4, 5, 6}\n\t\ta3 := []int{7, 8, 9}\n\t\ts1 := schema.StreamReaderFromArray(a1)\n\t\ts2 := schema.StreamReaderFromArray(a2)\n\t\ts3 := schema.StreamReaderFromArray(a3)\n\n\t\tsp1 := streamReaderPacker[int]{sr: s1}\n\t\tsp2 := streamReaderPacker[int]{sr: s2}\n\t\tsp3 := streamReaderPacker[int]{sr: s3}\n\n\t\tsp := sp1.merge([]streamReader{sp2, sp3})\n\n\t\tsr, ok := unpackStreamReader[int](sp)\n\t\tif !ok {\n\t\t\tt.Fatal(\"unexpected\")\n\t\t}\n\n\t\tdefer sr.Close()\n\n\t\tvar result []int\n\t\tfor {\n\t\t\tchunk, err := sr.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.Nil(t, err)\n\t\t\tresult = append(result, chunk)\n\t\t}\n\t\tif !reflect.DeepEqual(result, append(append(a1, a2...), a3...)) {\n\t\t\tt.Fatalf(\"result: %v error\", result)\n\t\t}\n\t})\n\n\tt.Run(\"unpack_to_father_type\", func(t *testing.T) {\n\t\ta1 := []*doctor{{say: \"a\"}, {say: \"b\"}, {say: \"c\"}}\n\t\ta2 := []*doctor{{say: \"d\"}, {say: \"e\"}, {say: \"f\"}}\n\t\ta3 := []*doctor{{say: \"g\"}, {say: \"h\"}, {say: \"i\"}}\n\t\ts1 := schema.StreamReaderFromArray(a1)\n\t\ts2 := schema.StreamReaderFromArray(a2)\n\t\ts3 := schema.StreamReaderFromArray(a3)\n\n\t\tsp1 := streamReaderPacker[*doctor]{sr: s1}\n\t\tsp2 := streamReaderPacker[*doctor]{sr: s2}\n\t\tsp3 := streamReaderPacker[*doctor]{sr: s3}\n\n\t\tsp := sp1.merge([]streamReader{sp2, sp3})\n\n\t\tsr, ok := unpackStreamReader[person](sp)\n\t\tassert.True(t, ok)\n\n\t\tdefer sr.Close()\n\n\t\tvar result []person\n\t\tfor {\n\t\t\tchunk, err := sr.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.Nil(t, err)\n\t\t\tresult = append(result, chunk)\n\t\t}\n\n\t\tbaseline := append(append(a1, a2...), a3...)\n\n\t\tassert.Len(t, result, len(baseline))\n\n\t\tfor idx := range result {\n\t\t\tassert.Equal(t, baseline[idx].say, result[idx].Say())\n\t\t}\n\t})\n}\n"
  },
  {
    "path": "compose/tool_node.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"runtime/debug\"\n\t\"sync\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/internal/safe\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype toolsNodeOptions struct {\n\tToolOptions []tool.Option\n\tToolList    []tool.BaseTool\n}\n\n// ToolsNodeOption is the option func type for ToolsNode.\ntype ToolsNodeOption func(o *toolsNodeOptions)\n\n// WithToolOption adds tool options to the ToolsNode.\nfunc WithToolOption(opts ...tool.Option) ToolsNodeOption {\n\treturn func(o *toolsNodeOptions) {\n\t\to.ToolOptions = append(o.ToolOptions, opts...)\n\t}\n}\n\n// WithToolList sets the tool list for the ToolsNode.\nfunc WithToolList(tool ...tool.BaseTool) ToolsNodeOption {\n\treturn func(o *toolsNodeOptions) {\n\t\to.ToolList = tool\n\t}\n}\n\n// ToolsNode represents a node capable of executing tools within a graph.\n// The Graph Node interface is defined as follows:\n//\n//\tInvoke(ctx context.Context, input *schema.Message, opts ...ToolsNodeOption) ([]*schema.Message, error)\n//\tStream(ctx context.Context, input *schema.Message, opts ...ToolsNodeOption) (*schema.StreamReader[[]*schema.Message], error)\n//\n// Input: An AssistantMessage containing ToolCalls\n// Output: An array of ToolMessage where the order of elements corresponds to the order of ToolCalls in the input\ntype ToolsNode struct {\n\ttuple                             *toolsTuple\n\tunknownToolHandler                func(ctx context.Context, name, input string) (string, error)\n\texecuteSequentially               bool\n\ttoolArgumentsHandler              func(ctx context.Context, name, input string) (string, error)\n\ttoolCallMiddlewares               []InvokableToolMiddleware\n\tstreamToolCallMiddlewares         []StreamableToolMiddleware\n\tenhancedToolCallMiddlewares       []EnhancedInvokableToolMiddleware\n\tenhancedStreamToolCallMiddlewares []EnhancedStreamableToolMiddleware\n}\n\n// ToolInput represents the input parameters for a tool call execution.\ntype ToolInput struct {\n\t// Name is the name of the tool to be executed.\n\tName string\n\t// Arguments contains the arguments for the tool call.\n\tArguments string\n\t// CallID is the unique identifier for this tool call.\n\tCallID string\n\t// CallOptions contains tool options for the execution.\n\tCallOptions []tool.Option\n}\n\n// ToolOutput represents the result of a non-streaming tool call execution.\ntype ToolOutput struct {\n\t// Result contains the string output from the tool execution.\n\tResult string\n}\n\n// StreamToolOutput represents the result of a streaming tool call execution.\ntype StreamToolOutput struct {\n\t// Result is a stream reader that provides access to the tool's streaming output.\n\tResult *schema.StreamReader[string]\n}\n\n// EnhancedInvokableToolOutput represents the result of a non-streaming enhanced tool call execution.\n// It supports returning structured multimodal content (text, images, audio, video, files) from tools.\ntype EnhancedInvokableToolOutput struct {\n\t// Result contains the structured multimodal output from the tool execution.\n\tResult *schema.ToolResult\n}\n\n// EnhancedStreamableToolOutput represents the result of a streaming enhanced tool call execution.\n// It provides a stream reader for accessing multimodal content progressively.\ntype EnhancedStreamableToolOutput struct {\n\t// Result is a stream reader that provides access to the tool's streaming multimodal output.\n\tResult *schema.StreamReader[*schema.ToolResult]\n}\n\n// InvokableToolEndpoint is the function signature for non-streaming tool calls.\ntype InvokableToolEndpoint func(ctx context.Context, input *ToolInput) (*ToolOutput, error)\n\n// StreamableToolEndpoint is the function signature for streaming tool calls.\ntype StreamableToolEndpoint func(ctx context.Context, input *ToolInput) (*StreamToolOutput, error)\n\ntype EnhancedInvokableToolEndpoint func(ctx context.Context, input *ToolInput) (*EnhancedInvokableToolOutput, error)\n\ntype EnhancedStreamableToolEndpoint func(ctx context.Context, input *ToolInput) (*EnhancedStreamableToolOutput, error)\n\n// InvokableToolMiddleware is a function that wraps InvokableToolEndpoint to add custom processing logic.\n// It can be used to intercept, modify, or enhance tool call execution for non-streaming tools.\ntype InvokableToolMiddleware func(InvokableToolEndpoint) InvokableToolEndpoint\n\n// StreamableToolMiddleware is a function that wraps StreamableToolEndpoint to add custom processing logic.\n// It can be used to intercept, modify, or enhance tool call execution for streaming tools.\ntype StreamableToolMiddleware func(StreamableToolEndpoint) StreamableToolEndpoint\n\ntype EnhancedInvokableToolMiddleware func(EnhancedInvokableToolEndpoint) EnhancedInvokableToolEndpoint\n\ntype EnhancedStreamableToolMiddleware func(EnhancedStreamableToolEndpoint) EnhancedStreamableToolEndpoint\n\n// ToolMiddleware groups middleware hooks for invokable and streamable tool calls.\ntype ToolMiddleware struct {\n\t// Invokable contains middleware function for non-streaming tool calls.\n\t// Note: This middleware only applies to tools that implement the InvokableTool interface.\n\tInvokable InvokableToolMiddleware\n\n\t// Streamable contains middleware function for streaming tool calls.\n\t// Note: This middleware only applies to tools that implement the StreamableTool interface.\n\tStreamable StreamableToolMiddleware\n\n\t// EnhancedInvokable contains middleware function for non-streaming enhanced tool calls.\n\t// Note: This middleware only applies to tools that implement the EnhancedInvokableTool interface.\n\tEnhancedInvokable EnhancedInvokableToolMiddleware\n\n\t// EnhancedStreamable contains middleware function for streaming enhanced tool calls.\n\t// Note: This middleware only applies to tools that implement the EnhancedStreamableTool interface.\n\tEnhancedStreamable EnhancedStreamableToolMiddleware\n}\n\n// ToolsNodeConfig is the config for ToolsNode.\ntype ToolsNodeConfig struct {\n\t// Tools specify the list of tools can be called which are BaseTool but must implement InvokableTool or StreamableTool.\n\tTools []tool.BaseTool\n\n\t// UnknownToolsHandler handles tool calls for non-existent tools when LLM hallucinates.\n\t// This field is optional. When not set, calling a non-existent tool will result in an error.\n\t// When provided, if the LLM attempts to call a tool that doesn't exist in the Tools list,\n\t// this handler will be invoked instead of returning an error, allowing graceful handling of hallucinated tools.\n\t// Parameters:\n\t//   - ctx: The context for the tool call\n\t//   - name: The name of the non-existent tool\n\t//   - input: The tool call input generated by llm\n\t// Returns:\n\t//   - string: The response to be returned as if the tool was executed\n\t//   - error: Any error that occurred during handling\n\tUnknownToolsHandler func(ctx context.Context, name, input string) (string, error)\n\n\t// ExecuteSequentially determines whether tool calls should be executed sequentially (in order) or in parallel.\n\t// When set to true, tool calls will be executed one after another in the order they appear in the input message.\n\t// When set to false (default), tool calls will be executed in parallel.\n\tExecuteSequentially bool\n\n\t// ToolArgumentsHandler allows handling of tool arguments before execution.\n\t// When provided, this function will be called for each tool call to process the arguments.\n\t// Parameters:\n\t//   - ctx: The context for the tool call\n\t//   - name: The name of the tool being called\n\t//   - arguments: The original arguments string for the tool\n\t// Returns:\n\t//   - string: The processed arguments string to be used for tool execution\n\t//   - error: Any error that occurred during preprocessing\n\tToolArgumentsHandler func(ctx context.Context, name, arguments string) (string, error)\n\n\t// ToolCallMiddlewares configures middleware for tool calls.\n\t// Each element can contain Invokable and/or Streamable middleware.\n\t// Invokable middleware only applies to tools implementing InvokableTool interface.\n\t// Streamable middleware only applies to tools implementing StreamableTool interface.\n\tToolCallMiddlewares []ToolMiddleware\n}\n\n// NewToolNode creates a new ToolsNode.\n// e.g.\n//\n//\tconf := &ToolsNodeConfig{\n//\t\tTools: []tool.BaseTool{invokableTool1, streamableTool2},\n//\t}\n//\ttoolsNode, err := NewToolNode(ctx, conf)\nfunc NewToolNode(ctx context.Context, conf *ToolsNodeConfig) (*ToolsNode, error) {\n\tvar middlewares []InvokableToolMiddleware\n\tvar streamMiddlewares []StreamableToolMiddleware\n\tvar enhancedInvokableMiddlewares []EnhancedInvokableToolMiddleware\n\tvar enhancedStreamableMiddlewares []EnhancedStreamableToolMiddleware\n\n\tfor _, m := range conf.ToolCallMiddlewares {\n\t\tif m.Invokable != nil {\n\t\t\tmiddlewares = append(middlewares, m.Invokable)\n\t\t}\n\t\tif m.Streamable != nil {\n\t\t\tstreamMiddlewares = append(streamMiddlewares, m.Streamable)\n\t\t}\n\t\tif m.EnhancedInvokable != nil {\n\t\t\tenhancedInvokableMiddlewares = append(enhancedInvokableMiddlewares, m.EnhancedInvokable)\n\t\t}\n\t\tif m.EnhancedStreamable != nil {\n\t\t\tenhancedStreamableMiddlewares = append(enhancedStreamableMiddlewares, m.EnhancedStreamable)\n\t\t}\n\t}\n\n\ttuple, err := convTools(ctx, conf.Tools, middlewares, streamMiddlewares, enhancedInvokableMiddlewares, enhancedStreamableMiddlewares)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &ToolsNode{\n\t\ttuple:                             tuple,\n\t\tunknownToolHandler:                conf.UnknownToolsHandler,\n\t\texecuteSequentially:               conf.ExecuteSequentially,\n\t\ttoolArgumentsHandler:              conf.ToolArgumentsHandler,\n\t\ttoolCallMiddlewares:               middlewares,\n\t\tstreamToolCallMiddlewares:         streamMiddlewares,\n\t\tenhancedToolCallMiddlewares:       enhancedInvokableMiddlewares,\n\t\tenhancedStreamToolCallMiddlewares: enhancedStreamableMiddlewares,\n\t}, nil\n}\n\n// ToolsInterruptAndRerunExtra carries interrupt metadata for ToolsNode reruns.\ntype ToolsInterruptAndRerunExtra struct {\n\t// ToolCalls contains all tool calls from the original assistant message.\n\tToolCalls []schema.ToolCall\n\n\t// ExecutedTools maps tool call IDs to their string output for successfully executed standard tools.\n\tExecutedTools map[string]string\n\n\t// ExecutedEnhancedTools maps tool call IDs to their structured multimodal output for successfully executed enhanced tools.\n\tExecutedEnhancedTools map[string]*schema.ToolResult\n\n\t// RerunTools contains the IDs of tool calls that need to be re-executed.\n\tRerunTools []string\n\n\t// RerunExtraMap stores additional metadata for each tool call that needs rerun, keyed by tool call ID.\n\tRerunExtraMap map[string]any\n}\n\nfunc init() {\n\tschema.RegisterName[*ToolsInterruptAndRerunExtra](\"_eino_compose_tools_interrupt_and_rerun_extra\")\n\tschema.RegisterName[*toolsInterruptAndRerunState](\"_eino_compose_tools_interrupt_and_rerun_state\")\n}\n\ntype toolsInterruptAndRerunState struct {\n\tInput                 *schema.Message\n\tExecutedTools         map[string]string\n\tExecutedEnhancedTools map[string]*schema.ToolResult\n\tRerunTools            []string\n}\n\ntype toolsTuple struct {\n\tindexes                     map[string]int\n\tmeta                        []*executorMeta\n\tendpoints                   []InvokableToolEndpoint\n\tstreamEndpoints             []StreamableToolEndpoint\n\tenhancedInvokableEndpoints  []EnhancedInvokableToolEndpoint\n\tenhancedStreamableEndpoints []EnhancedStreamableToolEndpoint\n}\n\nfunc convTools(ctx context.Context, tools []tool.BaseTool, ms []InvokableToolMiddleware, sms []StreamableToolMiddleware,\n\tems []EnhancedInvokableToolMiddleware, esms []EnhancedStreamableToolMiddleware) (*toolsTuple, error) {\n\tret := &toolsTuple{\n\t\tindexes:                     make(map[string]int),\n\t\tmeta:                        make([]*executorMeta, len(tools)),\n\t\tendpoints:                   make([]InvokableToolEndpoint, len(tools)),\n\t\tstreamEndpoints:             make([]StreamableToolEndpoint, len(tools)),\n\t\tenhancedInvokableEndpoints:  make([]EnhancedInvokableToolEndpoint, len(tools)),\n\t\tenhancedStreamableEndpoints: make([]EnhancedStreamableToolEndpoint, len(tools)),\n\t}\n\tfor idx, bt := range tools {\n\t\ttl, err := bt.Info(ctx)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"(NewToolNode) failed to get tool info at idx= %d: %w\", idx, err)\n\t\t}\n\n\t\ttoolName := tl.Name\n\t\tvar (\n\t\t\tst     tool.StreamableTool\n\t\t\tit     tool.InvokableTool\n\t\t\teiTool tool.EnhancedInvokableTool\n\t\t\tesTool tool.EnhancedStreamableTool\n\n\t\t\tinvokable          InvokableToolEndpoint\n\t\t\tstreamable         StreamableToolEndpoint\n\t\t\tenhancedInvokable  EnhancedInvokableToolEndpoint\n\t\t\tenhancedStreamable EnhancedStreamableToolEndpoint\n\n\t\t\tok   bool\n\t\t\tmeta *executorMeta\n\t\t)\n\n\t\tmeta = parseExecutorInfoFromComponent(components.ComponentOfTool, bt)\n\n\t\tif st, ok = bt.(tool.StreamableTool); ok {\n\t\t\tstreamable = wrapStreamToolCall(st, sms, !meta.isComponentCallbackEnabled)\n\t\t}\n\n\t\tif it, ok = bt.(tool.InvokableTool); ok {\n\t\t\tinvokable = wrapToolCall(it, ms, !meta.isComponentCallbackEnabled)\n\t\t}\n\n\t\tif eiTool, ok = bt.(tool.EnhancedInvokableTool); ok {\n\t\t\tenhancedInvokable = wrapEnhancedInvokableToolCall(eiTool, ems, !meta.isComponentCallbackEnabled)\n\t\t}\n\n\t\tif esTool, ok = bt.(tool.EnhancedStreamableTool); ok {\n\t\t\tenhancedStreamable = wrapEnhancedStreamableToolCall(esTool, esms, !meta.isComponentCallbackEnabled)\n\t\t}\n\n\t\tif st == nil && it == nil && eiTool == nil && esTool == nil {\n\t\t\treturn nil, fmt.Errorf(\"tool %s is not invokable, streamable, enhanced invokable or enhanced streamable\", toolName)\n\t\t}\n\t\tif streamable == nil && invokable != nil {\n\t\t\tstreamable = invokableToStreamable(invokable)\n\t\t}\n\t\tif invokable == nil && streamable != nil {\n\t\t\tinvokable = streamableToInvokable(streamable)\n\t\t}\n\n\t\tif enhancedStreamable == nil && enhancedInvokable != nil {\n\t\t\tenhancedStreamable = enhancedInvokableToEnhancedStreamable(enhancedInvokable)\n\t\t}\n\t\tif enhancedInvokable == nil && enhancedStreamable != nil {\n\t\t\tenhancedInvokable = enhancedStreamableToEnhancedInvokable(enhancedStreamable)\n\t\t}\n\n\t\tret.indexes[toolName] = idx\n\t\tret.meta[idx] = meta\n\t\tret.endpoints[idx] = invokable\n\t\tret.streamEndpoints[idx] = streamable\n\t\tret.enhancedInvokableEndpoints[idx] = enhancedInvokable\n\t\tret.enhancedStreamableEndpoints[idx] = enhancedStreamable\n\t}\n\treturn ret, nil\n}\n\nfunc wrapToolCall(it tool.InvokableTool, middlewares []InvokableToolMiddleware, needCallback bool) InvokableToolEndpoint {\n\tmiddleware := func(next InvokableToolEndpoint) InvokableToolEndpoint {\n\t\tfor i := len(middlewares) - 1; i >= 0; i-- {\n\t\t\tnext = middlewares[i](next)\n\t\t}\n\t\treturn next\n\t}\n\tif needCallback {\n\t\tit = &invokableToolWithCallback{it: it}\n\t}\n\treturn middleware(func(ctx context.Context, input *ToolInput) (*ToolOutput, error) {\n\t\tresult, err := it.InvokableRun(ctx, input.Arguments, input.CallOptions...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &ToolOutput{Result: result}, nil\n\t})\n}\n\nfunc wrapStreamToolCall(st tool.StreamableTool, middlewares []StreamableToolMiddleware, needCallback bool) StreamableToolEndpoint {\n\tmiddleware := func(next StreamableToolEndpoint) StreamableToolEndpoint {\n\t\tfor i := len(middlewares) - 1; i >= 0; i-- {\n\t\t\tnext = middlewares[i](next)\n\t\t}\n\t\treturn next\n\t}\n\tif needCallback {\n\t\tst = &streamableToolWithCallback{st: st}\n\t}\n\treturn middleware(func(ctx context.Context, input *ToolInput) (*StreamToolOutput, error) {\n\t\tresult, err := st.StreamableRun(ctx, input.Arguments, input.CallOptions...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &StreamToolOutput{Result: result}, nil\n\t})\n}\n\nfunc wrapEnhancedInvokableToolCall(eiTool tool.EnhancedInvokableTool, middlewares []EnhancedInvokableToolMiddleware, needCallback bool) EnhancedInvokableToolEndpoint {\n\tmiddleware := func(next EnhancedInvokableToolEndpoint) EnhancedInvokableToolEndpoint {\n\t\tfor i := len(middlewares) - 1; i >= 0; i-- {\n\t\t\tnext = middlewares[i](next)\n\t\t}\n\t\treturn next\n\t}\n\tif needCallback {\n\t\teiTool = &enhancedInvokableToolWithCallback{eiTool: eiTool}\n\t}\n\treturn middleware(func(ctx context.Context, input *ToolInput) (*EnhancedInvokableToolOutput, error) {\n\t\tresult, err := eiTool.InvokableRun(ctx, &schema.ToolArgument{Text: input.Arguments}, input.CallOptions...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &EnhancedInvokableToolOutput{Result: result}, nil\n\t})\n}\n\nfunc wrapEnhancedStreamableToolCall(est tool.EnhancedStreamableTool, middlewares []EnhancedStreamableToolMiddleware, needCallback bool) EnhancedStreamableToolEndpoint {\n\tmiddleware := func(next EnhancedStreamableToolEndpoint) EnhancedStreamableToolEndpoint {\n\t\tfor i := len(middlewares) - 1; i >= 0; i-- {\n\t\t\tnext = middlewares[i](next)\n\t\t}\n\t\treturn next\n\t}\n\tif needCallback {\n\t\test = &enhancedStreamableToolWithCallback{est: est}\n\t}\n\treturn middleware(func(ctx context.Context, input *ToolInput) (*EnhancedStreamableToolOutput, error) {\n\t\tresult, err := est.StreamableRun(ctx, &schema.ToolArgument{Text: input.Arguments}, input.CallOptions...)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &EnhancedStreamableToolOutput{Result: result}, nil\n\t})\n}\n\ntype invokableToolWithCallback struct {\n\tit tool.InvokableTool\n}\n\nfunc (i *invokableToolWithCallback) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn i.it.Info(ctx)\n}\n\nfunc (i *invokableToolWithCallback) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\treturn invokeWithCallbacks(i.it.InvokableRun)(ctx, argumentsInJSON, opts...)\n}\n\ntype streamableToolWithCallback struct {\n\tst tool.StreamableTool\n}\n\nfunc (s *streamableToolWithCallback) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn s.st.Info(ctx)\n}\n\nfunc (s *streamableToolWithCallback) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\treturn streamWithCallbacks(s.st.StreamableRun)(ctx, argumentsInJSON, opts...)\n}\n\ntype enhancedInvokableToolWithCallback struct {\n\teiTool tool.EnhancedInvokableTool\n}\n\nfunc (e *enhancedInvokableToolWithCallback) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn e.eiTool.Info(ctx)\n}\n\nfunc (e *enhancedInvokableToolWithCallback) InvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\treturn invokeEnhancedWithCallbacks(e.eiTool.InvokableRun)(ctx, toolArgument, opts...)\n}\n\ntype enhancedStreamableToolWithCallback struct {\n\test tool.EnhancedStreamableTool\n}\n\nfunc (e *enhancedStreamableToolWithCallback) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn e.est.Info(ctx)\n}\n\nfunc (e *enhancedStreamableToolWithCallback) StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\treturn streamEnhancedWithCallbacks(e.est.StreamableRun)(ctx, toolArgument, opts...)\n}\n\nfunc streamableToInvokable(e StreamableToolEndpoint) InvokableToolEndpoint {\n\treturn func(ctx context.Context, input *ToolInput) (*ToolOutput, error) {\n\t\tso, err := e(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\to, err := concatStreamReader(so.Result)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to concat StreamableTool output message stream: %w\", err)\n\t\t}\n\t\treturn &ToolOutput{Result: o}, nil\n\t}\n}\n\nfunc invokableToStreamable(e InvokableToolEndpoint) StreamableToolEndpoint {\n\treturn func(ctx context.Context, input *ToolInput) (*StreamToolOutput, error) {\n\t\to, err := e(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &StreamToolOutput{Result: schema.StreamReaderFromArray([]string{o.Result})}, nil\n\t}\n}\n\nfunc enhancedStreamableToEnhancedInvokable(e EnhancedStreamableToolEndpoint) EnhancedInvokableToolEndpoint {\n\treturn func(ctx context.Context, input *ToolInput) (*EnhancedInvokableToolOutput, error) {\n\t\tso, err := e(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\to, err := concatStreamReader(so.Result)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to concat EnhancedStreamableTool output message stream: %w\", err)\n\t\t}\n\t\treturn &EnhancedInvokableToolOutput{Result: o}, nil\n\t}\n}\n\nfunc enhancedInvokableToEnhancedStreamable(e EnhancedInvokableToolEndpoint) EnhancedStreamableToolEndpoint {\n\treturn func(ctx context.Context, input *ToolInput) (*EnhancedStreamableToolOutput, error) {\n\t\to, err := e(ctx, input)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{o.Result})}, nil\n\t}\n}\n\nfunc invokeEnhancedWithCallbacks(i func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error)) func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) {\n\treturn runWithCallbacks(i, onStart[*schema.ToolArgument], onEnd[*schema.ToolResult], onError)\n}\n\nfunc streamEnhancedWithCallbacks(s func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error)) func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\treturn runWithCallbacks(s, onStart[*schema.ToolArgument], onEndWithStreamOutput[*schema.ToolResult], onError)\n}\n\ntype toolCallTask struct {\n\t// in\n\tendpoint                   InvokableToolEndpoint\n\tstreamEndpoint             StreamableToolEndpoint\n\tenhancedInvokableEndpoint  EnhancedInvokableToolEndpoint\n\tenhancedStreamableEndpoint EnhancedStreamableToolEndpoint\n\tmeta                       *executorMeta\n\tname                       string\n\targ                        string\n\tcallID                     string\n\tuseEnhanced                bool\n\n\t// out\n\texecuted        bool\n\toutput          string\n\tsOutput         *schema.StreamReader[string]\n\tenhancedOutput  *schema.ToolResult\n\tenhancedSOutput *schema.StreamReader[*schema.ToolResult]\n\terr             error\n}\n\nfunc (tn *ToolsNode) genToolCallTasks(ctx context.Context, tuple *toolsTuple,\n\tinput *schema.Message, executedTools map[string]string, executedEnhancedTools map[string]*schema.ToolResult, isStream bool) ([]toolCallTask, error) {\n\n\tif input.Role != schema.Assistant {\n\t\treturn nil, fmt.Errorf(\"expected message role is Assistant, got %s\", input.Role)\n\t}\n\n\tn := len(input.ToolCalls)\n\tif n == 0 {\n\t\treturn nil, errors.New(\"no tool call found in input message\")\n\t}\n\n\ttoolCallTasks := make([]toolCallTask, n)\n\n\tfor i := 0; i < n; i++ {\n\t\ttoolCall := input.ToolCalls[i]\n\t\tif enhancedResult, executed := executedEnhancedTools[toolCall.ID]; executed {\n\t\t\ttoolCallTasks[i].name = toolCall.Function.Name\n\t\t\ttoolCallTasks[i].arg = toolCall.Function.Arguments\n\t\t\ttoolCallTasks[i].callID = toolCall.ID\n\t\t\ttoolCallTasks[i].executed = true\n\t\t\ttoolCallTasks[i].useEnhanced = true\n\t\t\tif isStream {\n\t\t\t\ttoolCallTasks[i].enhancedSOutput = schema.StreamReaderFromArray([]*schema.ToolResult{enhancedResult})\n\t\t\t} else {\n\t\t\t\ttoolCallTasks[i].enhancedOutput = enhancedResult\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\t\tif result, executed := executedTools[toolCall.ID]; executed {\n\t\t\ttoolCallTasks[i].name = toolCall.Function.Name\n\t\t\ttoolCallTasks[i].arg = toolCall.Function.Arguments\n\t\t\ttoolCallTasks[i].callID = toolCall.ID\n\t\t\ttoolCallTasks[i].executed = true\n\t\t\ttoolCallTasks[i].useEnhanced = false\n\t\t\tif isStream {\n\t\t\t\ttoolCallTasks[i].sOutput = schema.StreamReaderFromArray([]string{result})\n\t\t\t} else {\n\t\t\t\ttoolCallTasks[i].output = result\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\t\tindex, ok := tuple.indexes[toolCall.Function.Name]\n\t\tif !ok {\n\t\t\tif tn.unknownToolHandler == nil {\n\t\t\t\treturn nil, fmt.Errorf(\"tool %s not found in toolsNode indexes\", toolCall.Function.Name)\n\t\t\t}\n\t\t\ttoolCallTasks[i] = newUnknownToolTask(toolCall.Function.Name, toolCall.Function.Arguments, toolCall.ID, tn.unknownToolHandler)\n\t\t} else {\n\t\t\ttoolCallTasks[i].meta = tuple.meta[index]\n\t\t\ttoolCallTasks[i].name = toolCall.Function.Name\n\t\t\ttoolCallTasks[i].callID = toolCall.ID\n\n\t\t\tif tuple.enhancedInvokableEndpoints[index] != nil && tuple.enhancedStreamableEndpoints[index] != nil {\n\t\t\t\ttoolCallTasks[i].enhancedInvokableEndpoint = tuple.enhancedInvokableEndpoints[index]\n\t\t\t\ttoolCallTasks[i].enhancedStreamableEndpoint = tuple.enhancedStreamableEndpoints[index]\n\t\t\t\ttoolCallTasks[i].useEnhanced = true\n\t\t\t} else {\n\t\t\t\ttoolCallTasks[i].endpoint = tuple.endpoints[index]\n\t\t\t\ttoolCallTasks[i].streamEndpoint = tuple.streamEndpoints[index]\n\t\t\t\ttoolCallTasks[i].useEnhanced = false\n\t\t\t}\n\n\t\t\tif tn.toolArgumentsHandler != nil {\n\t\t\t\targ, err := tn.toolArgumentsHandler(ctx, toolCall.Function.Name, toolCall.Function.Arguments)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, fmt.Errorf(\"failed to executed tool[name:%s arguments:%s] arguments handler: %w\", toolCall.Function.Name, toolCall.Function.Arguments, err)\n\t\t\t\t}\n\t\t\t\ttoolCallTasks[i].arg = arg\n\t\t\t} else {\n\t\t\t\ttoolCallTasks[i].arg = toolCall.Function.Arguments\n\t\t\t}\n\t\t}\n\t}\n\n\treturn toolCallTasks, nil\n}\n\nfunc newUnknownToolTask(name, arg, callID string, unknownToolHandler func(ctx context.Context, name, input string) (string, error)) toolCallTask {\n\tendpoint := func(ctx context.Context, input *ToolInput) (*ToolOutput, error) {\n\t\tresult, err := unknownToolHandler(ctx, input.Name, input.Arguments)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &ToolOutput{\n\t\t\tResult: result,\n\t\t}, nil\n\t}\n\treturn toolCallTask{\n\t\tendpoint:       endpoint,\n\t\tstreamEndpoint: invokableToStreamable(endpoint),\n\t\tmeta: &executorMeta{\n\t\t\tcomponent:                  components.ComponentOfTool,\n\t\t\tisComponentCallbackEnabled: false,\n\t\t\tcomponentImplType:          \"UnknownTool\",\n\t\t},\n\t\tname:   name,\n\t\targ:    arg,\n\t\tcallID: callID,\n\t}\n}\n\nfunc runToolCallTaskByInvoke(ctx context.Context, task *toolCallTask, opts ...tool.Option) {\n\tif task.executed {\n\t\treturn\n\t}\n\tctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{\n\t\tName:      task.name,\n\t\tType:      task.meta.componentImplType,\n\t\tComponent: task.meta.component,\n\t})\n\n\tctx = setToolCallInfo(ctx, &toolCallInfo{toolCallID: task.callID})\n\tctx = appendToolAddressSegment(ctx, task.name, task.callID)\n\n\tif task.useEnhanced {\n\t\tenhancedOutput, err := task.enhancedInvokableEndpoint(ctx, &ToolInput{\n\t\t\tName:        task.name,\n\t\t\tArguments:   task.arg,\n\t\t\tCallID:      task.callID,\n\t\t\tCallOptions: opts,\n\t\t})\n\t\tif err != nil {\n\t\t\ttask.err = err\n\t\t} else {\n\t\t\ttask.enhancedOutput = enhancedOutput.Result\n\t\t\ttask.executed = true\n\t\t}\n\t} else {\n\t\toutput, err := task.endpoint(ctx, &ToolInput{\n\t\t\tName:        task.name,\n\t\t\tArguments:   task.arg,\n\t\t\tCallID:      task.callID,\n\t\t\tCallOptions: opts,\n\t\t})\n\t\tif err != nil {\n\t\t\ttask.err = err\n\t\t} else {\n\t\t\ttask.output = output.Result\n\t\t\ttask.executed = true\n\t\t}\n\t}\n}\n\nfunc runToolCallTaskByStream(ctx context.Context, task *toolCallTask, opts ...tool.Option) {\n\tctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{\n\t\tName:      task.name,\n\t\tType:      task.meta.componentImplType,\n\t\tComponent: task.meta.component,\n\t})\n\n\tctx = setToolCallInfo(ctx, &toolCallInfo{toolCallID: task.callID})\n\tctx = appendToolAddressSegment(ctx, task.name, task.callID)\n\n\tif task.useEnhanced {\n\t\tenhancedOutput, err := task.enhancedStreamableEndpoint(ctx, &ToolInput{\n\t\t\tName:        task.name,\n\t\t\tArguments:   task.arg,\n\t\t\tCallID:      task.callID,\n\t\t\tCallOptions: opts,\n\t\t})\n\t\tif err != nil {\n\t\t\ttask.err = err\n\t\t} else {\n\t\t\ttask.enhancedSOutput = enhancedOutput.Result\n\t\t\ttask.executed = true\n\t\t}\n\t} else {\n\t\toutput, err := task.streamEndpoint(ctx, &ToolInput{\n\t\t\tName:        task.name,\n\t\t\tArguments:   task.arg,\n\t\t\tCallID:      task.callID,\n\t\t\tCallOptions: opts,\n\t\t})\n\t\tif err != nil {\n\t\t\ttask.err = err\n\t\t} else {\n\t\t\ttask.sOutput = output.Result\n\t\t\ttask.executed = true\n\t\t}\n\t}\n}\n\nfunc sequentialRunToolCall(ctx context.Context,\n\trun func(ctx2 context.Context, callTask *toolCallTask, opts ...tool.Option),\n\ttasks []toolCallTask, opts ...tool.Option) {\n\n\tfor i := range tasks {\n\t\tif tasks[i].executed {\n\t\t\tcontinue\n\t\t}\n\t\trun(ctx, &tasks[i], opts...)\n\t}\n}\n\nfunc parallelRunToolCall(ctx context.Context,\n\trun func(ctx2 context.Context, callTask *toolCallTask, opts ...tool.Option),\n\ttasks []toolCallTask, opts ...tool.Option) {\n\n\tif len(tasks) == 1 {\n\t\trun(ctx, &tasks[0], opts...)\n\t\treturn\n\t}\n\n\tvar wg sync.WaitGroup\n\tfor i := 1; i < len(tasks); i++ {\n\t\tif tasks[i].executed {\n\t\t\tcontinue\n\t\t}\n\t\twg.Add(1)\n\t\tgo func(ctx_ context.Context, t *toolCallTask, opts ...tool.Option) {\n\t\t\tdefer wg.Done()\n\t\t\tdefer func() {\n\t\t\t\tpanicErr := recover()\n\t\t\t\tif panicErr != nil {\n\t\t\t\t\tt.err = safe.NewPanicErr(panicErr, debug.Stack())\n\t\t\t\t}\n\t\t\t}()\n\t\t\trun(ctx_, t, opts...)\n\t\t}(ctx, &tasks[i], opts...)\n\t}\n\n\tif !tasks[0].executed {\n\t\trun(ctx, &tasks[0], opts...)\n\t}\n\n\twg.Wait()\n}\n\n// Invoke calls the tools and collects the results of invokable tools.\n// it's parallel if there are multiple tool calls in the input message.\nfunc (tn *ToolsNode) Invoke(ctx context.Context, input *schema.Message,\n\topts ...ToolsNodeOption) ([]*schema.Message, error) {\n\n\topt := getToolsNodeOptions(opts...)\n\ttuple := tn.tuple\n\tif opt.ToolList != nil {\n\t\tvar err error\n\t\ttuple, err = convTools(ctx, opt.ToolList, tn.toolCallMiddlewares, tn.streamToolCallMiddlewares, tn.enhancedToolCallMiddlewares, tn.enhancedStreamToolCallMiddlewares)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to convert tool list from call option: %w\", err)\n\t\t}\n\t}\n\n\tvar executedTools map[string]string\n\tvar executedEnhancedTools map[string]*schema.ToolResult\n\tif wasInterrupted, hasState, tnState := GetInterruptState[*toolsInterruptAndRerunState](ctx); wasInterrupted && hasState {\n\t\tinput = tnState.Input\n\t\tif tnState.ExecutedTools != nil {\n\t\t\texecutedTools = tnState.ExecutedTools\n\t\t}\n\t\tif tnState.ExecutedEnhancedTools != nil {\n\t\t\texecutedEnhancedTools = tnState.ExecutedEnhancedTools\n\t\t}\n\t}\n\n\ttasks, err := tn.genToolCallTasks(ctx, tuple, input, executedTools, executedEnhancedTools, false)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tif tn.executeSequentially {\n\t\tsequentialRunToolCall(ctx, runToolCallTaskByInvoke, tasks, opt.ToolOptions...)\n\t} else {\n\t\tparallelRunToolCall(ctx, runToolCallTaskByInvoke, tasks, opt.ToolOptions...)\n\t}\n\n\tn := len(tasks)\n\toutput := make([]*schema.Message, n)\n\n\trerunExtra := &ToolsInterruptAndRerunExtra{\n\t\tToolCalls:             input.ToolCalls,\n\t\tExecutedTools:         make(map[string]string),\n\t\tExecutedEnhancedTools: make(map[string]*schema.ToolResult),\n\t\tRerunExtraMap:         make(map[string]any),\n\t}\n\trerunState := &toolsInterruptAndRerunState{\n\t\tInput:                 input,\n\t\tExecutedTools:         make(map[string]string),\n\t\tExecutedEnhancedTools: make(map[string]*schema.ToolResult),\n\t}\n\n\tvar errs []error\n\tfor i := 0; i < n; i++ {\n\t\tif tasks[i].err != nil {\n\t\t\tinfo, ok := IsInterruptRerunError(tasks[i].err)\n\t\t\tif !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"failed to invoke tool[name:%s id:%s]: %w\", tasks[i].name, tasks[i].callID, tasks[i].err)\n\t\t\t}\n\n\t\t\trerunExtra.RerunTools = append(rerunExtra.RerunTools, tasks[i].callID)\n\t\t\trerunState.RerunTools = append(rerunState.RerunTools, tasks[i].callID)\n\t\t\tif info != nil {\n\t\t\t\trerunExtra.RerunExtraMap[tasks[i].callID] = info\n\t\t\t}\n\n\t\t\tiErr := WrapInterruptAndRerunIfNeeded(ctx,\n\t\t\t\tAddressSegment{ID: tasks[i].callID, Type: AddressSegmentTool}, tasks[i].err)\n\t\t\terrs = append(errs, iErr)\n\t\t\tcontinue\n\t\t}\n\t\tif tasks[i].executed {\n\t\t\tif tasks[i].useEnhanced {\n\t\t\t\trerunExtra.ExecutedEnhancedTools[tasks[i].callID] = tasks[i].enhancedOutput\n\t\t\t\trerunState.ExecutedEnhancedTools[tasks[i].callID] = tasks[i].enhancedOutput\n\t\t\t} else {\n\t\t\t\trerunExtra.ExecutedTools[tasks[i].callID] = tasks[i].output\n\t\t\t\trerunState.ExecutedTools[tasks[i].callID] = tasks[i].output\n\t\t\t}\n\t\t}\n\n\t\tif len(errs) == 0 {\n\t\t\tif tasks[i].useEnhanced {\n\t\t\t\toutput[i] = schema.ToolMessage(\"\", tasks[i].callID, schema.WithToolName(tasks[i].name))\n\t\t\t\toutput[i].UserInputMultiContent, err = tasks[i].enhancedOutput.ToMessageInputParts()\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\toutput[i] = schema.ToolMessage(tasks[i].output, tasks[i].callID, schema.WithToolName(tasks[i].name))\n\t\t\t}\n\n\t\t}\n\t}\n\tif len(errs) > 0 {\n\t\treturn nil, CompositeInterrupt(ctx, rerunExtra, rerunState, errs...)\n\t}\n\n\treturn output, nil\n}\n\n// Stream calls the tools and collects the results of stream readers.\n// it's parallel if there are multiple tool calls in the input message.\nfunc (tn *ToolsNode) Stream(ctx context.Context, input *schema.Message,\n\topts ...ToolsNodeOption) (*schema.StreamReader[[]*schema.Message], error) {\n\n\topt := getToolsNodeOptions(opts...)\n\ttuple := tn.tuple\n\tif opt.ToolList != nil {\n\t\tvar err error\n\t\ttuple, err = convTools(ctx, opt.ToolList, tn.toolCallMiddlewares, tn.streamToolCallMiddlewares, tn.enhancedToolCallMiddlewares, tn.enhancedStreamToolCallMiddlewares)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to convert tool list from call option: %w\", err)\n\t\t}\n\t}\n\n\tvar executedTools map[string]string\n\tvar executedEnhancedTools map[string]*schema.ToolResult\n\tif wasInterrupted, hasState, tnState := GetInterruptState[*toolsInterruptAndRerunState](ctx); wasInterrupted && hasState {\n\t\tinput = tnState.Input\n\t\tif tnState.ExecutedTools != nil {\n\t\t\texecutedTools = tnState.ExecutedTools\n\t\t}\n\t\tif tnState.ExecutedEnhancedTools != nil {\n\t\t\texecutedEnhancedTools = tnState.ExecutedEnhancedTools\n\t\t}\n\t}\n\n\ttasks, err := tn.genToolCallTasks(ctx, tuple, input, executedTools, executedEnhancedTools, true)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tif tn.executeSequentially {\n\t\tsequentialRunToolCall(ctx, runToolCallTaskByStream, tasks, opt.ToolOptions...)\n\t} else {\n\t\tparallelRunToolCall(ctx, runToolCallTaskByStream, tasks, opt.ToolOptions...)\n\t}\n\n\tn := len(tasks)\n\n\trerunExtra := &ToolsInterruptAndRerunExtra{\n\t\tToolCalls:             input.ToolCalls,\n\t\tExecutedTools:         make(map[string]string),\n\t\tExecutedEnhancedTools: make(map[string]*schema.ToolResult),\n\t\tRerunExtraMap:         make(map[string]any),\n\t}\n\trerunState := &toolsInterruptAndRerunState{\n\t\tInput:                 input,\n\t\tExecutedTools:         make(map[string]string),\n\t\tExecutedEnhancedTools: make(map[string]*schema.ToolResult),\n\t}\n\tvar errs []error\n\t// check rerun\n\tfor i := 0; i < n; i++ {\n\t\tif tasks[i].err != nil {\n\t\t\tinfo, ok := IsInterruptRerunError(tasks[i].err)\n\t\t\tif !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"failed to stream tool call %s: %w\", tasks[i].callID, tasks[i].err)\n\t\t\t}\n\n\t\t\trerunExtra.RerunTools = append(rerunExtra.RerunTools, tasks[i].callID)\n\t\t\trerunState.RerunTools = append(rerunState.RerunTools, tasks[i].callID)\n\t\t\tif info != nil {\n\t\t\t\trerunExtra.RerunExtraMap[tasks[i].callID] = info\n\t\t\t}\n\t\t\tiErr := WrapInterruptAndRerunIfNeeded(ctx,\n\t\t\t\tAddressSegment{ID: tasks[i].callID, Type: AddressSegmentTool}, tasks[i].err)\n\t\t\terrs = append(errs, iErr)\n\t\t\tcontinue\n\t\t}\n\t}\n\n\tif len(errs) > 0 {\n\t\t// concat and save tool output\n\t\tfor _, t := range tasks {\n\t\t\tif t.executed {\n\t\t\t\tif t.useEnhanced {\n\t\t\t\t\teo, err_ := concatStreamReader(t.enhancedSOutput)\n\t\t\t\t\tif err_ != nil {\n\t\t\t\t\t\treturn nil, fmt.Errorf(\"failed to concat enhanced tool[name:%s id:%s]'s stream output: %w\", t.name, t.callID, err_)\n\t\t\t\t\t}\n\t\t\t\t\trerunExtra.ExecutedEnhancedTools[t.callID] = eo\n\t\t\t\t\trerunState.ExecutedEnhancedTools[t.callID] = eo\n\n\t\t\t\t} else {\n\t\t\t\t\to, err_ := concatStreamReader(t.sOutput)\n\t\t\t\t\tif err_ != nil {\n\t\t\t\t\t\treturn nil, fmt.Errorf(\"failed to concat tool[name:%s id:%s]'s stream output: %w\", t.name, t.callID, err_)\n\t\t\t\t\t}\n\t\t\t\t\trerunExtra.ExecutedTools[t.callID] = o\n\t\t\t\t\trerunState.ExecutedTools[t.callID] = o\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t\treturn nil, CompositeInterrupt(ctx, rerunExtra, rerunState, errs...)\n\t}\n\n\t// common return\n\tsOutput := make([]*schema.StreamReader[[]*schema.Message], n)\n\tfor i := 0; i < n; i++ {\n\t\tindex := i\n\t\tcallID := tasks[i].callID\n\t\tcallName := tasks[i].name\n\t\tif tasks[i].useEnhanced {\n\t\t\tcvt := func(tr *schema.ToolResult) ([]*schema.Message, error) {\n\t\t\t\tret := make([]*schema.Message, n)\n\t\t\t\tret[index] = schema.ToolMessage(\"\", callID, schema.WithToolName(callName))\n\t\t\t\tret[index].UserInputMultiContent, err = tr.ToMessageInputParts()\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\treturn ret, nil\n\t\t\t}\n\t\t\tsOutput[i] = schema.StreamReaderWithConvert(tasks[i].enhancedSOutput, cvt)\n\t\t} else {\n\t\t\tcvt := func(s string) ([]*schema.Message, error) {\n\t\t\t\tret := make([]*schema.Message, n)\n\t\t\t\tret[index] = schema.ToolMessage(s, callID, schema.WithToolName(callName))\n\t\t\t\treturn ret, nil\n\t\t\t}\n\t\t\tsOutput[i] = schema.StreamReaderWithConvert(tasks[i].sOutput, cvt)\n\t\t}\n\n\t}\n\treturn schema.MergeStreamReaders(sOutput), nil\n}\n\n// GetType returns the component type string for the Tools node.\nfunc (tn *ToolsNode) GetType() string {\n\treturn \"\"\n}\n\nfunc getToolsNodeOptions(opts ...ToolsNodeOption) *toolsNodeOptions {\n\to := &toolsNodeOptions{\n\t\tToolOptions: make([]tool.Option, 0),\n\t}\n\tfor _, opt := range opts {\n\t\topt(o)\n\t}\n\treturn o\n}\n\ntype toolCallInfoKey struct{}\ntype toolCallInfo struct {\n\ttoolCallID string\n}\n\nfunc setToolCallInfo(ctx context.Context, toolCallInfo *toolCallInfo) context.Context {\n\treturn context.WithValue(ctx, toolCallInfoKey{}, toolCallInfo)\n}\n\n// GetToolCallID gets the current tool call id from the context.\nfunc GetToolCallID(ctx context.Context) string {\n\tv := ctx.Value(toolCallInfoKey{})\n\tif v == nil {\n\t\treturn \"\"\n\t}\n\n\tinfo, ok := v.(*toolCallInfo)\n\tif !ok {\n\t\treturn \"\"\n\t}\n\n\treturn info.toolCallID\n}\n"
  },
  {
    "path": "compose/tool_node_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"io\"\n\t\"strconv\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/internal\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nconst (\n\ttoolNameOfUserCompany = \"user_company\"\n\ttoolIDOfUserCompany   = \"call_TRZhlagwBS0LpWbWPeZOvIXc\"\n\n\ttoolNameOfUserSalary = \"user_salary\"\n\ttoolIDOfUserSalary   = \"call_AqfoRW6fuF98k0o7696k2nzm\"\n)\n\nfunc TestToolsNode(t *testing.T) {\n\tvar err error\n\tctx := context.Background()\n\n\tuserCompanyToolInfo := &schema.ToolInfo{\n\t\tName: toolNameOfUserCompany,\n\t\tDesc: \"Query user's company and position information based on user's name and email\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"name\": {\n\t\t\t\t\tType: \"string\",\n\t\t\t\t\tDesc: \"User's name\",\n\t\t\t\t},\n\t\t\t\t\"email\": {\n\t\t\t\t\tType: \"string\",\n\t\t\t\t\tDesc: \"User's email\",\n\t\t\t\t},\n\t\t\t}),\n\t}\n\n\tuserSalaryToolInfo := &schema.ToolInfo{\n\t\tName: toolNameOfUserSalary,\n\t\tDesc: \"Query user's salary information based on user's name and email\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"name\": {\n\t\t\t\t\tType: \"string\",\n\t\t\t\t\tDesc: \"User's name\",\n\t\t\t\t},\n\t\t\t\t\"email\": {\n\t\t\t\t\tType: \"string\",\n\t\t\t\t\tDesc: \"User's email\",\n\t\t\t\t},\n\t\t\t}),\n\t}\n\n\tt.Run(\"success\", func(t *testing.T) {\n\t\tconst (\n\t\t\tnodeOfTools = \"tools\"\n\t\t\tnodeOfModel = \"model\"\n\t\t)\n\t\tg := NewGraph[[]*schema.Message, []*schema.Message]()\n\n\t\terr = g.AddChatModelNode(nodeOfModel, &mockIntentChatModel{})\n\t\tassert.NoError(t, err)\n\n\t\tui := newTool(userCompanyToolInfo, queryUserCompany)\n\t\tus := newStreamableTool(userSalaryToolInfo, queryUserSalary)\n\n\t\ttoolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{ui, us},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddToolsNode(nodeOfTools, toolsNode)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddEdge(START, nodeOfModel)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddEdge(nodeOfModel, nodeOfTools)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddEdge(nodeOfTools, END)\n\t\tassert.NoError(t, err)\n\n\t\tr, err := g.Compile(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tout, err := r.Invoke(ctx, []*schema.Message{})\n\t\tassert.NoError(t, err)\n\n\t\tmsg := findMsgByToolCallID(out, toolIDOfUserCompany)\n\t\tassert.Equal(t, toolIDOfUserCompany, msg.ToolCallID)\n\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"gender\":\"male\",\"company\":\"bytedance\",\"position\":\"CEO\"}`,\n\t\t\tmsg.Content)\n\n\t\tmsg = findMsgByToolCallID(out, toolIDOfUserSalary)\n\t\tassert.Equal(t, toolIDOfUserSalary, msg.ToolCallID)\n\t\tassert.Contains(t, msg.Content,\n\t\t\t`{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":5000}{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":3000}{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":2000}`)\n\n\t\t// 测试流式调用\n\t\treader, err := r.Stream(ctx, []*schema.Message{})\n\t\tassert.NoError(t, err)\n\t\tloops := 0\n\t\tuserSalaryTimes := 0\n\n\t\tdefer reader.Close()\n\n\t\tvar arrMsgs [][]*schema.Message\n\t\tfor ; loops < 10; loops++ {\n\t\t\tmsgs, err := reader.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tarrMsgs = append(arrMsgs, msgs)\n\n\t\t\tassert.NoError(t, err)\n\n\t\t\tassert.Len(t, msgs, 2)\n\t\t\tif msg := findMsgByToolCallID(out, toolIDOfUserCompany); msg != nil {\n\t\t\t\tassert.Equal(t, schema.Tool, msg.Role)\n\t\t\t\tassert.Equal(t, toolIDOfUserCompany, msg.ToolCallID)\n\t\t\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"gender\":\"male\",\"company\":\"bytedance\",\"position\":\"CEO\"}`,\n\t\t\t\t\tmsg.Content)\n\t\t\t} else if msg := findMsgByToolCallID(out, toolIDOfUserSalary); msg != nil {\n\t\t\t\tassert.Equal(t, schema.Tool, msg.Role)\n\t\t\t\tassert.Equal(t, toolIDOfUserSalary, msg.ToolCallID)\n\n\t\t\t\tswitch userSalaryTimes {\n\t\t\t\tcase 0:\n\t\t\t\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":5000}`,\n\t\t\t\t\t\tmsg.Content)\n\t\t\t\tcase 1:\n\t\t\t\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":3000}`,\n\t\t\t\t\t\tmsg.Content)\n\t\t\t\tcase 2:\n\t\t\t\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":2000}`,\n\t\t\t\t\t\tmsg.Content)\n\t\t\t\t}\n\n\t\t\t\tuserSalaryTimes++\n\t\t\t} else {\n\t\t\t\tassert.Fail(t, \"unexpected tool name\")\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 4, loops)\n\n\t\tmsgs, err_ := schema.ConcatMessageArray(arrMsgs)\n\t\tassert.NoError(t, err_)\n\t\tmsg = findMsgByToolCallID(msgs, toolIDOfUserCompany)\n\t\tmsg = findMsgByToolCallID(msgs, toolIDOfUserSalary)\n\n\t\tsr, sw := schema.Pipe[[]*schema.Message](2)\n\t\tsw.Send([]*schema.Message{\n\t\t\t{\n\t\t\t\tRole:    schema.User,\n\t\t\t\tContent: `hi, how are you`,\n\t\t\t},\n\t\t}, nil)\n\t\tsw.Send([]*schema.Message{\n\t\t\t{\n\t\t\t\tRole:    schema.User,\n\t\t\t\tContent: `i'm fine'`,\n\t\t\t},\n\t\t}, nil)\n\t\tsw.Close()\n\n\t\treader, err = r.Transform(ctx, sr)\n\t\tassert.NoError(t, err)\n\n\t\tdefer reader.Close()\n\n\t\tloops = 0\n\t\tuserSalaryTimes = 0\n\n\t\tfor ; loops < 10; loops++ {\n\t\t\tmsgs, err := reader.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tassert.NoError(t, err)\n\n\t\t\tassert.Len(t, msgs, 2)\n\t\t\tif msg := findMsgByToolCallID(out, toolIDOfUserCompany); msg != nil {\n\t\t\t\tassert.Equal(t, schema.Tool, msg.Role)\n\t\t\t\tassert.Equal(t, toolIDOfUserCompany, msg.ToolCallID)\n\t\t\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"gender\":\"male\",\"company\":\"bytedance\",\"position\":\"CEO\"}`,\n\t\t\t\t\tmsg.Content)\n\t\t\t} else if msg := findMsgByToolCallID(out, toolIDOfUserSalary); msg != nil {\n\t\t\t\tassert.Equal(t, schema.Tool, msg.Role)\n\t\t\t\tassert.Equal(t, toolIDOfUserSalary, msg.ToolCallID)\n\n\t\t\t\tswitch userSalaryTimes {\n\t\t\t\tcase 0:\n\t\t\t\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":5000}`,\n\t\t\t\t\t\tmsg.Content)\n\t\t\t\tcase 1:\n\t\t\t\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":3000}`,\n\t\t\t\t\t\tmsg.Content)\n\t\t\t\tcase 2:\n\t\t\t\t\tassert.JSONEq(t, `{\"user_id\":\"zhangsan-zhangsan@bytedance.com\",\"salary\":2000}`,\n\t\t\t\t\t\tmsg.Content)\n\t\t\t\t}\n\n\t\t\t\tuserSalaryTimes++\n\t\t\t} else {\n\t\t\t\tassert.Fail(t, \"unexpected tool name\")\n\t\t\t}\n\t\t}\n\n\t\tassert.Equal(t, 4, loops)\n\t})\n\n\tt.Run(\"order_consistency\", func(t *testing.T) {\n\t\t// Create a ToolsNode with multiple tools\n\t\tui := newTool(userCompanyToolInfo, queryUserCompany)\n\t\tus := newTool(userSalaryToolInfo, queryUserSalary)\n\n\t\ttoolsNode, err_ := NewToolNode(context.Background(), &ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{ui, us},\n\t\t})\n\t\tassert.NoError(t, err_)\n\n\t\t// Create an input message with multiple tool calls in a specific order\n\t\tinput := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: toolIDOfUserSalary,\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      toolNameOfUserSalary,\n\t\t\t\t\t\tArguments: `{\"name\": \"zhangsan\", \"email\": \"zhangsan@bytedance.com\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t{\n\t\t\t\t\tID: toolIDOfUserCompany,\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      toolNameOfUserCompany,\n\t\t\t\t\t\tArguments: `{\"name\": \"zhangsan\", \"email\": \"zhangsan@bytedance.com\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t// Invoke the ToolsNode\n\t\toutput, err_ := toolsNode.Invoke(context.Background(), input)\n\t\tassert.NoError(t, err_)\n\n\t\t// Verify the order of output messages matches the order of input tool calls\n\t\tassert.Equal(t, 2, len(output))\n\t\tassert.Equal(t, toolIDOfUserSalary, output[0].ToolCallID)\n\t\tassert.Equal(t, toolIDOfUserCompany, output[1].ToolCallID)\n\n\t\t// Test with Stream method as well\n\t\tstreamer, err_ := toolsNode.Stream(context.Background(), input)\n\t\tassert.NoError(t, err_)\n\t\tdefer streamer.Close()\n\n\t\t// Collect all stream outputs\n\t\tvar streamOutputs [][]*schema.Message\n\t\tfor {\n\t\t\tchunk, err__ := streamer.Recv()\n\t\t\tif err__ == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err__)\n\t\t\tstreamOutputs = append(streamOutputs, chunk)\n\t\t}\n\n\t\t// Verify each chunk maintains the correct order\n\t\tfor _, chunk := range streamOutputs {\n\t\t\tif chunk[0] != nil {\n\t\t\t\tassert.Equal(t, toolIDOfUserSalary, chunk[0].ToolCallID)\n\t\t\t}\n\t\t\tif chunk[1] != nil {\n\t\t\t\tassert.Equal(t, toolIDOfUserCompany, chunk[1].ToolCallID)\n\t\t\t}\n\t\t}\n\n\t\t// Concatenate all stream outputs and verify final result\n\t\tconcatenated, err_ := schema.ConcatMessageArray(streamOutputs)\n\t\tassert.NoError(t, err_)\n\t\tassert.Equal(t, 2, len(concatenated))\n\t\tassert.Equal(t, toolIDOfUserSalary, concatenated[0].ToolCallID)\n\t\tassert.Equal(t, toolIDOfUserCompany, concatenated[1].ToolCallID)\n\t})\n}\n\ntype userCompanyRequest struct {\n\tName  string `json:\"name\"`\n\tEmail string `json:\"email\"`\n}\n\ntype userCompanyResponse struct {\n\tUserID   string `json:\"user_id\"`\n\tGender   string `json:\"gender\"`\n\tCompany  string `json:\"company\"`\n\tPosition string `json:\"position\"`\n}\n\nfunc queryUserCompany(ctx context.Context, req *userCompanyRequest) (resp *userCompanyResponse, err error) {\n\tcallID := GetToolCallID(ctx)\n\tif callID != toolIDOfUserCompany {\n\t\treturn nil, fmt.Errorf(\"invalid tool call id= %s\", callID)\n\t}\n\n\treturn &userCompanyResponse{\n\t\tUserID:   fmt.Sprintf(\"%v-%v\", req.Name, req.Email),\n\t\tGender:   \"male\",\n\t\tCompany:  \"bytedance\",\n\t\tPosition: \"CEO\",\n\t}, nil\n}\n\ntype userSalaryRequest struct {\n\tName  string `json:\"name\"`\n\tEmail string `json:\"email\"`\n}\n\ntype userSalaryResponse struct {\n\tUserID string `json:\"user_id\"`\n\tSalary int    `json:\"salary\"`\n}\n\nfunc queryUserSalary(ctx context.Context, req *userSalaryRequest) (resp *schema.StreamReader[*userSalaryResponse], err error) {\n\tcallID := GetToolCallID(ctx)\n\tif callID != toolIDOfUserSalary {\n\t\treturn nil, fmt.Errorf(\"invalid tool call id= %s\", callID)\n\t}\n\n\tsr, sw := schema.Pipe[*userSalaryResponse](10)\n\tsw.Send(&userSalaryResponse{\n\t\tUserID: fmt.Sprintf(\"%v-%v\", req.Name, req.Email),\n\t\tSalary: 5000,\n\t}, nil)\n\n\tsw.Send(&userSalaryResponse{\n\t\tUserID: fmt.Sprintf(\"%v-%v\", req.Name, req.Email),\n\t\tSalary: 3000,\n\t}, nil)\n\n\tsw.Send(&userSalaryResponse{\n\t\tUserID: fmt.Sprintf(\"%v-%v\", req.Name, req.Email),\n\t\tSalary: 2000,\n\t}, nil)\n\tsw.Close()\n\treturn sr, nil\n}\n\ntype mockIntentChatModel struct{}\n\nfunc (m *mockIntentChatModel) BindTools(tools []*schema.ToolInfo) error {\n\treturn nil\n}\n\nfunc (m *mockIntentChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\treturn &schema.Message{\n\t\tRole:    schema.Assistant,\n\t\tContent: \"\",\n\t\tToolCalls: []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: toolIDOfUserCompany,\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      toolNameOfUserCompany,\n\t\t\t\t\tArguments: `{\"name\": \"zhangsan\", \"email\": \"zhangsan@bytedance.com\"}`,\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tID: toolIDOfUserSalary,\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      toolNameOfUserSalary,\n\t\t\t\t\tArguments: `{\"name\": \"zhangsan\", \"email\": \"zhangsan@bytedance.com\"}`,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc (m *mockIntentChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tsr, sw := schema.Pipe[*schema.Message](2)\n\tsw.Send(&schema.Message{\n\t\tRole:    schema.Assistant,\n\t\tContent: \"\",\n\t\tToolCalls: []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: toolIDOfUserCompany,\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      toolNameOfUserCompany,\n\t\t\t\t\tArguments: `{\"name\": \"zhangsan\", \"email\": \"zhangsan@bytedance.com\"}`,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}, nil)\n\n\tsw.Send(&schema.Message{\n\t\tRole:    schema.Assistant,\n\t\tContent: \"\",\n\t\tToolCalls: []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: toolIDOfUserSalary,\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      toolNameOfUserSalary,\n\t\t\t\t\tArguments: `{\"name\": \"zhangsan\", \"email\": \"zhangsan@bytedance.com\"}`,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}, nil)\n\n\tsw.Close()\n\n\treturn sr, nil\n}\n\nfunc TestToolsNodeOptions(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"tool_option\", func(t *testing.T) {\n\n\t\tg := NewGraph[*schema.Message, []*schema.Message]()\n\n\t\tmt := &mockTool{}\n\n\t\ttn, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{mt},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddToolsNode(\"tools\", tn)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddEdge(START, \"tools\")\n\t\tassert.NoError(t, err)\n\t\terr = g.AddEdge(\"tools\", END)\n\t\tassert.NoError(t, err)\n\n\t\tr, err := g.Compile(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tout, err := r.Invoke(ctx, &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: toolIDOfUserCompany,\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      \"mock_tool\",\n\t\t\t\t\t\tArguments: `{\"name\": \"jack\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, WithToolsNodeOption(WithToolOption(WithAge(10))))\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, out, 1)\n\t\tassert.JSONEq(t, `{\"echo\": \"jack: 10\"}`, out[0].Content)\n\n\t\toutMessages := make([][]*schema.Message, 0)\n\t\toutStream, err := r.Stream(ctx, &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: toolIDOfUserCompany,\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      \"mock_tool\",\n\t\t\t\t\t\tArguments: `{\"name\": \"jack\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, WithToolsNodeOption(WithToolOption(WithAge(10))))\n\n\t\tassert.NoError(t, err)\n\n\t\tfor {\n\t\t\tmsgs, err := outStream.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\toutMessages = append(outMessages, msgs)\n\t\t}\n\t\toutStream.Close()\n\n\t\tmsgs, err := internal.ConcatItems(outMessages)\n\t\tassert.NoError(t, err)\n\n\t\tassert.Len(t, msgs, 1)\n\t\tassert.JSONEq(t, `{\"echo\":\"jack: 10\"}`, msgs[0].Content)\n\t})\n\tt.Run(\"tool_list\", func(t *testing.T) {\n\n\t\tg := NewGraph[*schema.Message, []*schema.Message]()\n\n\t\tmt := &mockTool{}\n\n\t\ttn, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddToolsNode(\"tools\", tn)\n\t\tassert.NoError(t, err)\n\n\t\terr = g.AddEdge(START, \"tools\")\n\t\tassert.NoError(t, err)\n\t\terr = g.AddEdge(\"tools\", END)\n\t\tassert.NoError(t, err)\n\n\t\tr, err := g.Compile(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tout, err := r.Invoke(ctx, &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: toolIDOfUserCompany,\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      \"mock_tool\",\n\t\t\t\t\t\tArguments: `{\"name\": \"jack\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, WithToolsNodeOption(WithToolList(mt), WithToolOption(WithAge(10))))\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, out, 1)\n\t\tassert.JSONEq(t, `{\"echo\": \"jack: 10\"}`, out[0].Content)\n\n\t\toutMessages := make([][]*schema.Message, 0)\n\t\toutStream, err := r.Stream(ctx, &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tID: toolIDOfUserCompany,\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      \"mock_tool\",\n\t\t\t\t\t\tArguments: `{\"name\": \"jack\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, WithToolsNodeOption(WithToolList(mt), WithToolOption(WithAge(10))))\n\n\t\tassert.NoError(t, err)\n\n\t\tfor {\n\t\t\tmsgs, err := outStream.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\toutMessages = append(outMessages, msgs)\n\t\t}\n\t\toutStream.Close()\n\n\t\tmsgs, err := internal.ConcatItems(outMessages)\n\t\tassert.NoError(t, err)\n\n\t\tassert.Len(t, msgs, 1)\n\t\tassert.JSONEq(t, `{\"echo\":\"jack: 10\"}`, msgs[0].Content)\n\t})\n\n}\n\nfunc findMsgByToolCallID(msgs []*schema.Message, toolCallID string) *schema.Message {\n\tfor _, msg := range msgs {\n\t\tif msg.ToolCallID == toolCallID {\n\t\t\treturn msg\n\t\t}\n\t}\n\n\treturn nil\n}\n\ntype mockToolOptions struct {\n\tAge int\n}\n\nfunc WithAge(age int) tool.Option {\n\treturn tool.WrapImplSpecificOptFn(func(o *mockToolOptions) {\n\t\to.Age = age\n\t})\n}\n\ntype mockToolRequest struct {\n\tName string `json:\"name\"`\n}\n\ntype mockToolResponse struct {\n\tEcho string `json:\"echo\"`\n}\n\ntype mockTool struct{}\n\nfunc (m *mockTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"mock_tool\",\n\t\tDesc: \"mock tool\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"name\": {\n\t\t\t\t\tType:     \"string\",\n\t\t\t\t\tDesc:     \"name\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t},\n\t\t\t}),\n\t}, nil\n}\n\nfunc (m *mockTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\topt := tool.GetImplSpecificOptions(&mockToolOptions{}, opts...)\n\n\treq := &mockToolRequest{}\n\n\tif e := sonic.UnmarshalString(argumentsInJSON, req); e != nil {\n\t\treturn \"\", e\n\t}\n\n\tresp := &mockToolResponse{\n\t\tEcho: fmt.Sprintf(\"%v: %v\", req.Name, opt.Age),\n\t}\n\n\treturn sonic.MarshalString(resp)\n}\n\nfunc (m *mockTool) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\tsr, sw := schema.Pipe[string](1)\n\tgo func() {\n\t\tdefer sw.Close()\n\n\t\topt := tool.GetImplSpecificOptions(&mockToolOptions{}, opts...)\n\n\t\treq := &mockToolRequest{}\n\n\t\tif e := sonic.UnmarshalString(argumentsInJSON, req); e != nil {\n\t\t\tsw.Send(\"\", e)\n\t\t\treturn\n\t\t}\n\n\t\tresp := mockToolResponse{\n\t\t\tEcho: fmt.Sprintf(\"%v: %v\", req.Name, opt.Age),\n\t\t}\n\n\t\toutput, err := sonic.MarshalString(resp)\n\t\tif err != nil {\n\t\t\tsw.Send(\"\", err)\n\t\t\treturn\n\t\t}\n\n\t\tfor i := 0; i < len(output); i++ {\n\t\t\tsw.Send(string(output[i]), nil)\n\t\t}\n\t}()\n\n\treturn sr, nil\n}\n\nfunc TestUnknownTool(t *testing.T) {\n\tctx := context.Background()\n\ttn, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\tTools: nil,\n\t\tUnknownToolsHandler: func(ctx context.Context, name, input string) (string, error) {\n\t\t\treturn \"unknown\", nil\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tinput := &schema.Message{\n\t\tRole: schema.Assistant,\n\t\tToolCalls: []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: \"1\",\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      \"unknown1\",\n\t\t\t\t\tArguments: `arg1`,\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tID: \"2\",\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      \"unknown2\",\n\t\t\t\t\tArguments: `arg2`,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\n\texpected := []*schema.Message{\n\t\t{\n\t\t\tRole:       schema.Tool,\n\t\t\tContent:    \"unknown\",\n\t\t\tToolCallID: \"1\",\n\t\t\tToolName:   \"unknown1\",\n\t\t},\n\t\t{\n\t\t\tRole:       schema.Tool,\n\t\t\tContent:    \"unknown\",\n\t\t\tToolCallID: \"2\",\n\t\t\tToolName:   \"unknown2\",\n\t\t},\n\t}\n\n\tresult, err := tn.Invoke(ctx, input)\n\tassert.NoError(t, err)\n\tassert.Equal(t, expected, result)\n\n\tstreamResult, err := tn.Stream(ctx, input)\n\tassert.NoError(t, err)\n\tresult = make([]*schema.Message, 2)\n\tfor {\n\t\tchunk, err := streamResult.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tfor i := range chunk {\n\t\t\tif chunk[i] != nil {\n\t\t\t\tresult[i] = chunk[i]\n\t\t\t}\n\t\t}\n\t}\n\tassert.Equal(t, expected, result)\n}\n\nfunc TestToolRerun(t *testing.T) {\n\ttype myToolRerunState struct {\n\t\tIn *schema.Message\n\t}\n\n\tschema.Register[myToolRerunState]()\n\n\ttc := []schema.ToolCall{\n\t\t{\n\t\t\tID: \"3\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"tool3\",\n\t\t\t\tArguments: \"input\",\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tID: \"4\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"tool4\",\n\t\t\t\tArguments: \"input\",\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tID: \"1\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"tool1\",\n\t\t\t\tArguments: \"input\",\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tID: \"2\",\n\t\t\tFunction: schema.FunctionCall{\n\t\t\t\tName:      \"tool2\",\n\t\t\t\tArguments: \"input\",\n\t\t\t},\n\t\t},\n\t}\n\tg := NewGraph[*schema.Message, string](WithGenLocalState(func(ctx context.Context) (state *myToolRerunState) {\n\t\treturn &myToolRerunState{In: &schema.Message{Role: schema.Assistant, ToolCalls: tc}}\n\t}))\n\tctx := context.Background()\n\ttn, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\tTools: []tool.BaseTool{&myTool1{}, &myTool2{}, &myTool3{t: t}, &myTool4{t: t}},\n\t})\n\tassert.NoError(t, err)\n\tassert.NoError(t, g.AddToolsNode(\"tool node\", tn))\n\tassert.NoError(t, g.AddLambdaNode(\"lambda\", InvokableLambda(func(ctx context.Context, input []*schema.Message) (output string, err error) {\n\t\tcontents := make([]string, len(input))\n\t\tfor _, m := range input {\n\t\t\tcallID := m.ToolCallID\n\t\t\tcallIDInt, err := strconv.Atoi(callID)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", err\n\t\t\t}\n\t\t\tcontents[callIDInt-1] = m.Content\n\t\t}\n\t\tsb := strings.Builder{}\n\t\tfor _, m := range contents {\n\t\t\tsb.WriteString(m)\n\t\t}\n\t\treturn sb.String(), nil\n\t})))\n\tassert.NoError(t, g.AddEdge(START, \"tool node\"))\n\tassert.NoError(t, g.AddEdge(\"tool node\", \"lambda\"))\n\tassert.NoError(t, g.AddEdge(\"lambda\", END))\n\n\tr, err := g.Compile(ctx, WithCheckPointStore(&inMemoryStore{m: map[string][]byte{}}))\n\tassert.NoError(t, err)\n\n\t_, err = r.Stream(ctx, &schema.Message{Role: schema.Assistant, ToolCalls: tc}, WithCheckPointID(\"1\"))\n\tinfo, ok := ExtractInterruptInfo(err)\n\tassert.True(t, ok)\n\tassert.Equal(t, []string{\"tool node\"}, info.RerunNodes)\n\tassert.Equal(t, &ToolsInterruptAndRerunExtra{\n\t\tToolCalls:     tc,\n\t\tRerunTools:    []string{\"1\", \"2\"},\n\t\tRerunExtraMap: map[string]any{\"1\": \"tool1 rerun extra\", \"2\": \"tool2 rerun extra\"},\n\t\tExecutedTools: map[string]string{\n\t\t\t\"3\": \"tool3 input: input\",\n\t\t\t\"4\": \"tool4 input: input\",\n\t\t},\n\t\tExecutedEnhancedTools: make(map[string]*schema.ToolResult),\n\t}, info.RerunNodesExtra[\"tool node\"])\n\n\tsr, err := r.Stream(ctx, nil, WithCheckPointID(\"1\"))\n\tassert.NoError(t, err)\n\tresult, err := concatStreamReader(sr)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"tool1 input: inputtool2 input: inputtool3 input: inputtool4 input: input\", result)\n}\n\nfunc TestToolMiddleware(t *testing.T) {\n\tctx := context.Background()\n\tt3 := &myTool3{t: t}\n\tt4 := &myTool4{t: t}\n\ttn, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\tTools: []tool.BaseTool{t3, t4},\n\t\tToolCallMiddlewares: []ToolMiddleware{\n\t\t\t{\n\t\t\t\tInvokable: func(endpoint InvokableToolEndpoint) InvokableToolEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, input *ToolInput) (*ToolOutput, error) {\n\t\t\t\t\t\t_, err := endpoint(ctx, input)\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn &ToolOutput{Result: \"middleware1\"}, nil\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tStreamable: func(endpoint StreamableToolEndpoint) StreamableToolEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, input *ToolInput) (*StreamToolOutput, error) {\n\t\t\t\t\t\t_, err := endpoint(ctx, input)\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\treturn nil, err\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn &StreamToolOutput{Result: schema.StreamReaderFromArray([]string{\"middleware2\"})}, nil\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tmessages, err := tn.Invoke(ctx, schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{ID: \"1\", Function: schema.FunctionCall{Name: \"tool3\", Arguments: \"\"}},\n\t\t{ID: \"2\", Function: schema.FunctionCall{Name: \"tool4\", Arguments: \"\"}},\n\t}))\n\tassert.NoError(t, err)\n\tassert.Len(t, messages, 2)\n\tassert.Equal(t, \"middleware1\", messages[0].Content)\n\tassert.Equal(t, \"middleware2\", messages[1].Content)\n\n\tt3.times, t4.times = 0, 0 // reset t3 t4\n\tmessageStreams, err := tn.Stream(ctx, schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t{ID: \"1\", Function: schema.FunctionCall{Name: \"tool3\", Arguments: \"\"}},\n\t\t{ID: \"2\", Function: schema.FunctionCall{Name: \"tool4\", Arguments: \"\"}},\n\t}))\n\tassert.NoError(t, err)\n\tvar messageArray [][]*schema.Message\n\tfor {\n\t\tchunk, err := messageStreams.Recv()\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tassert.NoError(t, err)\n\t\tmessageArray = append(messageArray, chunk)\n\t}\n\tmessages, err = schema.ConcatMessageArray(messageArray)\n\tassert.Len(t, messages, 2)\n\tassert.Equal(t, \"middleware1\", messages[0].Content)\n\tassert.Equal(t, \"middleware2\", messages[1].Content)\n}\n\ntype myTool1 struct {\n\ttimes uint\n}\n\nfunc (m *myTool1) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"tool1\"}, nil\n}\n\nfunc (m *myTool1) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tif m.times == 0 {\n\t\tm.times++\n\t\treturn \"\", tool.Interrupt(ctx, \"tool1 rerun extra\")\n\t}\n\treturn \"tool1 input: \" + argumentsInJSON, nil\n}\n\ntype myTool2 struct {\n\ttimes uint\n}\n\nfunc (m *myTool2) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"tool2\"}, nil\n}\n\nfunc (m *myTool2) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\tif m.times == 0 {\n\t\tm.times++\n\t\treturn nil, tool.Interrupt(ctx, \"tool2 rerun extra\")\n\t}\n\treturn schema.StreamReaderFromArray([]string{\"tool2 input: \", argumentsInJSON}), nil\n}\n\ntype myTool3 struct {\n\tt     *testing.T\n\ttimes int\n}\n\nfunc (m *myTool3) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"tool3\"}, nil\n}\n\nfunc (m *myTool3) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\tassert.Equal(m.t, 0, m.times)\n\tm.times++\n\treturn \"tool3 input: \" + argumentsInJSON, nil\n}\n\ntype myTool4 struct {\n\tt     *testing.T\n\ttimes int\n}\n\nfunc (m *myTool4) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"tool4\"}, nil\n}\n\nfunc (m *myTool4) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) {\n\tassert.Equal(m.t, 0, m.times)\n\tm.times++\n\treturn schema.StreamReaderFromArray([]string{\"tool4 input: \", argumentsInJSON}), nil\n}\n\nfunc newTool[I, O any](info *schema.ToolInfo, f func(ctx context.Context, in I) (O, error)) tool.InvokableTool {\n\treturn &invokableTool[I, O]{\n\t\tinfo: info,\n\t\tfn:   f,\n\t}\n}\n\ntype invokableTool[I, O any] struct {\n\tinfo *schema.ToolInfo\n\tfn   func(ctx context.Context, in I) (O, error)\n}\n\nfunc (f *invokableTool[I, O]) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn f.info, nil\n}\n\nfunc (f *invokableTool[I, O]) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\tt := generic.NewInstance[I]()\n\terr := sonic.UnmarshalString(argumentsInJSON, t)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\to, err := f.fn(ctx, t)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn sonic.MarshalString(o)\n}\n\nfunc newStreamableTool[I, O any](info *schema.ToolInfo, f func(ctx context.Context, in I) (*schema.StreamReader[O], error)) tool.StreamableTool {\n\treturn &streamableTool[I, O]{\n\t\tinfo: info,\n\t\tfn:   f,\n\t}\n}\n\ntype streamableTool[I, O any] struct {\n\tinfo *schema.ToolInfo\n\tfn   func(ctx context.Context, in I) (*schema.StreamReader[O], error)\n}\n\nfunc (f *streamableTool[I, O]) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn f.info, nil\n}\nfunc (f *streamableTool[I, O]) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) {\n\tt := generic.NewInstance[I]()\n\terr := sonic.UnmarshalString(argumentsInJSON, t)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tsr, err := f.fn(ctx, t)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn schema.StreamReaderWithConvert(sr, func(o O) (string, error) {\n\t\treturn sonic.MarshalString(o)\n\t}), nil\n}\n\ntype enhancedInvokableTool struct {\n\tinfo *schema.ToolInfo\n\tfn   func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error)\n}\n\nfunc (e *enhancedInvokableTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn e.info, nil\n}\n\nfunc (e *enhancedInvokableTool) InvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) {\n\treturn e.fn(ctx, toolArgument)\n}\n\ntype enhancedStreamableTool struct {\n\tinfo *schema.ToolInfo\n\tfn   func(ctx context.Context, input *schema.ToolArgument) (*schema.StreamReader[*schema.ToolResult], error)\n}\n\nfunc (e *enhancedStreamableTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn e.info, nil\n}\n\nfunc (e *enhancedStreamableTool) StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {\n\treturn e.fn(ctx, toolArgument)\n}\n\nfunc TestEnhancedToolNode(t *testing.T) {\n\tctx := context.Background()\n\n\tenhancedInvokable := &enhancedInvokableTool{\n\t\tinfo: &schema.ToolInfo{\n\t\t\tName: \"enhanced_invokable_tool\",\n\t\t\tDesc: \"test enhanced invokable tool\",\n\t\t},\n\t\tfn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) {\n\t\t\treturn &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t{Type: schema.ToolPartTypeText, Text: \"invokable result: \" + input.Text},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tenhancedStreamable := &enhancedStreamableTool{\n\t\tinfo: &schema.ToolInfo{\n\t\t\tName: \"enhanced_streamable_tool\",\n\t\t\tDesc: \"test enhanced streamable tool\",\n\t\t},\n\t\tfn: func(ctx context.Context, input *schema.ToolArgument) (*schema.StreamReader[*schema.ToolResult], error) {\n\t\t\tresults := []*schema.ToolResult{\n\t\t\t\t{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: \"stream part 1: \" + input.Text}}},\n\t\t\t\t{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: \" stream part 2\"}}},\n\t\t\t}\n\t\t\treturn schema.StreamReaderFromArray(results), nil\n\t\t},\n\t}\n\n\ttoolNode, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\tTools: []tool.BaseTool{enhancedInvokable, enhancedStreamable},\n\t})\n\tassert.NoError(t, err)\n\tassert.NotNil(t, toolNode)\n\n\tt.Run(\"enhanced invokable tool\", func(t *testing.T) {\n\t\tinput := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: \"call1\",\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      \"enhanced_invokable_tool\",\n\t\t\t\t\tArguments: \"test input\",\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\toutput, err := toolNode.Invoke(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, output, 1)\n\t\tassert.Equal(t, schema.Tool, output[0].Role)\n\t\tassert.Equal(t, \"call1\", output[0].ToolCallID)\n\t})\n\n\tt.Run(\"enhanced streamable tool\", func(t *testing.T) {\n\t\tinput := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: \"call2\",\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      \"enhanced_streamable_tool\",\n\t\t\t\t\tArguments: \"test stream\",\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\tstreamReader, err := toolNode.Stream(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, streamReader)\n\n\t\tvar messages []*schema.Message\n\t\tfor {\n\t\t\tchunk, err := streamReader.Recv()\n\t\t\tif err != nil {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif chunk != nil {\n\t\t\t\tmessages = append(messages, chunk...)\n\t\t\t}\n\t\t}\n\t\tmessage, err := schema.ConcatMessages(messages)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, messages, 2)\n\t\tassert.Equal(t, schema.Tool, messages[0].Role)\n\t\tassert.Equal(t, \"call2\", messages[0].ToolCallID)\n\t\tassert.Contains(t, message.UserInputMultiContent[0].Text, \"stream part\")\n\t})\n}\n\nfunc TestEnhancedToolConversion(t *testing.T) {\n\tctx := context.Background()\n\n\tenhancedInvokable := &enhancedInvokableTool{\n\t\tinfo: &schema.ToolInfo{\n\t\t\tName: \"enhanced_only_invokable\",\n\t\t\tDesc: \"test enhanced invokable only\",\n\t\t},\n\t\tfn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) {\n\t\t\treturn &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t{Type: schema.ToolPartTypeText, Text: \"enhanced: \" + input.Text},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\ttoolNode, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\tTools: []tool.BaseTool{enhancedInvokable},\n\t})\n\tassert.NoError(t, err)\n\n\tt.Run(\"enhanced invokable auto-converts to streamable\", func(t *testing.T) {\n\t\tinput := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: \"call1\",\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      \"enhanced_only_invokable\",\n\t\t\t\t\tArguments: \"test\",\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\tstreamReader, err := toolNode.Stream(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, streamReader)\n\n\t\tvar messages []*schema.Message\n\t\tfor {\n\t\t\tchunk, err := streamReader.Recv()\n\t\t\tif err != nil {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif chunk != nil {\n\t\t\t\tmessages = append(messages, chunk...)\n\t\t\t}\n\t\t}\n\t\tassert.Len(t, messages, 1)\n\t})\n}\n\nfunc TestEnhancedToolMiddleware(t *testing.T) {\n\tctx := context.Background()\n\n\tvar invokableMiddlewareCalled bool\n\tvar streamableMiddlewareCalled bool\n\n\tenhancedInvokable := &enhancedInvokableTool{\n\t\tinfo: &schema.ToolInfo{\n\t\t\tName: \"enhanced_tool_with_middleware\",\n\t\t\tDesc: \"test enhanced tool with middleware\",\n\t\t},\n\t\tfn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) {\n\t\t\treturn &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t{Text: \"result\", Type: schema.ToolPartTypeText},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\ttoolNode, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\tTools: []tool.BaseTool{enhancedInvokable},\n\t\tToolCallMiddlewares: []ToolMiddleware{\n\t\t\t{\n\t\t\t\tEnhancedInvokable: func(next EnhancedInvokableToolEndpoint) EnhancedInvokableToolEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, input *ToolInput) (*EnhancedInvokableToolOutput, error) {\n\t\t\t\t\t\tinvokableMiddlewareCalled = true\n\t\t\t\t\t\treturn next(ctx, input)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t\tEnhancedStreamable: func(next EnhancedStreamableToolEndpoint) EnhancedStreamableToolEndpoint {\n\t\t\t\t\treturn func(ctx context.Context, input *ToolInput) (*EnhancedStreamableToolOutput, error) {\n\t\t\t\t\t\tstreamableMiddlewareCalled = true\n\t\t\t\t\t\treturn next(ctx, input)\n\t\t\t\t\t}\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\tt.Run(\"enhanced invokable middleware\", func(t *testing.T) {\n\t\tinvokableMiddlewareCalled = false\n\t\tinput := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: \"call1\",\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      \"enhanced_tool_with_middleware\",\n\t\t\t\t\tArguments: \"test\",\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\t_, err := toolNode.Invoke(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.True(t, invokableMiddlewareCalled)\n\t})\n\n\tt.Run(\"enhanced streamable middleware\", func(t *testing.T) {\n\t\tstreamableMiddlewareCalled = false\n\t\tinput := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: \"call2\",\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      \"enhanced_tool_with_middleware\",\n\t\t\t\t\tArguments: \"test\",\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\tstreamReader, err := toolNode.Stream(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tfor {\n\t\t\t_, err := streamReader.Recv()\n\t\t\tif err != nil {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tassert.False(t, streamableMiddlewareCalled)\n\t})\n}\n\nfunc TestEnhancedToolPriority(t *testing.T) {\n\tctx := context.Background()\n\n\tenhancedInvokable := &enhancedInvokableTool{\n\t\tinfo: &schema.ToolInfo{\n\t\t\tName: \"test_tool\",\n\t\t\tDesc: \"test tool with both enhanced and regular\",\n\t\t},\n\t\tfn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) {\n\t\t\treturn &schema.ToolResult{\n\t\t\t\tParts: []schema.ToolOutputPart{\n\t\t\t\t\t{Text: \"enhanced result\", Type: schema.ToolPartTypeText},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\ttoolNode, err := NewToolNode(ctx, &ToolsNodeConfig{\n\t\tTools: []tool.BaseTool{enhancedInvokable},\n\t})\n\tassert.NoError(t, err)\n\n\tt.Run(\"enhanced tool is used when available\", func(t *testing.T) {\n\t\tinput := schema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t{\n\t\t\t\tID: \"call1\",\n\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\tName:      \"test_tool\",\n\t\t\t\t\tArguments: \"test\",\n\t\t\t\t},\n\t\t\t},\n\t\t})\n\n\t\toutput, err := toolNode.Invoke(ctx, input)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, output, 1)\n\t\tassert.Contains(t, output[0].UserInputMultiContent[0].Text, \"enhanced result\")\n\t})\n}\n"
  },
  {
    "path": "compose/types.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"github.com/cloudwego/eino/components\"\n)\n\ntype component = components.Component\n\n// built-in component types in graph node.\n// it represents the type of the most primitive executable object provided by the user.\nconst (\n\tComponentOfUnknown     component = \"Unknown\"\n\tComponentOfGraph       component = \"Graph\"\n\tComponentOfWorkflow    component = \"Workflow\"\n\tComponentOfChain       component = \"Chain\"\n\tComponentOfPassthrough component = \"Passthrough\"\n\tComponentOfToolsNode   component = \"ToolsNode\"\n\tComponentOfLambda      component = \"Lambda\"\n)\n\n// NodeTriggerMode controls the triggering mode of graph nodes.\ntype NodeTriggerMode string\n\nconst (\n\t// AnyPredecessor means that the node will be triggered when any of its predecessors is included in the previous completed super step.\n\t// Ref:https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/#runtime-engine\n\tAnyPredecessor NodeTriggerMode = \"any_predecessor\"\n\t// AllPredecessor means that the current node will only be triggered when all of its predecessor nodes have finished running.\n\tAllPredecessor NodeTriggerMode = \"all_predecessor\"\n)\n"
  },
  {
    "path": "compose/types_composable.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"reflect\"\n)\n\n// AnyGraph the identifiers for composable and compilable Graph[I, O]、Chain[I, O] in Eino.\ntype AnyGraph interface {\n\tgetGenericHelper() *genericHelper\n\tcompile(ctx context.Context, options *graphCompileOptions) (*composableRunnable, error)\n\tinputType() reflect.Type\n\toutputType() reflect.Type\n\tcomponent() component\n}\n"
  },
  {
    "path": "compose/types_lambda.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Invoke is the type of the invokable lambda function.\ntype Invoke[I, O, TOption any] func(ctx context.Context, input I, opts ...TOption) (output O, err error)\n\n// Stream is the type of the streamable lambda function.\ntype Stream[I, O, TOption any] func(ctx context.Context,\n\tinput I, opts ...TOption) (output *schema.StreamReader[O], err error)\n\n// Collect is the type of the collectable lambda function.\ntype Collect[I, O, TOption any] func(ctx context.Context,\n\tinput *schema.StreamReader[I], opts ...TOption) (output O, err error)\n\n// Transform is the type of the transformable lambda function.\ntype Transform[I, O, TOption any] func(ctx context.Context,\n\tinput *schema.StreamReader[I], opts ...TOption) (output *schema.StreamReader[O], err error)\n\n// InvokeWOOpt is the type of the invokable lambda function without options.\ntype InvokeWOOpt[I, O any] func(ctx context.Context, input I) (output O, err error)\n\n// StreamWOOpt is the type of the streamable lambda function without options.\ntype StreamWOOpt[I, O any] func(ctx context.Context,\n\tinput I) (output *schema.StreamReader[O], err error)\n\n// CollectWOOpt is the type of the collectable lambda function without options.\ntype CollectWOOpt[I, O any] func(ctx context.Context,\n\tinput *schema.StreamReader[I]) (output O, err error)\n\n// TransformWOOpts is the type of the transformable lambda function without options.\ntype TransformWOOpts[I, O any] func(ctx context.Context,\n\tinput *schema.StreamReader[I]) (output *schema.StreamReader[O], err error)\n\n// Lambda is the node that wraps the user provided lambda function.\n// It can be used as a node in Graph or Chain (include Parallel and Branch).\n// Create a Lambda by using AnyLambda/InvokableLambda/StreamableLambda/CollectableLambda/TransformableLambda.\n// eg.\n//\n//\tlambda := compose.InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n//\t\treturn input, nil\n//\t})\ntype Lambda struct {\n\texecutor *composableRunnable\n}\n\ntype lambdaOpts struct {\n\t// same as executorMeta.isComponentCallbackEnabled\n\t// indicates whether the executable lambda user provided could execute the callback aspect itself.\n\t// if it could, the callback in the corresponding graph node won't be executed anymore\n\tenableComponentCallback bool\n\n\t// same as executorMeta.componentImplType\n\t// for AnyLambda, the value comes from the user's explicit config\n\t// if componentImplType is empty, then the class name or func name in the instance will be inferred, but no guarantee.\n\tcomponentImplType string\n}\n\n// LambdaOpt is the option for creating a Lambda.\ntype LambdaOpt func(o *lambdaOpts)\n\n// WithLambdaCallbackEnable enables the callback aspect of the lambda function.\nfunc WithLambdaCallbackEnable(y bool) LambdaOpt {\n\treturn func(o *lambdaOpts) {\n\t\to.enableComponentCallback = y\n\t}\n}\n\n// WithLambdaType sets the type of the lambda function.\nfunc WithLambdaType(t string) LambdaOpt {\n\treturn func(o *lambdaOpts) {\n\t\to.componentImplType = t\n\t}\n}\n\ntype unreachableOption struct{}\n\n// InvokableLambdaWithOption creates a Lambda with invokable lambda function and options.\nfunc InvokableLambdaWithOption[I, O, TOption any](i Invoke[I, O, TOption], opts ...LambdaOpt) *Lambda {\n\treturn anyLambda(i, nil, nil, nil, opts...)\n}\n\n// InvokableLambda creates a Lambda with invokable lambda function without options.\nfunc InvokableLambda[I, O any](i InvokeWOOpt[I, O], opts ...LambdaOpt) *Lambda {\n\tf := func(ctx context.Context, input I, opts_ ...unreachableOption) (output O, err error) {\n\t\treturn i(ctx, input)\n\t}\n\n\treturn anyLambda(f, nil, nil, nil, opts...)\n}\n\n// StreamableLambdaWithOption creates a Lambda with streamable lambda function and options.\nfunc StreamableLambdaWithOption[I, O, TOption any](s Stream[I, O, TOption], opts ...LambdaOpt) *Lambda {\n\treturn anyLambda(nil, s, nil, nil, opts...)\n}\n\n// StreamableLambda creates a Lambda with streamable lambda function without options.\nfunc StreamableLambda[I, O any](s StreamWOOpt[I, O], opts ...LambdaOpt) *Lambda {\n\tf := func(ctx context.Context, input I, opts_ ...unreachableOption) (\n\t\toutput *schema.StreamReader[O], err error) {\n\n\t\treturn s(ctx, input)\n\t}\n\n\treturn anyLambda(nil, f, nil, nil, opts...)\n}\n\n// CollectableLambdaWithOption creates a Lambda with collectable lambda function and options.\nfunc CollectableLambdaWithOption[I, O, TOption any](c Collect[I, O, TOption], opts ...LambdaOpt) *Lambda {\n\treturn anyLambda(nil, nil, c, nil, opts...)\n}\n\n// CollectableLambda creates a Lambda with collectable lambda function without options.\nfunc CollectableLambda[I, O any](c CollectWOOpt[I, O], opts ...LambdaOpt) *Lambda {\n\tf := func(ctx context.Context, input *schema.StreamReader[I],\n\t\topts_ ...unreachableOption) (output O, err error) {\n\n\t\treturn c(ctx, input)\n\t}\n\n\treturn anyLambda(nil, nil, f, nil, opts...)\n}\n\n// TransformableLambdaWithOption creates a Lambda with transformable lambda function and options.\nfunc TransformableLambdaWithOption[I, O, TOption any](t Transform[I, O, TOption], opts ...LambdaOpt) *Lambda {\n\treturn anyLambda(nil, nil, nil, t, opts...)\n}\n\n// TransformableLambda creates a Lambda with transformable lambda function without options.\nfunc TransformableLambda[I, O any](t TransformWOOpts[I, O], opts ...LambdaOpt) *Lambda {\n\n\tf := func(ctx context.Context, input *schema.StreamReader[I],\n\t\topts_ ...unreachableOption) (output *schema.StreamReader[O], err error) {\n\n\t\treturn t(ctx, input)\n\t}\n\n\treturn anyLambda(nil, nil, nil, f, opts...)\n}\n\n// AnyLambda creates a Lambda with any lambda function.\n// you can only implement one or more of the four lambda functions, and the rest use nil.\n// eg.\n//\n//\tinvokeFunc := func(ctx context.Context, input string, opts ...myOption) (output string, err error) {\n//\t\t// ...\n//\t}\n//\tstreamFunc := func(ctx context.Context, input string, opts ...myOption) (output *schema.StreamReader[string], err error) {\n//\t\t// ...\n//\t}\n//\n// lambda := compose.AnyLambda(invokeFunc, streamFunc, nil, nil)\nfunc AnyLambda[I, O, TOption any](i Invoke[I, O, TOption], s Stream[I, O, TOption],\n\tc Collect[I, O, TOption], t Transform[I, O, TOption], opts ...LambdaOpt) (*Lambda, error) {\n\n\tif i == nil && s == nil && c == nil && t == nil {\n\t\treturn nil, fmt.Errorf(\"needs to have at least one of four lambda types: invoke/stream/collect/transform, got none\")\n\t}\n\n\treturn anyLambda(i, s, c, t, opts...), nil\n}\n\nfunc anyLambda[I, O, TOption any](i Invoke[I, O, TOption], s Stream[I, O, TOption],\n\tc Collect[I, O, TOption], t Transform[I, O, TOption], opts ...LambdaOpt) *Lambda {\n\n\topt := getLambdaOpt(opts...)\n\n\texecutor := runnableLambda(i, s, c, t,\n\t\t!opt.enableComponentCallback,\n\t)\n\texecutor.meta = &executorMeta{\n\t\tcomponent:                  ComponentOfLambda,\n\t\tisComponentCallbackEnabled: opt.enableComponentCallback,\n\t\tcomponentImplType:          opt.componentImplType,\n\t}\n\n\treturn &Lambda{\n\t\texecutor: executor,\n\t}\n}\n\nfunc getLambdaOpt(opts ...LambdaOpt) *lambdaOpts {\n\topt := &lambdaOpts{\n\t\tenableComponentCallback: false,\n\t\tcomponentImplType:       \"\",\n\t}\n\n\tfor _, optFn := range opts {\n\t\toptFn(opt)\n\t}\n\treturn opt\n}\n\n// ToList creates a Lambda that converts input I to a []I.\n// It's useful when you want to convert a single input to a list of inputs.\n// eg.\n//\n//\tlambda := compose.ToList[*schema.Message]()\n//\tchain := compose.NewChain[[]*schema.Message, []*schema.Message]()\n//\n//\tchain.AddChatModel(chatModel) // chatModel returns *schema.Message, but we need []*schema.Message\n//\tchain.AddLambda(lambda) // convert *schema.Message to []*schema.Message\nfunc ToList[I any](opts ...LambdaOpt) *Lambda {\n\ti := func(ctx context.Context, input I, opts_ ...unreachableOption) (output []I, err error) {\n\t\treturn []I{input}, nil\n\t}\n\n\tf := func(ctx context.Context, inputS *schema.StreamReader[I], opts_ ...unreachableOption) (outputS *schema.StreamReader[[]I], err error) {\n\t\treturn schema.StreamReaderWithConvert(inputS, func(i I) ([]I, error) {\n\t\t\treturn []I{i}, nil\n\t\t}), nil\n\t}\n\n\treturn anyLambda(i, nil, nil, f, opts...)\n}\n\n// MessageParser creates a lambda that parses a message into an object T, usually used after a chatmodel.\n// usage:\n//\n//\tparser := schema.NewMessageJSONParser[MyStruct](&schema.MessageJSONParseConfig{\n//\t\tParseFrom: schema.MessageParseFromContent,\n//\t})\n//\tparserLambda := MessageParser(parser)\n//\n//\tchain := NewChain[*schema.Message, MyStruct]()\n//\tchain.AppendChatModel(chatModel)\n//\tchain.AppendLambda(parserLambda)\n//\n//\tr, err := chain.Compile(context.Background())\n//\n//\t// parsed is a MyStruct object\n//\tparsed, err := r.Invoke(context.Background(), &schema.Message{\n//\t\tRole:    schema.MessageRoleUser,\n//\t\tContent: \"return a json string for my struct\",\n//\t})\nfunc MessageParser[T any](p schema.MessageParser[T], opts ...LambdaOpt) *Lambda {\n\ti := func(ctx context.Context, input *schema.Message, opts_ ...unreachableOption) (output T, err error) {\n\t\treturn p.Parse(ctx, input)\n\t}\n\n\topts = append([]LambdaOpt{WithLambdaType(\"MessageParse\")}, opts...)\n\n\treturn anyLambda(i, nil, nil, nil, opts...)\n}\n"
  },
  {
    "path": "compose/types_lambda_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestLambda(t *testing.T) {\n\tt.Run(\"InvokableLambda\", func(t *testing.T) {\n\t\tld := InvokableLambdaWithOption(\n\t\t\tfunc(ctx context.Context, input string, opts ...any) (output string, err error) {\n\t\t\t\treturn \"good\", nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(false),\n\t\t\tWithLambdaType(\"ForTest\"),\n\t\t)\n\n\t\tassert.Equal(t, false, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"ForTest\", ld.executor.meta.componentImplType)\n\n\t\tld = InvokableLambda(\n\t\t\tfunc(ctx context.Context, input string) (output string, err error) {\n\t\t\t\treturn \"good\", nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(false),\n\t\t\tWithLambdaType(\"ForTest\"),\n\t\t)\n\n\t\tassert.Equal(t, false, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"ForTest\", ld.executor.meta.componentImplType)\n\t})\n\n\tt.Run(\"StreamableLambda\", func(t *testing.T) {\n\t\tld := StreamableLambdaWithOption(\n\t\t\tfunc(ctx context.Context, input string, opts ...any) (output *schema.StreamReader[string], err error) {\n\t\t\t\tsr, sw := schema.Pipe[string](1)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(false),\n\t\t\tWithLambdaType(\"ForTest\"),\n\t\t)\n\n\t\tassert.Equal(t, false, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"ForTest\", ld.executor.meta.componentImplType)\n\n\t\tld = StreamableLambda(\n\t\t\tfunc(ctx context.Context, input string) (output *schema.StreamReader[string], err error) {\n\t\t\t\tsr, sw := schema.Pipe[string](1)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(false),\n\t\t\tWithLambdaType(\"ForTest\"),\n\t\t)\n\n\t\tassert.Equal(t, false, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"ForTest\", ld.executor.meta.componentImplType)\n\t})\n\n\tt.Run(\"CollectableLambda\", func(t *testing.T) {\n\t\tld := CollectableLambdaWithOption(\n\t\t\tfunc(ctx context.Context, input *schema.StreamReader[string], opts ...any) (output string, err error) {\n\t\t\t\treturn \"good\", nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(true),\n\t\t)\n\n\t\tassert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"\", ld.executor.meta.componentImplType)\n\n\t\tld = CollectableLambda(\n\t\t\tfunc(ctx context.Context, input *schema.StreamReader[string]) (output string, err error) {\n\t\t\t\treturn \"good\", nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(true),\n\t\t)\n\n\t\tassert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"\", ld.executor.meta.componentImplType)\n\t})\n\n\tt.Run(\"TransformableLambda\", func(t *testing.T) {\n\t\tld := TransformableLambdaWithOption(\n\t\t\tfunc(ctx context.Context, input *schema.StreamReader[string], opts ...any) (output *schema.StreamReader[string], err error) {\n\t\t\t\tsr, sw := schema.Pipe[string](1)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(true),\n\t\t)\n\n\t\tassert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"\", ld.executor.meta.componentImplType)\n\n\t\tld = TransformableLambda(\n\t\t\tfunc(ctx context.Context, input *schema.StreamReader[string]) (output *schema.StreamReader[string], err error) {\n\t\t\t\tsr, sw := schema.Pipe[string](1)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(true),\n\t\t)\n\n\t\tassert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"\", ld.executor.meta.componentImplType)\n\t})\n\n\tt.Run(\"AnyLambda\", func(t *testing.T) {\n\t\tld, err := AnyLambda[string, string](\n\t\t\tfunc(ctx context.Context, input string, opts ...any) (output string, err error) {\n\t\t\t\treturn \"good\", nil\n\t\t\t},\n\t\t\tfunc(ctx context.Context, input string, opts ...any) (output *schema.StreamReader[string], err error) {\n\t\t\t\tsr, sw := schema.Pipe[string](1)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t},\n\t\t\tfunc(ctx context.Context, input *schema.StreamReader[string], opts ...any) (output string, err error) {\n\t\t\t\treturn \"good\", nil\n\t\t\t},\n\t\t\tfunc(ctx context.Context, input *schema.StreamReader[string], opts ...any) (output *schema.StreamReader[string], err error) {\n\t\t\t\tsr, sw := schema.Pipe[string](1)\n\t\t\t\tsw.Close()\n\t\t\t\treturn sr, nil\n\t\t\t},\n\t\t\tWithLambdaCallbackEnable(true),\n\t\t\tWithLambdaType(\"ForTest\"),\n\t\t)\n\t\tassert.NoError(t, err)\n\n\t\tassert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled)\n\t\tassert.Equal(t, ComponentOfLambda, ld.executor.meta.component)\n\t\tassert.Equal(t, \"ForTest\", ld.executor.meta.componentImplType)\n\t})\n}\n\ntype TestStructForParse struct {\n\tID int `json:\"id\"`\n}\n\nfunc TestMessageParser(t *testing.T) {\n\tt.Run(\"parse from content\", func(t *testing.T) {\n\t\tparser := schema.NewMessageJSONParser[TestStructForParse](&schema.MessageJSONParseConfig{\n\t\t\tParseFrom: schema.MessageParseFromContent,\n\t\t})\n\n\t\tparserLambda := MessageParser(parser)\n\n\t\tchain := NewChain[*schema.Message, TestStructForParse]()\n\t\tchain.AppendLambda(parserLambda)\n\n\t\tr, err := chain.Compile(context.Background())\n\t\tassert.Nil(t, err)\n\n\t\tparsed, err := r.Invoke(context.Background(), &schema.Message{\n\t\t\tContent: `{\"id\": 1}`,\n\t\t})\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, 1, parsed.ID)\n\t})\n\n\tt.Run(\"parse from tool call\", func(t *testing.T) {\n\t\tparser := schema.NewMessageJSONParser[*TestStructForParse](&schema.MessageJSONParseConfig{\n\t\t\tParseFrom: schema.MessageParseFromToolCall,\n\t\t})\n\n\t\tparserLambda := MessageParser(parser)\n\n\t\tchain := NewChain[*schema.Message, *TestStructForParse]()\n\t\tchain.AppendLambda(parserLambda)\n\n\t\tr, err := chain.Compile(context.Background())\n\t\tassert.Nil(t, err)\n\n\t\tparsed, err := r.Invoke(context.Background(), &schema.Message{\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{Function: schema.FunctionCall{Arguments: `{\"id\": 1}`}},\n\t\t\t},\n\t\t})\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, 1, parsed.ID)\n\t})\n}\n"
  },
  {
    "path": "compose/utils.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\ticb \"github.com/cloudwego/eino/internal/callbacks\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype on[T any] func(context.Context, T) (context.Context, T)\n\nfunc onStart[T any](ctx context.Context, input T) (context.Context, T) {\n\treturn icb.On(ctx, input, icb.OnStartHandle[T], callbacks.TimingOnStart, true)\n}\n\nfunc onEnd[T any](ctx context.Context, output T) (context.Context, T) {\n\treturn icb.On(ctx, output, icb.OnEndHandle[T], callbacks.TimingOnEnd, false)\n}\n\nfunc onStartWithStreamInput[T any](ctx context.Context, input *schema.StreamReader[T]) (\n\tcontext.Context, *schema.StreamReader[T]) {\n\n\treturn icb.On(ctx, input, icb.OnStartWithStreamInputHandle[T], callbacks.TimingOnStartWithStreamInput, true)\n}\n\nfunc genericOnStartWithStreamInputHandle(ctx context.Context, input streamReader,\n\trunInfo *icb.RunInfo, handlers []icb.Handler) (context.Context, streamReader) {\n\n\thandlers = generic.Reverse(handlers)\n\n\tcpy := input.copy\n\n\thandle := func(ctx context.Context, handler icb.Handler, in streamReader) context.Context {\n\t\tin_, ok := unpackStreamReader[icb.CallbackInput](in)\n\t\tif !ok {\n\t\t\tpanic(\"impossible\")\n\t\t}\n\n\t\treturn handler.OnStartWithStreamInput(ctx, runInfo, in_)\n\t}\n\n\treturn icb.OnWithStreamHandle(ctx, input, handlers, cpy, handle)\n}\n\nfunc genericOnStartWithStreamInput(ctx context.Context, input streamReader) (context.Context, streamReader) {\n\treturn icb.On(ctx, input, genericOnStartWithStreamInputHandle, callbacks.TimingOnStartWithStreamInput, true)\n}\n\nfunc onEndWithStreamOutput[T any](ctx context.Context, output *schema.StreamReader[T]) (\n\tcontext.Context, *schema.StreamReader[T]) {\n\n\treturn icb.On(ctx, output, icb.OnEndWithStreamOutputHandle[T], callbacks.TimingOnEndWithStreamOutput, false)\n}\n\nfunc genericOnEndWithStreamOutputHandle(ctx context.Context, output streamReader,\n\trunInfo *icb.RunInfo, handlers []icb.Handler) (context.Context, streamReader) {\n\n\tcpy := output.copy\n\n\thandle := func(ctx context.Context, handler icb.Handler, out streamReader) context.Context {\n\t\tout_, ok := unpackStreamReader[icb.CallbackOutput](out)\n\t\tif !ok {\n\t\t\tpanic(\"impossible\")\n\t\t}\n\n\t\treturn handler.OnEndWithStreamOutput(ctx, runInfo, out_)\n\t}\n\n\treturn icb.OnWithStreamHandle(ctx, output, handlers, cpy, handle)\n}\n\nfunc genericOnEndWithStreamOutput(ctx context.Context, output streamReader) (context.Context, streamReader) {\n\treturn icb.On(ctx, output, genericOnEndWithStreamOutputHandle, callbacks.TimingOnEndWithStreamOutput, false)\n}\n\nfunc onError(ctx context.Context, err error) (context.Context, error) {\n\treturn icb.On(ctx, err, icb.OnErrorHandle, callbacks.TimingOnError, false)\n}\n\nfunc runWithCallbacks[I, O, TOption any](r func(context.Context, I, ...TOption) (O, error),\n\tonStart on[I], onEnd on[O], onError on[error]) func(context.Context, I, ...TOption) (O, error) {\n\n\treturn func(ctx context.Context, input I, opts ...TOption) (output O, err error) {\n\t\tctx, input = onStart(ctx, input)\n\n\t\toutput, err = r(ctx, input, opts...)\n\t\tif err != nil {\n\t\t\tctx, err = onError(ctx, err)\n\t\t\treturn output, err\n\t\t}\n\n\t\tctx, output = onEnd(ctx, output)\n\n\t\treturn output, nil\n\t}\n}\n\nfunc invokeWithCallbacks[I, O, TOption any](i Invoke[I, O, TOption]) Invoke[I, O, TOption] {\n\treturn runWithCallbacks(i, onStart[I], onEnd[O], onError)\n}\n\nfunc onGraphStart(ctx context.Context, input any, isStream bool) (context.Context, any) {\n\tif isStream {\n\t\treturn genericOnStartWithStreamInput(ctx, input.(streamReader))\n\t}\n\treturn onStart(ctx, input)\n}\n\nfunc onGraphEnd(ctx context.Context, output any, isStream bool) (context.Context, any) {\n\tif isStream {\n\t\treturn genericOnEndWithStreamOutput(ctx, output.(streamReader))\n\t}\n\treturn onEnd(ctx, output)\n}\n\nfunc onGraphError(ctx context.Context, err error) (context.Context, error) {\n\treturn onError(ctx, err)\n}\n\nfunc streamWithCallbacks[I, O, TOption any](s Stream[I, O, TOption]) Stream[I, O, TOption] {\n\treturn runWithCallbacks(s, onStart[I], onEndWithStreamOutput[O], onError)\n}\n\nfunc collectWithCallbacks[I, O, TOption any](c Collect[I, O, TOption]) Collect[I, O, TOption] {\n\treturn runWithCallbacks(c, onStartWithStreamInput[I], onEnd[O], onError)\n}\n\nfunc transformWithCallbacks[I, O, TOption any](t Transform[I, O, TOption]) Transform[I, O, TOption] {\n\treturn runWithCallbacks(t, onStartWithStreamInput[I], onEndWithStreamOutput[O], onError)\n}\n\nfunc initGraphCallbacks(ctx context.Context, info *nodeInfo, meta *executorMeta, opts ...Option) context.Context {\n\tri := &callbacks.RunInfo{}\n\tif meta != nil {\n\t\tri.Component = meta.component\n\t\tri.Type = meta.componentImplType\n\t}\n\n\tif info != nil {\n\t\tri.Name = info.name\n\t}\n\n\tvar cbs []callbacks.Handler\n\tfor i := range opts {\n\t\tif len(opts[i].handler) != 0 && len(opts[i].paths) == 0 {\n\t\t\tcbs = append(cbs, opts[i].handler...)\n\t\t}\n\t}\n\n\tif len(cbs) == 0 {\n\t\treturn icb.ReuseHandlers(ctx, ri)\n\t}\n\n\treturn icb.AppendHandlers(ctx, ri, cbs...)\n}\n\nfunc initNodeCallbacks(ctx context.Context, key string, info *nodeInfo, meta *executorMeta, opts ...Option) context.Context {\n\tri := &callbacks.RunInfo{}\n\tif meta != nil {\n\t\tri.Component = meta.component\n\t\tri.Type = meta.componentImplType\n\t}\n\n\tif info != nil {\n\t\tri.Name = info.name\n\t}\n\n\tvar cbs []callbacks.Handler\n\tfor i := range opts {\n\t\tif len(opts[i].handler) != 0 {\n\t\t\tif len(opts[i].paths) != 0 {\n\t\t\t\tfor _, k := range opts[i].paths {\n\t\t\t\t\tif len(k.path) == 1 && k.path[0] == key {\n\t\t\t\t\t\tcbs = append(cbs, opts[i].handler...)\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tif len(cbs) == 0 {\n\t\treturn icb.ReuseHandlers(ctx, ri)\n\t}\n\n\treturn icb.AppendHandlers(ctx, ri, cbs...)\n}\n\nfunc streamChunkConvertForCBOutput[O any](o O) (callbacks.CallbackOutput, error) {\n\treturn o, nil\n}\n\nfunc streamChunkConvertForCBInput[I any](i I) (callbacks.CallbackInput, error) {\n\treturn i, nil\n}\n\nfunc toAnyList[T any](in []T) []any {\n\tret := make([]any, len(in))\n\tfor i := range in {\n\t\tret[i] = in[i]\n\t}\n\treturn ret\n}\n\ntype assignableType uint8\n\nconst (\n\tassignableTypeMustNot assignableType = iota\n\tassignableTypeMust\n\tassignableTypeMay\n)\n\nfunc checkAssignable(input, arg reflect.Type) assignableType {\n\tif arg == nil || input == nil {\n\t\treturn assignableTypeMustNot\n\t}\n\n\tif arg == input {\n\t\treturn assignableTypeMust\n\t}\n\n\tif arg.Kind() == reflect.Interface && input.Implements(arg) {\n\t\treturn assignableTypeMust\n\t}\n\tif input.Kind() == reflect.Interface {\n\t\tif arg.Implements(input) {\n\t\t\treturn assignableTypeMay\n\t\t}\n\t\treturn assignableTypeMustNot\n\t}\n\n\treturn assignableTypeMustNot\n}\n\nfunc extractOption(nodes map[string]*chanCall, opts ...Option) (map[string][]any, error) {\n\toptMap := map[string][]any{}\n\tfor _, opt := range opts {\n\t\tif len(opt.paths) == 0 {\n\t\t\t// common, discard callback, filter option by type\n\t\t\tif len(opt.options) == 0 {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tfor name, c := range nodes {\n\t\t\t\tif c.action.optionType == nil {\n\t\t\t\t\t// subgraph\n\t\t\t\t\toptMap[name] = append(optMap[name], opt)\n\t\t\t\t} else if reflect.TypeOf(opt.options[0]) == c.action.optionType { // assume that types of options are the same\n\t\t\t\t\toptMap[name] = append(optMap[name], opt.options...)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t\tfor _, path := range opt.paths {\n\t\t\tif len(path.path) == 0 {\n\t\t\t\treturn nil, fmt.Errorf(\"call option has designated an empty path\")\n\t\t\t}\n\n\t\t\tvar curNode *chanCall\n\t\t\tvar ok bool\n\t\t\tif curNode, ok = nodes[path.path[0]]; !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"option has designated an unknown node: %s\", path)\n\t\t\t}\n\t\t\tcurNodeKey := path.path[0]\n\n\t\t\tif len(path.path) == 1 {\n\t\t\t\tif len(opt.options) == 0 {\n\t\t\t\t\t// sub graph common callbacks has been added to ctx in initNodeCallback and won't be passed to subgraph only pass options\n\t\t\t\t\t// node callback also won't be passed\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\t\t\t\tif curNode.action.optionType == nil {\n\t\t\t\t\tnOpt := opt.deepCopy()\n\t\t\t\t\tnOpt.paths = []*NodePath{}\n\t\t\t\t\toptMap[curNodeKey] = append(optMap[curNodeKey], nOpt)\n\t\t\t\t} else {\n\t\t\t\t\t// designate to component\n\t\t\t\t\tif curNode.action.optionType != reflect.TypeOf(opt.options[0]) { // assume that types of options are the same\n\t\t\t\t\t\treturn nil, fmt.Errorf(\"option type[%s] is different from which the designated node[%s] expects[%s]\",\n\t\t\t\t\t\t\treflect.TypeOf(opt.options[0]).String(), path, curNode.action.optionType.String())\n\t\t\t\t\t}\n\t\t\t\t\toptMap[curNodeKey] = append(optMap[curNodeKey], opt.options...)\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tif curNode.action.optionType != nil {\n\t\t\t\t\t// component\n\t\t\t\t\treturn nil, fmt.Errorf(\"cannot designate sub path of a component, path:%s\", path)\n\t\t\t\t}\n\t\t\t\t// designate to sub graph's nodes\n\t\t\t\tnOpt := opt.deepCopy()\n\t\t\t\tnOpt.paths = []*NodePath{NewNodePath(path.path[1:]...)}\n\t\t\t\toptMap[curNodeKey] = append(optMap[curNodeKey], nOpt)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn optMap, nil\n}\n\nfunc mapToList(m map[string]any) []any {\n\tret := make([]any, 0, len(m))\n\tfor _, v := range m {\n\t\tret = append(ret, v)\n\t}\n\treturn ret\n}\n"
  },
  {
    "path": "compose/utils_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\ntype good interface {\n\tThisIsGood() bool\n}\n\ntype good2 interface {\n\tThisIsGood2() bool\n}\n\ntype good3 interface {\n\tThisIsGood() bool\n}\n\ntype goodImpl struct{}\n\nfunc (g *goodImpl) ThisIsGood() bool {\n\treturn true\n}\n\ntype goodNotImpl struct{}\n\nfunc TestValidateType(t *testing.T) {\n\n\tt.Run(\"equal_type\", func(t *testing.T) {\n\t\targ := generic.TypeOf[int]()\n\t\tinput := generic.TypeOf[int]()\n\n\t\tresult := checkAssignable(input, arg)\n\t\tassert.Equal(t, assignableTypeMust, result)\n\t})\n\n\tt.Run(\"unequal_type\", func(t *testing.T) {\n\t\targ := generic.TypeOf[int]()\n\t\tinput := generic.TypeOf[string]()\n\n\t\tresult := checkAssignable(input, arg)\n\t\tassert.Equal(t, assignableTypeMustNot, result)\n\t})\n\n\tt.Run(\"implement_interface\", func(t *testing.T) {\n\t\targ := generic.TypeOf[good]()\n\t\tinput := generic.TypeOf[*goodImpl]()\n\n\t\tresult := checkAssignable(input, arg)\n\t\tassert.Equal(t, assignableTypeMust, result)\n\t})\n\n\tt.Run(\"may_implement_interface\", func(t *testing.T) {\n\t\targ := generic.TypeOf[*goodImpl]()\n\t\tinput := generic.TypeOf[good]()\n\n\t\tresult := checkAssignable(input, arg)\n\t\tassert.Equal(t, assignableTypeMay, result)\n\t})\n\n\tt.Run(\"not_implement_interface\", func(t *testing.T) {\n\t\targ := generic.TypeOf[good]()\n\t\tinput := generic.TypeOf[*goodNotImpl]()\n\n\t\tresult := checkAssignable(input, arg)\n\t\tassert.Equal(t, assignableTypeMustNot, result)\n\t})\n\n\tt.Run(\"interface_unequal_interface\", func(t *testing.T) {\n\t\targ := generic.TypeOf[good]()\n\t\tinput := generic.TypeOf[good2]()\n\n\t\tresult := checkAssignable(input, arg)\n\t\tassert.Equal(t, assignableTypeMustNot, result)\n\t})\n\n\tt.Run(\"interface_equal_interface\", func(t *testing.T) {\n\t\targ := generic.TypeOf[good]()\n\t\tinput := generic.TypeOf[good3]()\n\n\t\tresult := checkAssignable(input, arg)\n\t\tassert.Equal(t, assignableTypeMust, result)\n\t})\n}\n\nfunc TestStreamChunkConvert(t *testing.T) {\n\to, err := streamChunkConvertForCBOutput(1)\n\tassert.Nil(t, err)\n\tassert.Equal(t, o, 1)\n\n\ti, err := streamChunkConvertForCBInput(1)\n\tassert.Nil(t, err)\n\tassert.Equal(t, i, 1)\n}\n"
  },
  {
    "path": "compose/values_merge.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal\"\n)\n\n// RegisterValuesMergeFunc registers a function to merge outputs from multiple nodes when fan-in.\n// It's used to define how to merge for a specific type.\n// For maps that already have a default merge function, you don't need to register a new one unless you want to customize the merge logic.\nfunc RegisterValuesMergeFunc[T any](fn func([]T) (T, error)) {\n\tinternal.RegisterValuesMergeFunc(fn)\n}\n\ntype mergeOptions struct {\n\tstreamMergeWithSourceEOF bool\n\tnames                    []string\n}\n\n// the caller should ensure len(vs) > 1\nfunc mergeValues(vs []any, opts *mergeOptions) (any, error) {\n\tv0 := reflect.ValueOf(vs[0])\n\tt0 := v0.Type()\n\n\tif fn := internal.GetMergeFunc(t0); fn != nil {\n\t\treturn fn(vs)\n\t}\n\n\t// merge StreamReaders\n\tif s, ok := vs[0].(streamReader); ok {\n\t\tt := s.getChunkType()\n\t\tif internal.GetMergeFunc(t) == nil {\n\t\t\treturn nil, fmt.Errorf(\"(mergeValues | stream type)\"+\n\t\t\t\t\" unsupported chunk type: %v\", t)\n\t\t}\n\n\t\tss := make([]streamReader, len(vs)-1)\n\t\tfor i := 0; i < len(ss); i++ {\n\t\t\tsri, ok_ := vs[i+1].(streamReader)\n\t\t\tif !ok_ {\n\t\t\t\treturn nil, fmt.Errorf(\"(mergeStream) unexpected type. \"+\n\t\t\t\t\t\"expect: %v, got: %v\", t0, reflect.TypeOf(vs[i]))\n\t\t\t}\n\n\t\t\tif st := sri.getChunkType(); st != t {\n\t\t\t\treturn nil, fmt.Errorf(\"(mergeStream) chunk type mismatch. \"+\n\t\t\t\t\t\"expect: %v, got: %v\", t, st)\n\t\t\t}\n\n\t\t\tss[i] = sri\n\t\t}\n\n\t\tif opts != nil && opts.streamMergeWithSourceEOF {\n\t\t\tms := s.mergeWithNames(ss, opts.names)\n\t\t\treturn ms, nil\n\t\t}\n\n\t\tms := s.merge(ss)\n\n\t\treturn ms, nil\n\t}\n\n\treturn nil, fmt.Errorf(\"(mergeValues) unsupported type: %v\", t0)\n}\n"
  },
  {
    "path": "compose/values_merge_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"fmt\"\n\t\"io\"\n\t\"sort\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc Test_mergeValues(t *testing.T) {\n\tt.Run(\"merge maps\", func(t *testing.T) {\n\t\tm1 := map[int]int{1: 1, 2: 2, 3: 3, 4: 4}\n\t\tm2 := map[int]int{5: 5, 6: 6, 7: 7, 8: 8}\n\t\tm3 := map[int]int{9: 9, 10: 10, 11: 11}\n\n\t\tt.Run(\"regular\", func(t *testing.T) {\n\t\t\tmergedM, err := mergeValues([]any{m1, m2, m3}, nil)\n\t\t\tassert.NoError(t, err)\n\n\t\t\tm := mergedM.(map[int]int)\n\n\t\t\t// len(m) == len(m1) + len(m2) + len(m3)\n\t\t\tassert.Equal(t, len(m), len(m1)+len(m2)+len(m3))\n\t\t})\n\n\t\tt.Run(\"duplicated key\", func(t *testing.T) {\n\t\t\t_, err := mergeValues([]any{m1, m2, m3, map[int]int{1: 1}}, nil)\n\t\t\tassert.ErrorContains(t, err, \"duplicated key\")\n\t\t})\n\n\t\tt.Run(\"type mismatch\", func(t *testing.T) {\n\t\t\t_, err := mergeValues([]any{m1, m2, m3, map[int]string{1: \"1\"}}, nil)\n\t\t\tassert.ErrorContains(t, err, \"type mismatch\")\n\t\t})\n\t})\n\n\tt.Run(\"merge stream\", func(t *testing.T) {\n\t\tass := []any{\n\t\t\tpackStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{1: \"1\"}})),\n\t\t\tpackStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{2: \"2\"}})),\n\t\t\tpackStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{3: \"3\", 4: \"4\"}})),\n\t\t}\n\t\tisr, err := mergeValues(ass, nil)\n\t\trequire.NoError(t, err)\n\t\tret, ok := unpackStreamReader[map[int]string](isr.(streamReader))\n\t\trequire.True(t, ok)\n\t\tdefer ret.Close()\n\n\t\tgot := make(map[int]string)\n\t\tfor i := 0; i < 3; i++ {\n\t\t\tm, err := ret.Recv()\n\t\t\trequire.NoError(t, err)\n\t\t\tfor k, v := range m {\n\t\t\t\tgot[k] = v\n\t\t\t}\n\t\t}\n\t\t_, err = ret.Recv()\n\t\trequire.ErrorIs(t, err, io.EOF)\n\n\t\tassert.Equal(t, map[int]string{\n\t\t\t1: \"1\",\n\t\t\t2: \"2\",\n\t\t\t3: \"3\",\n\t\t\t4: \"4\",\n\t\t}, got)\n\t})\n\n\tt.Run(\"merge stream with source EOF\", func(t *testing.T) {\n\t\tass := []any{\n\t\t\tpackStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{1: \"1\"}})),\n\t\t\tpackStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{2: \"2\"}})),\n\t\t\tpackStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{3: \"3\", 4: \"4\"}})),\n\t\t}\n\t\topts := &mergeOptions{\n\t\t\tstreamMergeWithSourceEOF: true,\n\t\t\tnames: []string{\n\t\t\t\t\"source0\",\n\t\t\t\t\"source1\",\n\t\t\t\t\"source2\",\n\t\t\t},\n\t\t}\n\t\tisr, err := mergeValues(ass, opts)\n\t\trequire.NoError(t, err)\n\t\tret, ok := unpackStreamReader[map[int]string](isr.(streamReader))\n\t\trequire.True(t, ok)\n\t\tdefer ret.Close()\n\n\t\tgot := make(map[int]string)\n\t\tendedSources := make(map[string]bool)\n\n\t\tfor {\n\t\t\tm, e := ret.Recv()\n\t\t\tif e != nil {\n\t\t\t\tif sourceName, ok_ := schema.GetSourceName(e); ok_ {\n\t\t\t\t\tt.Logf(\"Source '%s' ended\", sourceName)\n\t\t\t\t\tendedSources[sourceName] = true\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\t\t\t\tif e == io.EOF {\n\t\t\t\t\t// This EOF means all chunks from all sources that were not SourceEOF have been merged and sent.\n\t\t\t\t\t// Or, if all sources send SourceEOF first, this io.EOF means the merged stream itself is now empty.\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\trequire.NoError(t, e) // Fail on any other error\n\t\t\t}\n\t\t\t// If streamMergeWithSourceEOF is true, the final merged result comes as a single map chunk\n\t\t\t// after all SourceEOFs (if any non-empty streams existed) or directly if all streams were empty.\n\t\t\tfor k, v := range m {\n\t\t\t\tgot[k] = v\n\t\t\t}\n\t\t}\n\n\t\t// Check that all expected sources have ended if they were part of opts.names\n\t\tfor i := 0; i < len(ass); i++ {\n\t\t\texpectedSourceName := opts.names[i]\n\t\t\tassert.True(t, endedSources[expectedSourceName], \"Expected source %s to have sent SourceEOF\", expectedSourceName)\n\t\t}\n\n\t\t// The final 'got' map should contain all items because streamMergeWithSourceEOF merges them at the end.\n\t\tassert.Equal(t, map[int]string{\n\t\t\t1: \"1\",\n\t\t\t2: \"2\",\n\t\t\t3: \"3\",\n\t\t\t4: \"4\",\n\t\t}, got)\n\t})\n\n\ttype TestType struct {\n\t\tA int\n\t\tB []string\n\t}\n\n\tRegisterValuesMergeFunc(func(vs []*TestType) (*TestType, error) {\n\t\tret := &TestType{}\n\t\tfor _, v := range vs {\n\t\t\tif v == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif ret.A < 0 {\n\t\t\t\treturn nil, fmt.Errorf(\"test error: %v\", ret.A)\n\t\t\t}\n\t\t\tret.A += v.A\n\t\t\tret.B = append(ret.B, v.B...)\n\t\t}\n\t\tsort.Strings(ret.B)\n\t\treturn ret, nil\n\t})\n\n\tt.Run(\"custom merge\", func(t *testing.T) {\n\t\tt.Run(\"regular\", func(t *testing.T) {\n\t\t\tvs := []any{\n\t\t\t\t&TestType{A: 0, B: []string{}},\n\t\t\t\t&TestType{A: 1, B: []string{\"1\"}},\n\t\t\t\t&TestType{A: 2, B: []string{\"2\", \"22\"}},\n\t\t\t\t&TestType{A: 3, B: []string{\"3\", \"33\", \"333\"}},\n\t\t\t}\n\t\t\tret, err := mergeValues(vs, nil)\n\t\t\trequire.NoError(t, err)\n\t\t\tassert.Equal(t, &TestType{\n\t\t\t\tA: 6,\n\t\t\t\tB: []string{\"1\", \"2\", \"22\", \"3\", \"33\", \"333\"},\n\t\t\t}, ret)\n\t\t})\n\n\t\tt.Run(\"custom error\", func(t *testing.T) {\n\t\t\tvs := []any{\n\t\t\t\t&TestType{A: 0, B: []string{}},\n\t\t\t\t&TestType{A: 1, B: []string{\"1\"}},\n\t\t\t\t&TestType{A: -2, B: []string{\"2\", \"22\"}},\n\t\t\t\t&TestType{A: 3, B: []string{\"3\", \"33\", \"333\"}},\n\t\t\t}\n\t\t\t_, err := mergeValues(vs, nil)\n\t\t\trequire.ErrorContains(t, err, \"test error\")\n\t\t})\n\n\t\tt.Run(\"type mismatch\", func(t *testing.T) {\n\t\t\tvs := []any{\n\t\t\t\t&TestType{A: 0, B: []string{}},\n\t\t\t\t&TestType{A: 1, B: []string{\"1\"}},\n\t\t\t\t&TestType{A: 2, B: []string{\"2\", \"22\"}},\n\t\t\t\t\"test3\",\n\t\t\t}\n\t\t\t_, err := mergeValues(vs, nil)\n\t\t\trequire.ErrorContains(t, err, \"type mismatch\")\n\t\t})\n\n\t\tt.Run(\"stream\", func(t *testing.T) {\n\t\t\tass := []any{\n\t\t\t\tpackStreamReader(schema.StreamReaderFromArray([]*TestType{\n\t\t\t\t\t{A: 0, B: []string{}},\n\t\t\t\t})),\n\t\t\t\tpackStreamReader(schema.StreamReaderFromArray([]*TestType{\n\t\t\t\t\t{A: 1, B: []string{\"1\"}},\n\t\t\t\t})),\n\t\t\t\tpackStreamReader(schema.StreamReaderFromArray([]*TestType{\n\t\t\t\t\t{A: 2, B: []string{\"2\", \"22\"}},\n\t\t\t\t})),\n\t\t\t\tpackStreamReader(schema.StreamReaderFromArray([]*TestType{\n\t\t\t\t\t{A: 3, B: []string{\"3\", \"33\", \"333\"}},\n\t\t\t\t})),\n\t\t\t}\n\t\t\tisr, err := mergeValues(ass, nil)\n\t\t\trequire.NoError(t, err)\n\t\t\tret, ok := unpackStreamReader[*TestType](isr.(streamReader))\n\t\t\trequire.True(t, ok)\n\t\t\tdefer ret.Close()\n\n\t\t\tvar vs []any\n\t\t\tfor i := 0; i < 4; i++ {\n\t\t\t\tv, err := ret.Recv()\n\t\t\t\trequire.NoError(t, err)\n\t\t\t\tvs = append(vs, v)\n\t\t\t}\n\n\t\t\t_, err = ret.Recv()\n\t\t\trequire.ErrorIs(t, err, io.EOF)\n\n\t\t\tmerged, err := mergeValues(vs, nil)\n\t\t\trequire.NoError(t, err)\n\n\t\t\tassert.Equal(t, &TestType{\n\t\t\t\tA: 6,\n\t\t\t\tB: []string{\"1\", \"2\", \"22\", \"3\", \"33\", \"333\"},\n\t\t\t}, merged)\n\t\t})\n\t})\n\n\tt.Run(\"unregistered type\", func(t *testing.T) {\n\t\ttype Unregistered TestType\n\t\t_, err := mergeValues([]any{&Unregistered{}}, nil)\n\t\tassert.ErrorContains(t, err, \"unsupported type\")\n\t})\n\n}\n"
  },
  {
    "path": "compose/workflow.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// WorkflowNode is the node of the Workflow.\ntype WorkflowNode struct {\n\tg                *graph\n\tkey              string\n\taddInputs        []func() error\n\tstaticValues     map[string]any\n\tdependencySetter func(fromNodeKey string, typ dependencyType)\n\tmappedFieldPath  map[string]any\n}\n\n// Workflow is wrapper of graph, replacing AddEdge with declaring dependencies and field mappings between nodes.\n// Under the hood it uses NodeTriggerMode(AllPredecessor), so does not support cycles.\ntype Workflow[I, O any] struct {\n\tg                *graph\n\tworkflowNodes    map[string]*WorkflowNode\n\tworkflowBranches []*WorkflowBranch\n\tdependencies     map[string]map[string]dependencyType\n}\n\ntype dependencyType int\n\nconst (\n\tnormalDependency dependencyType = iota\n\tnoDirectDependency\n\tbranchDependency\n)\n\n// NewWorkflow creates a new Workflow.\nfunc NewWorkflow[I, O any](opts ...NewGraphOption) *Workflow[I, O] {\n\toptions := &newGraphOptions{}\n\tfor _, opt := range opts {\n\t\topt(options)\n\t}\n\n\twf := &Workflow[I, O]{\n\t\tg: newGraphFromGeneric[I, O](\n\t\t\tComponentOfWorkflow,\n\t\t\toptions.withState,\n\t\t\toptions.stateType,\n\t\t\topts,\n\t\t),\n\t\tworkflowNodes: make(map[string]*WorkflowNode),\n\t\tdependencies:  make(map[string]map[string]dependencyType),\n\t}\n\n\treturn wf\n}\n\n// Compile builds the workflow into a runnable graph.\nfunc (wf *Workflow[I, O]) Compile(ctx context.Context, opts ...GraphCompileOption) (Runnable[I, O], error) {\n\treturn compileAnyGraph[I, O](ctx, wf, opts...)\n}\n\n// AddChatModelNode adds a chat model node and returns it.\nfunc (wf *Workflow[I, O]) AddChatModelNode(key string, chatModel model.BaseChatModel, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddChatModelNode(key, chatModel, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddChatTemplateNode adds a chat template node and returns it.\nfunc (wf *Workflow[I, O]) AddChatTemplateNode(key string, chatTemplate prompt.ChatTemplate, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddChatTemplateNode(key, chatTemplate, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddToolsNode adds a tools node and returns it.\nfunc (wf *Workflow[I, O]) AddToolsNode(key string, tools *ToolsNode, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddToolsNode(key, tools, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddRetrieverNode adds a retriever node and returns it.\nfunc (wf *Workflow[I, O]) AddRetrieverNode(key string, retriever retriever.Retriever, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddRetrieverNode(key, retriever, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddEmbeddingNode adds an embedding node and returns it.\nfunc (wf *Workflow[I, O]) AddEmbeddingNode(key string, embedding embedding.Embedder, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddEmbeddingNode(key, embedding, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddIndexerNode adds an indexer node to the workflow and returns it.\nfunc (wf *Workflow[I, O]) AddIndexerNode(key string, indexer indexer.Indexer, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddIndexerNode(key, indexer, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddLoaderNode adds a document loader node to the workflow and returns it.\nfunc (wf *Workflow[I, O]) AddLoaderNode(key string, loader document.Loader, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddLoaderNode(key, loader, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddDocumentTransformerNode adds a document transformer node and returns it.\nfunc (wf *Workflow[I, O]) AddDocumentTransformerNode(key string, transformer document.Transformer, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddDocumentTransformerNode(key, transformer, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddGraphNode adds a nested graph node to the workflow and returns it.\nfunc (wf *Workflow[I, O]) AddGraphNode(key string, graph AnyGraph, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddGraphNode(key, graph, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddLambdaNode adds a lambda node to the workflow and returns it.\nfunc (wf *Workflow[I, O]) AddLambdaNode(key string, lambda *Lambda, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddLambdaNode(key, lambda, opts...)\n\treturn wf.initNode(key)\n}\n\n// End returns the WorkflowNode representing END node.\nfunc (wf *Workflow[I, O]) End() *WorkflowNode {\n\tif node, ok := wf.workflowNodes[END]; ok {\n\t\treturn node\n\t}\n\treturn wf.initNode(END)\n}\n\n// AddPassthroughNode adds a passthrough node to the workflow and returns it.\nfunc (wf *Workflow[I, O]) AddPassthroughNode(key string, opts ...GraphAddNodeOpt) *WorkflowNode {\n\t_ = wf.g.AddPassthroughNode(key, opts...)\n\treturn wf.initNode(key)\n}\n\n// AddInput creates both data and execution dependencies between nodes.\n// It configures how data flows from the predecessor node (fromNodeKey) to the current node,\n// and ensures the current node only executes after the predecessor completes.\n//\n// Parameters:\n//   - fromNodeKey: the key of the predecessor node\n//   - inputs: field mappings that specify how data should flow from the predecessor\n//     to the current node. If no mappings are provided, the entire output of the\n//     predecessor will be used as input.\n//\n// Example:\n//\n//\t// Map between specific field\n//\tnode.AddInput(\"userNode\", MapFields(\"user.name\", \"displayName\"))\n//\n//\t// Use entire output\n//\tnode.AddInput(\"dataNode\")\n//\n// Returns the current node for method chaining.\nfunc (n *WorkflowNode) AddInput(fromNodeKey string, inputs ...*FieldMapping) *WorkflowNode {\n\treturn n.addDependencyRelation(fromNodeKey, inputs, &workflowAddInputOpts{})\n}\n\ntype workflowAddInputOpts struct {\n\t// noDirectDependency indicates whether to create a data mapping without establishing\n\t// a direct execution dependency. When true, the current node can access data from\n\t// the predecessor node but its execution is not directly blocked by it.\n\tnoDirectDependency bool\n\t// dependencyWithoutInput indicates whether to create an execution dependency\n\t// without any data mapping. When true, the current node will wait for the\n\t// predecessor node to complete but won't receive any data from it.\n\tdependencyWithoutInput bool\n}\n\n// WorkflowAddInputOpt configures behavior of AddInputWithOptions.\ntype WorkflowAddInputOpt func(*workflowAddInputOpts)\n\nfunc getAddInputOpts(opts []WorkflowAddInputOpt) *workflowAddInputOpts {\n\topt := &workflowAddInputOpts{}\n\tfor _, o := range opts {\n\t\to(opt)\n\t}\n\treturn opt\n}\n\n// WithNoDirectDependency creates a data mapping without establishing a direct execution dependency.\n// The predecessor node will still complete before the current node executes, but through indirect\n// execution paths rather than a direct dependency.\n//\n// In a workflow graph, node dependencies typically serve two purposes:\n// 1. Execution order: determining when nodes should run\n// 2. Data flow: specifying how data passes between nodes\n//\n// This option separates these concerns by:\n//   - Creating data mapping from the predecessor to the current node\n//   - Relying on the predecessor's path to reach the current node through other nodes\n//     that have direct execution dependencies\n//\n// Example:\n//\n//\tnode.AddInputWithOptions(\"dataNode\", mappings, WithNoDirectDependency())\n//\n// Important:\n//\n//  1. Branch scenarios: When connecting nodes on different sides of a branch,\n//     WithNoDirectDependency MUST be used to let the branch itself handle the\n//     execution order, preventing incorrect dependencies that could bypass the branch.\n//\n//  2. Execution guarantee: The predecessor will still complete before the current\n//     node executes because the predecessor must have a path (through other nodes)\n//     that eventually reaches the current node.\n//\n//  3. Graph validity: There MUST be a path from the predecessor that eventually\n//     reaches the current node through other nodes with direct dependencies.\n//     This ensures the execution order while avoiding redundant direct dependencies.\n//\n// Common use cases:\n// - Cross-branch data access where the branch handles execution order\n// - Avoiding redundant dependencies when a path already exists\nfunc WithNoDirectDependency() WorkflowAddInputOpt {\n\treturn func(opt *workflowAddInputOpts) {\n\t\topt.noDirectDependency = true\n\t}\n}\n\n// AddInputWithOptions creates a dependency between nodes with custom configuration options.\n// It allows fine-grained control over both data flow and execution dependencies.\n//\n// Parameters:\n//   - fromNodeKey: the key of the predecessor node\n//   - inputs: field mappings that specify how data flows from the predecessor to the current node.\n//     If no mappings are provided, the entire output of the predecessor will be used as input.\n//   - opts: configuration options that control how the dependency is established\n//\n// Example:\n//\n//\t// Create data mapping without direct execution dependency\n//\tnode.AddInputWithOptions(\"dataNode\", mappings, WithNoDirectDependency())\n//\n// Returns the current node for method chaining.\nfunc (n *WorkflowNode) AddInputWithOptions(fromNodeKey string, inputs []*FieldMapping, opts ...WorkflowAddInputOpt) *WorkflowNode {\n\treturn n.addDependencyRelation(fromNodeKey, inputs, getAddInputOpts(opts))\n}\n\n// AddDependency creates an execution-only dependency between nodes.\n// The current node will wait for the predecessor node to complete before executing,\n// but no data will be passed between them.\n//\n// Parameters:\n//   - fromNodeKey: the key of the predecessor node that must complete before this node starts\n//\n// Example:\n//\n//\t// Wait for \"setupNode\" to complete before executing\n//\tnode.AddDependency(\"setupNode\")\n//\n// This is useful when:\n// - You need to ensure execution order without data transfer\n// - The predecessor performs setup or initialization that must complete first\n// - You want to explicitly separate execution dependencies from data flow\n//\n// Returns the current node for method chaining.\nfunc (n *WorkflowNode) AddDependency(fromNodeKey string) *WorkflowNode {\n\treturn n.addDependencyRelation(fromNodeKey, nil, &workflowAddInputOpts{dependencyWithoutInput: true})\n}\n\n// SetStaticValue sets a static value for a field path that will be available\n// during workflow execution. These values are determined at compile time and\n// remain constant throughout the workflow's lifecycle.\n//\n// Example:\n//\n//\tnode.SetStaticValue(FieldPath{\"query\"}, \"static query\")\nfunc (n *WorkflowNode) SetStaticValue(path FieldPath, value any) *WorkflowNode {\n\tn.staticValues[path.join()] = value\n\treturn n\n}\n\nfunc (n *WorkflowNode) addDependencyRelation(fromNodeKey string, inputs []*FieldMapping, options *workflowAddInputOpts) *WorkflowNode {\n\tfor _, input := range inputs {\n\t\tinput.fromNodeKey = fromNodeKey\n\t}\n\n\tif options.noDirectDependency {\n\t\tn.addInputs = append(n.addInputs, func() error {\n\t\t\tvar paths []FieldPath\n\t\t\tfor _, input := range inputs {\n\t\t\t\tpaths = append(paths, input.targetPath())\n\t\t\t}\n\t\t\tif err := n.checkAndAddMappedPath(paths); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\n\t\t\tif err := n.g.addEdgeWithMappings(fromNodeKey, n.key, true, false, inputs...); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tn.dependencySetter(fromNodeKey, noDirectDependency)\n\t\t\treturn nil\n\t\t})\n\t} else if options.dependencyWithoutInput {\n\t\tn.addInputs = append(n.addInputs, func() error {\n\t\t\tif len(inputs) > 0 {\n\t\t\t\treturn fmt.Errorf(\"dependency without input should not have inputs. node: %s, fromNode: %s, inputs: %v\", n.key, fromNodeKey, inputs)\n\t\t\t}\n\t\t\tif err := n.g.addEdgeWithMappings(fromNodeKey, n.key, false, true); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tn.dependencySetter(fromNodeKey, normalDependency)\n\t\t\treturn nil\n\t\t})\n\t} else {\n\t\tn.addInputs = append(n.addInputs, func() error {\n\t\t\tvar paths []FieldPath\n\t\t\tfor _, input := range inputs {\n\t\t\t\tpaths = append(paths, input.targetPath())\n\t\t\t}\n\t\t\tif err := n.checkAndAddMappedPath(paths); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\n\t\t\tif err := n.g.addEdgeWithMappings(fromNodeKey, n.key, false, false, inputs...); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tn.dependencySetter(fromNodeKey, normalDependency)\n\t\t\treturn nil\n\t\t})\n\t}\n\n\treturn n\n}\n\nfunc (n *WorkflowNode) checkAndAddMappedPath(paths []FieldPath) error {\n\tif v, ok := n.mappedFieldPath[\"\"]; ok {\n\t\tif _, ok = v.(struct{}); ok {\n\t\t\treturn fmt.Errorf(\"entire output has already been mapped for node: %s\", n.key)\n\t\t}\n\t} else {\n\t\tif len(paths) == 0 {\n\t\t\tn.mappedFieldPath[\"\"] = struct{}{}\n\t\t\treturn nil\n\t\t} else {\n\t\t\tn.mappedFieldPath[\"\"] = map[string]any{}\n\t\t}\n\t}\n\n\tfor _, targetPath := range paths {\n\t\tm := n.mappedFieldPath[\"\"].(map[string]any)\n\t\tvar traversed FieldPath\n\t\tfor i, path := range targetPath {\n\t\t\ttraversed = append(traversed, path)\n\t\t\tif v, ok := m[path]; ok {\n\t\t\t\tif _, ok = v.(struct{}); ok {\n\t\t\t\t\treturn fmt.Errorf(\"two terminal field paths conflict for node %s: %v, %v\", n.key, traversed, targetPath)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tif i < len(targetPath)-1 {\n\t\t\t\tm[path] = make(map[string]any)\n\t\t\t\tm = m[path].(map[string]any)\n\t\t\t} else {\n\t\t\t\tm[path] = struct{}{}\n\t\t\t}\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// WorkflowBranch represents a branch added to a workflow.\n// Each branch may define its own end nodes and mappings.\ntype WorkflowBranch struct {\n\tfromNodeKey string\n\t*GraphBranch\n}\n\n// AddBranch adds a branch to the workflow.\n//\n// End Nodes Field Mappings:\n// End nodes of the branch are required to define their own field mappings.\n// This is a key distinction between Graph's Branch and Workflow's Branch:\n// - Graph's Branch: Automatically passes its input to the selected node.\n// - Workflow's Branch: Does not pass its input to the selected node.\nfunc (wf *Workflow[I, O]) AddBranch(fromNodeKey string, branch *GraphBranch) *WorkflowBranch {\n\twb := &WorkflowBranch{\n\t\tfromNodeKey: fromNodeKey,\n\t\tGraphBranch: branch,\n\t}\n\n\twf.workflowBranches = append(wf.workflowBranches, wb)\n\treturn wb\n}\n\n// AddEnd connects a node to END with optional field mappings.\n// Deprecated: use *Workflow[I,O].End() to obtain a WorkflowNode instance for END, then work with it just like a normal WorkflowNode.\nfunc (wf *Workflow[I, O]) AddEnd(fromNodeKey string, inputs ...*FieldMapping) *Workflow[I, O] {\n\tfor _, input := range inputs {\n\t\tinput.fromNodeKey = fromNodeKey\n\t}\n\t_ = wf.g.addEdgeWithMappings(fromNodeKey, END, false, false, inputs...)\n\treturn wf\n}\n\nfunc (wf *Workflow[I, O]) compile(ctx context.Context, options *graphCompileOptions) (*composableRunnable, error) {\n\tif wf.g.buildError != nil {\n\t\treturn nil, wf.g.buildError\n\t}\n\n\tfor _, wb := range wf.workflowBranches {\n\t\tfor endNode := range wb.endNodes {\n\t\t\tif endNode == END {\n\t\t\t\tif _, ok := wf.dependencies[END]; !ok {\n\t\t\t\t\twf.dependencies[END] = make(map[string]dependencyType)\n\t\t\t\t}\n\t\t\t\twf.dependencies[END][wb.fromNodeKey] = branchDependency\n\t\t\t} else {\n\t\t\t\tn := wf.workflowNodes[endNode]\n\t\t\t\tn.dependencySetter(wb.fromNodeKey, branchDependency)\n\t\t\t}\n\t\t}\n\t\t_ = wf.g.addBranch(wb.fromNodeKey, wb.GraphBranch, true)\n\t}\n\n\tfor _, n := range wf.workflowNodes {\n\t\tfor _, addInput := range n.addInputs {\n\t\t\tif err := addInput(); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t\tn.addInputs = nil\n\t}\n\n\tfor _, n := range wf.workflowNodes {\n\t\tif len(n.staticValues) > 0 {\n\t\t\tvalue := make(map[string]any, len(n.staticValues))\n\t\t\tvar paths []FieldPath\n\t\t\tfor path, v := range n.staticValues {\n\t\t\t\tvalue[path] = v\n\t\t\t\tpaths = append(paths, splitFieldPath(path))\n\t\t\t}\n\n\t\t\tif err := n.checkAndAddMappedPath(paths); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\tpair := handlerPair{\n\t\t\t\tinvoke: func(in any) (any, error) {\n\t\t\t\t\tvalues := []any{in, value}\n\t\t\t\t\treturn mergeValues(values, nil)\n\t\t\t\t},\n\t\t\t\ttransform: func(in streamReader) streamReader {\n\t\t\t\t\tsr := schema.StreamReaderFromArray([]map[string]any{value})\n\t\t\t\t\tnewS, err := mergeValues([]any{in, packStreamReader(sr)}, nil)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\terrSR, errSW := schema.Pipe[map[string]any](1)\n\t\t\t\t\t\terrSW.Send(nil, err)\n\t\t\t\t\t\terrSW.Close()\n\t\t\t\t\t\treturn packStreamReader(errSR)\n\t\t\t\t\t}\n\n\t\t\t\t\treturn newS.(streamReader)\n\t\t\t\t},\n\t\t\t}\n\n\t\t\tfor i := range paths {\n\t\t\t\twf.g.fieldMappingRecords[n.key] = append(wf.g.fieldMappingRecords[n.key], ToFieldPath(paths[i]))\n\t\t\t}\n\n\t\t\twf.g.handlerPreNode[n.key] = []handlerPair{pair}\n\t\t}\n\t}\n\n\t// TODO: check indirect edges are legal\n\n\treturn wf.g.compile(ctx, options)\n}\n\nfunc (wf *Workflow[I, O]) initNode(key string) *WorkflowNode {\n\tn := &WorkflowNode{\n\t\tg:            wf.g,\n\t\tkey:          key,\n\t\tstaticValues: make(map[string]any),\n\t\tdependencySetter: func(fromNodeKey string, typ dependencyType) {\n\t\t\tif _, ok := wf.dependencies[key]; !ok {\n\t\t\t\twf.dependencies[key] = make(map[string]dependencyType)\n\t\t\t}\n\t\t\twf.dependencies[key][fromNodeKey] = typ\n\t\t},\n\t\tmappedFieldPath: make(map[string]any),\n\t}\n\twf.workflowNodes[key] = n\n\treturn n\n}\n\nfunc (wf *Workflow[I, O]) getGenericHelper() *genericHelper {\n\treturn wf.g.getGenericHelper()\n}\n\nfunc (wf *Workflow[I, O]) inputType() reflect.Type {\n\treturn wf.g.inputType()\n}\n\nfunc (wf *Workflow[I, O]) outputType() reflect.Type {\n\treturn wf.g.outputType()\n}\n\nfunc (wf *Workflow[I, O]) component() component {\n\treturn wf.g.component()\n}\n"
  },
  {
    "path": "compose/workflow_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage compose\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/internal/mock/components/embedding\"\n\t\"github.com/cloudwego/eino/internal/mock/components/indexer\"\n\t\"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestWorkflow(t *testing.T) {\n\tctx := context.Background()\n\n\ttype structA struct {\n\t\tField1 string\n\t\tField2 int\n\t\tField3 []any\n\t}\n\n\ttype structB struct {\n\t\tField1 string\n\t\tField2 int\n\t}\n\n\ttype structC struct {\n\t\tField1 string\n\t}\n\n\ttype structE struct {\n\t\tField1 string\n\t\tField2 string\n\t\tField3 []any\n\t}\n\n\ttype structF struct {\n\t\tField1    string\n\t\tField2    string\n\t\tField3    []any\n\t\tB         int\n\t\tStateTemp string\n\t}\n\tRegisterStreamChunkConcatFunc(func(ts []*structF) (*structF, error) {\n\t\tret := &structF{}\n\t\tfor _, tt := range ts {\n\t\t\tret.Field1 += tt.Field1\n\t\t\tret.Field2 += tt.Field2\n\t\t\tret.Field3 = append(ret.Field3, tt.Field3...)\n\t\t\tret.B += tt.B\n\t\t\tret.StateTemp += tt.StateTemp\n\t\t}\n\t\treturn ret, nil\n\t})\n\n\ttype state struct {\n\t\ttemp string\n\t}\n\n\ttype structEnd struct {\n\t\tField1 string\n\t}\n\n\tsubGraph := NewGraph[string, *structB]()\n\t_ = subGraph.AddLambdaNode(\n\t\t\"1\",\n\t\tInvokableLambda(func(ctx context.Context, input string) (*structB, error) {\n\t\t\treturn &structB{Field1: input, Field2: 33}, nil\n\t\t}),\n\t)\n\t_ = subGraph.AddEdge(START, \"1\")\n\t_ = subGraph.AddEdge(\"1\", END)\n\n\tsubChain := NewChain[any, *structC]().\n\t\tAppendLambda(InvokableLambda(func(_ context.Context, in any) (*structC, error) {\n\t\t\treturn &structC{Field1: fmt.Sprintf(\"%d\", in)}, nil\n\t\t}))\n\n\ttype struct2 struct {\n\t\tF map[string]any\n\t}\n\tsubWorkflow := NewWorkflow[[]any, []any]()\n\tsubWorkflow.AddLambdaNode(\n\t\t\"1\",\n\t\tInvokableLambda(func(_ context.Context, in []any) ([]any, error) {\n\t\t\treturn in, nil\n\t\t}),\n\t\tWithOutputKey(\"key\")).\n\t\tAddInput(START) // []any -> map[\"key\"][]any\n\tsubWorkflow.AddLambdaNode(\n\t\t\"2\",\n\t\tInvokableLambda(func(_ context.Context, in []any) ([]any, error) {\n\t\t\treturn in, nil\n\t\t}),\n\t\tWithInputKey(\"key\"),\n\t\tWithOutputKey(\"key1\")).\n\t\tAddInput(\"1\") // map[\"key\"][]any -> []any -> map[\"key1\"][]any\n\tsubWorkflow.AddLambdaNode(\n\t\t\"3\",\n\t\tInvokableLambda(func(_ context.Context, in struct2) (map[string]any, error) {\n\t\t\treturn in.F, nil\n\t\t}),\n\t).\n\t\tAddInput(\"2\", ToField(\"F\")) // map[\"key1\"][]any -> map[\"F\"]map[\"key1\"][]any -> struct2{F: map[\"key1\"]any} -> map[\"key1\"][]any\n\tsubWorkflow.AddLambdaNode(\n\t\t\"4\",\n\t\tInvokableLambda(func(_ context.Context, in []any) ([]any, error) {\n\t\t\treturn in, nil\n\t\t}),\n\t\tWithInputKey(\"key1\"),\n\t).\n\t\tAddInput(\"3\") // map[\"key1\"][]any -> []any\n\tsubWorkflow.End().AddInput(\"4\")\n\n\tw := NewWorkflow[*structA, *structEnd](WithGenLocalState(func(context.Context) *state { return &state{} }))\n\n\tw.\n\t\tAddGraphNode(\"B\", subGraph,\n\t\t\tWithStatePostHandler(func(ctx context.Context, out *structB, state *state) (*structB, error) {\n\t\t\t\tstate.temp = out.Field1\n\t\t\t\treturn out, nil\n\t\t\t})).\n\t\tAddInput(START, FromField(\"Field1\"))\n\n\tw.\n\t\tAddGraphNode(\"C\", subChain).\n\t\tAddInput(START, FromField(\"Field2\"))\n\n\tw.\n\t\tAddGraphNode(\"D\", subWorkflow).\n\t\tAddInput(START, FromField(\"Field3\"))\n\n\tw.\n\t\tAddLambdaNode(\n\t\t\t\"E\",\n\t\t\tTransformableLambda(func(_ context.Context, in *schema.StreamReader[structE]) (*schema.StreamReader[structE], error) {\n\t\t\t\treturn schema.StreamReaderWithConvert(in, func(in structE) (structE, error) {\n\t\t\t\t\tif len(in.Field1) > 0 {\n\t\t\t\t\t\tin.Field1 = \"E:\" + in.Field1\n\t\t\t\t\t}\n\t\t\t\t\tif len(in.Field2) > 0 {\n\t\t\t\t\t\tin.Field2 = \"E:\" + in.Field2\n\t\t\t\t\t}\n\n\t\t\t\t\treturn in, nil\n\t\t\t\t}), nil\n\t\t\t}),\n\t\t\tWithStreamStatePreHandler(func(ctx context.Context, in *schema.StreamReader[structE], state *state) (*schema.StreamReader[structE], error) {\n\t\t\t\ttemp := state.temp\n\t\t\t\treturn schema.StreamReaderWithConvert(in, func(v structE) (structE, error) {\n\t\t\t\t\tif len(v.Field3) > 0 {\n\t\t\t\t\t\tv.Field3 = append(v.Field3, \"Pre:\"+temp)\n\t\t\t\t\t}\n\n\t\t\t\t\treturn v, nil\n\t\t\t\t}), nil\n\t\t\t}),\n\t\t\tWithStreamStatePostHandler(func(ctx context.Context, out *schema.StreamReader[structE], state *state) (*schema.StreamReader[structE], error) {\n\t\t\t\treturn schema.StreamReaderWithConvert(out, func(v structE) (structE, error) {\n\t\t\t\t\tif len(v.Field1) > 0 {\n\t\t\t\t\t\tv.Field1 = v.Field1 + \"+Post\"\n\t\t\t\t\t}\n\t\t\t\t\treturn v, nil\n\t\t\t\t}), nil\n\t\t\t})).\n\t\tAddInput(\"B\", MapFields(\"Field1\", \"Field1\")).\n\t\tAddInput(\"C\", MapFields(\"Field1\", \"Field2\")).\n\t\tAddInput(\"D\", ToField(\"Field3\"))\n\n\tw.\n\t\tAddLambdaNode(\n\t\t\t\"F\",\n\t\t\tInvokableLambda(func(ctx context.Context, in *structF) (string, error) {\n\t\t\t\treturn fmt.Sprintf(\"%v_%v_%v_%v_%v\", in.Field1, in.Field2, in.Field3, in.B, in.StateTemp), nil\n\t\t\t}),\n\t\t\tWithStatePreHandler(func(ctx context.Context, in *structF, state *state) (*structF, error) {\n\t\t\t\tin.StateTemp = state.temp\n\t\t\t\treturn in, nil\n\t\t\t}),\n\t\t).\n\t\tAddInput(\"B\", MapFields(\"Field2\", \"B\")).\n\t\tAddInput(\"E\",\n\t\t\tMapFields(\"Field1\", \"Field1\"),\n\t\t\tMapFields(\"Field2\", \"Field2\"),\n\t\t\tMapFields(\"Field3\", \"Field3\"),\n\t\t)\n\n\tw.End().AddInput(\"F\", ToField(\"Field1\"))\n\n\tcompiled, err := w.Compile(ctx)\n\tassert.NoError(t, err)\n\n\tinput := &structA{\n\t\tField1: \"1\",\n\t\tField2: 2,\n\t\tField3: []any{\n\t\t\t1, \"good\",\n\t\t},\n\t}\n\tout, err := compiled.Invoke(ctx, input)\n\tassert.NoError(t, err)\n\tassert.Equal(t, &structEnd{\"E:1+Post_E:2_[1 good Pre:1]_33_1\"}, out)\n\n\toutStream, err := compiled.Stream(ctx, input)\n\tassert.NoError(t, err)\n\tdefer outStream.Close()\n\tfor {\n\t\tchunk, err := outStream.Recv()\n\t\tif err != nil {\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tt.Error(err)\n\t\t\treturn\n\t\t}\n\n\t\tassert.Equal(t, &structEnd{\"E:1+Post_E:2_[1 good Pre:1]_33_1\"}, chunk)\n\t}\n}\n\nfunc TestWorkflowWithMap(t *testing.T) {\n\tctx := context.Background()\n\n\ttype structA struct {\n\t\tF1 any\n\t}\n\n\twf := NewWorkflow[map[string]any, map[string]any]()\n\twf.AddLambdaNode(\"lambda1\", InvokableLambda(func(ctx context.Context, in map[string]any) (map[string]any, error) {\n\t\treturn in, nil\n\t})).AddInput(START, MapFields(\"map_key\", \"lambda1_key\"))\n\twf.AddLambdaNode(\"lambda2\", InvokableLambda(func(ctx context.Context, in *structA) (*structA, error) {\n\t\treturn in, nil\n\t})).AddInput(START, MapFields(\"map_key\", \"F1\"))\n\twf.End().AddInput(\"lambda1\", MapFields(\"lambda1_key\", \"end_lambda1\"))\n\twf.End().AddInput(\"lambda2\", MapFields(\"F1\", \"end_lambda2\"))\n\tr, err := wf.Compile(ctx)\n\tassert.NoError(t, err)\n\tout, err := r.Invoke(ctx, map[string]any{\"map_key\": \"value\"})\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\"end_lambda1\": \"value\", \"end_lambda2\": \"value\"}, out)\n}\n\nfunc TestWorkflowWithNestedFieldMappings(t *testing.T) {\n\tctx := context.Background()\n\n\ttype structA struct {\n\t\tF1 string\n\t}\n\n\ttype structB struct {\n\t\tF1 *structA\n\t\tF2 map[string]any\n\t\tF3 int\n\t\tF4 any\n\t\tF5 map[string]structA\n\t\tF6 structA\n\t}\n\n\tt.Run(\"from struct.struct.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[*structB, string]()\n\t\twf.End().AddInput(START, FromFieldPath([]string{\"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, &structB{\n\t\t\tF1: &structA{\n\t\t\t\tF1: \"hello\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\n\t\twf = NewWorkflow[*structB, string]()\n\t\twf.End().AddInput(START, FromFieldPath([]string{\"F1\", \"F2\"}))\n\t\t_, err = wf.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"has no field[F2]\")\n\t})\n\n\tt.Run(\"to struct.(non-ptr)struct.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, *structB]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F6\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, &structB{\n\t\t\tF6: structA{\n\t\t\t\tF1: \"hello\",\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"to map.(non-ptr)struct.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]structA]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"key\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]structA{\n\t\t\t\"key\": {\n\t\t\t\tF1: \"hello\",\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"from map.map.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]map[string]string, string]()\n\t\twf.End().AddInput(START, FromFieldPath([]string{\"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, map[string]map[string]string{\n\t\t\t\"F1\": {\n\t\t\t\t\"F1\": \"hello\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\n\t\t_, err = r.Invoke(ctx, map[string]map[string]string{\n\t\t\t\"F1\": {\n\t\t\t\t\"F2\": \"hello\",\n\t\t\t},\n\t\t})\n\n\t\tvar ie *internalError\n\t\tassert.True(t, errors.As(err, &ie))\n\t\tvar myErr *errMapKeyNotFound\n\t\tassert.True(t, errors.As(ie.origError, &myErr))\n\t})\n\n\tt.Run(\"from struct.map.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[*structB, string]()\n\t\twf.End().AddInput(START, FromFieldPath([]string{\"F2\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, &structB{\n\t\t\tF2: map[string]any{\n\t\t\t\t\"F1\": \"hello\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\n\t\t_, err = r.Invoke(ctx, &structB{\n\t\t\tF2: map[string]any{\n\t\t\t\t\"F2\": \"hello\",\n\t\t\t},\n\t\t})\n\t\tvar ie *internalError\n\t\tassert.True(t, errors.As(err, &ie))\n\t\tvar myErr *errMapKeyNotFound\n\t\tassert.True(t, errors.As(ie.origError, &myErr))\n\t})\n\n\tt.Run(\"from map.struct.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]*structA, string]()\n\t\twf.End().AddInput(START, FromFieldPath([]string{\"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, map[string]*structA{\n\t\t\t\"F1\": {\n\t\t\t\tF1: \"hello\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\n\t\twf = NewWorkflow[map[string]*structA, string]()\n\t\twf.End().AddInput(START, FromFieldPath([]string{\"F1\", \"F2\"}))\n\t\t_, err = wf.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"has no field[F2]\")\n\t})\n\n\tt.Run(\"from map[string]any.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, string]()\n\t\twf.End().AddInput(START, FromFieldPath([]string{\"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, map[string]any{\n\t\t\t\"F1\": &structA{\n\t\t\t\tF1: \"hello\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\n\t\tout, err = r.Invoke(ctx, map[string]any{\n\t\t\t\"F1\": map[string]any{\n\t\t\t\t\"F1\": \"hello\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\n\t\t_, err = r.Invoke(ctx, map[string]any{\n\t\t\t\"F1\": 1,\n\t\t})\n\n\t\tvar ie *internalError\n\t\tassert.True(t, errors.As(err, &ie))\n\t\tvar myErr *errInterfaceNotValidForFieldMapping\n\t\tassert.True(t, errors.As(ie.origError, &myErr))\n\t})\n\n\tt.Run(\"to struct.struct.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, *structB]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, &structB{\n\t\t\tF1: &structA{\n\t\t\t\tF1: \"hello\",\n\t\t\t},\n\t\t}, out)\n\n\t\twf = NewWorkflow[string, *structB]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F1\", \"F2\"}))\n\t\t_, err = wf.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"has no field[F2]\")\n\t})\n\n\tt.Run(\"to map.map.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]map[string]string]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]map[string]string{\n\t\t\t\"F1\": {\n\t\t\t\t\"F1\": \"hello\",\n\t\t\t},\n\t\t}, out)\n\n\t\twf1 := NewWorkflow[string, map[string]map[string]int]()\n\t\twf1.End().AddInput(START, ToFieldPath([]string{\"F1\", \"F1\"}))\n\t\t_, err = wf1.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"field[string]-[int] is absolutely not assignable\")\n\t})\n\n\tt.Run(\"to struct.map.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, *structB]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F2\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, &structB{\n\t\t\tF2: map[string]any{\n\t\t\t\t\"F1\": \"hello\",\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"to map.struct.struct.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]*structB]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F1\", \"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]*structB{\n\t\t\t\"F1\": {\n\t\t\t\tF1: &structA{\n\t\t\t\t\tF1: \"hello\",\n\t\t\t\t},\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"to struct.map.struct.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, *structB]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F5\", \"key\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, &structB{\n\t\t\tF5: map[string]structA{\n\t\t\t\t\"key\": {\n\t\t\t\t\tF1: \"hello\",\n\t\t\t\t},\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"to map.map.struct(non-ptr).field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]map[string]structA]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"key1\", \"key2\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]map[string]structA{\n\t\t\t\"key1\": {\n\t\t\t\t\"key2\": {\n\t\t\t\t\tF1: \"hello\",\n\t\t\t\t},\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"to struct.int.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, *structB]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F3\", \"F1\", \"F1\"}))\n\t\t_, err := wf.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"type[int] is not valid\")\n\t})\n\n\tt.Run(\"to struct.any.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, *structB]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"F4\", \"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, &structB{\n\t\t\tF4: map[string]any{\n\t\t\t\t\"F1\": map[string]any{\n\t\t\t\t\t\"F1\": \"hello\",\n\t\t\t\t},\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"to map.any.any.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]any]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"Key1\", \"Key2\", \"Key3\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\t\"Key1\": map[string]any{\n\t\t\t\t\"Key2\": map[string]any{\n\t\t\t\t\t\"Key3\": \"hello\",\n\t\t\t\t},\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"to any\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, any]()\n\t\twf.End().AddInput(START)\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\t})\n\n\tt.Run(\"to any.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, any]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"Key1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\t\"Key1\": \"hello\",\n\t\t}, out)\n\t})\n\n\tt.Run(\"to interface.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]fmt.Stringer]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"Key1\", \"A\"}))\n\t\t_, err := wf.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"static check failed for mapping [from start to Key1\\u001FA(field)], \"+\n\t\t\t\"the successor has intermediate interface type fmt.Stringer\")\n\t})\n\n\tt.Run(\"both to map.any, and to map.any.field\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]any]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"Key1\"}), ToFieldPath([]string{\"Key1\", \"Key2\"}))\n\t\t_, err := wf.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"two terminal field paths conflict\")\n\t})\n\n\tt.Run(\"to map.any.any.field1, and to map.any.any.field2\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]any]()\n\t\twf.End().AddInput(START, ToFieldPath([]string{\"Key1\", \"Key2\", \"key3\"}), ToFieldPath([]string{\"Key1\", \"Key2\", \"key4\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\t\"Key1\": map[string]any{\n\t\t\t\t\"Key2\": map[string]any{\n\t\t\t\t\t\"key3\": \"hello\",\n\t\t\t\t\t\"key4\": \"hello\",\n\t\t\t\t},\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"from nested to nested\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, *structB]()\n\t\twf.End().AddInput(START, MapFieldPaths([]string{\"key1\", \"key2\"}, []string{\"F1\", \"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, map[string]any{\n\t\t\t\"key1\": map[string]any{\n\t\t\t\t\"key2\": \"hello\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, &structB{\n\t\t\tF1: &structA{\n\t\t\t\tF1: \"hello\",\n\t\t\t},\n\t\t}, out)\n\t})\n\n\tt.Run(\"from nested to normal\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, *structA]()\n\t\twf.End().AddInput(START, MapFieldPaths(FieldPath{\"key1\", \"key2\"}, FieldPath{\"F1\"}))\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, map[string]any{\n\t\t\t\"key1\": map[string]any{\n\t\t\t\t\"key2\": \"hello\",\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, &structA{\n\t\t\tF1: \"hello\",\n\t\t}, out)\n\t})\n}\n\nfunc TestWorkflowCompile(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\n\tt.Run(\"compile without add end\", func(t *testing.T) {\n\t\tw := NewWorkflow[*schema.Message, []*schema.Message]()\n\t\tw.AddToolsNode(\"1\", &ToolsNode{}).AddInput(START)\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"end node not set\")\n\t})\n\n\tt.Run(\"type mismatch\", func(t *testing.T) {\n\t\tw := NewWorkflow[string, string]()\n\t\tw.AddToolsNode(\"1\", &ToolsNode{}).AddInput(START)\n\t\tw.End().AddInput(\"1\")\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \" mismatch\")\n\t})\n\n\tt.Run(\"predecessor's output not struct/struct ptr/map, mapping has FromField\", func(t *testing.T) {\n\t\tw := NewWorkflow[[]*schema.Document, []string]()\n\n\t\tw.AddIndexerNode(\"indexer\", indexer.NewMockIndexer(ctrl)).AddInput(START, FromField(\"F1\"))\n\t\tw.End().AddInput(\"indexer\")\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"predecessor output type should be struct\")\n\t})\n\n\tt.Run(\"successor's input not struct/struct ptr/map, mapping has ToField\", func(t *testing.T) {\n\t\tw := NewWorkflow[[]string, [][]float64]()\n\t\tw.AddEmbeddingNode(\"embedder\", embedding.NewMockEmbedder(ctrl)).AddInput(START, ToField(\"F1\"))\n\t\tw.End().AddInput(\"embedder\")\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"successor input type should be struct\")\n\t})\n\n\tt.Run(\"map to non existing field in predecessor\", func(t *testing.T) {\n\t\tw := NewWorkflow[*schema.Message, []*schema.Message]()\n\t\tw.AddToolsNode(\"tools_node\", &ToolsNode{}).AddInput(START, FromField(\"non_exist\"))\n\t\tw.End().AddInput(\"tools_node\")\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"type[schema.Message] has no field[non_exist]\")\n\t})\n\n\tt.Run(\"map to not exported field in successor\", func(t *testing.T) {\n\t\tw := NewWorkflow[string, *FieldMapping]()\n\t\tw.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\treturn input, nil\n\t\t})).AddInput(START)\n\t\tw.End().AddInput(\"1\", ToField(\"to\"))\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"has an unexported field[to]\")\n\t})\n\n\tt.Run(\"map from not exported field in predecessor\", func(t *testing.T) {\n\t\tw := NewWorkflow[*FieldMapping, string]()\n\t\tw.End().AddInput(START, FromField(\"from\"))\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"has an unexported field[from]\")\n\t})\n\n\tt.Run(\"duplicate node key\", func(t *testing.T) {\n\t\tw := NewWorkflow[[]*schema.Message, []*schema.Message]()\n\t\tw.AddChatModelNode(\"1\", model.NewMockChatModel(ctrl)).AddInput(START)\n\t\tw.AddToolsNode(\"1\", &ToolsNode{}).AddInput(\"1\")\n\t\tw.End().AddInput(\"1\")\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"node '1' already present\")\n\t})\n\n\tt.Run(\"from non-existing node\", func(t *testing.T) {\n\t\tw := NewWorkflow[*schema.Message, []*schema.Message]()\n\t\tw.AddToolsNode(\"1\", &ToolsNode{}).AddInput(START)\n\t\tw.End().AddInput(\"2\")\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"edge start node '2' needs to be added to graph first\")\n\t})\n\n\tt.Run(\"to map with non-string key type\", func(t *testing.T) {\n\t\tw := NewWorkflow[string, map[int]any]()\n\t\tw.End().AddInput(START, ToField(\"1\"))\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"type[map[int]interface {}] is not a map with string or string alias key\")\n\n\t\ttype stringAlias string\n\t\tw1 := NewWorkflow[string, map[stringAlias]any]()\n\t\tw1.End().AddInput(START, ToField(\"1\"))\n\t\tr, err := w1.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[stringAlias]any{\n\t\t\t\"1\": \"hello\",\n\t\t}, out)\n\t})\n\n\tt.Run(\"from map with non-string key type\", func(t *testing.T) {\n\t\tw := NewWorkflow[map[int]any, string]()\n\t\tw.End().AddInput(START, FromField(\"1\"))\n\t\t_, err := w.Compile(ctx)\n\t\tassert.ErrorContains(t, err, \"type[map[int]interface {}] is not a map with string or string alias key\")\n\t\ttype stringAlias string\n\t\tw1 := NewWorkflow[map[stringAlias]any, string]()\n\t\tw1.End().AddInput(START, FromField(\"1\"))\n\t\tr, err := w1.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, map[stringAlias]any{\n\t\t\t\"1\": \"hello\",\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\t})\n}\n\nfunc TestFanInToSameDest(t *testing.T) {\n\tt.Run(\"traditional outputKey fan-in with map[string]any\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, []*schema.Message]()\n\t\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in, nil\n\t\t}), WithOutputKey(\"q1\")).AddInput(START)\n\t\twf.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in, nil\n\t\t}), WithOutputKey(\"q2\")).AddInput(START)\n\t\twf.AddChatTemplateNode(\"prompt\", prompt.FromMessages(schema.Jinja2, schema.UserMessage(\"{{q1}}_{{q2}}\"))).\n\t\t\tAddInput(\"1\", MapFields(\"q1\", \"q1\")).\n\t\t\tAddInput(\"2\", MapFields(\"q2\", \"q2\"))\n\t\twf.End().AddInput(\"prompt\")\n\t\tc, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tout, err := c.Invoke(context.Background(), \"query\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []*schema.Message{{Role: schema.User, Content: \"query_query\"}}, out)\n\t})\n\n\tt.Run(\"fan-in to a field of map\", func(t *testing.T) {\n\t\ttype dest struct {\n\t\t\tF map[string]any\n\t\t}\n\n\t\ttype in struct {\n\t\t\tA string\n\t\t\tB int\n\t\t}\n\n\t\twf := NewWorkflow[in, dest]()\n\t\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in, nil\n\t\t}), WithOutputKey(\"A\")).AddInput(START, FromField(\"A\"))\n\t\twf.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, in int) (output int, err error) {\n\t\t\treturn in, nil\n\t\t}), WithOutputKey(\"B\")).AddInput(START, FromField(\"B\"))\n\t\twf.End().AddInput(\"1\", ToField(\"F\")).AddInput(\"2\", ToField(\"F\"))\n\t\t_, err := wf.Compile(context.Background())\n\t\tassert.ErrorContains(t, err, \"two terminal field paths conflict for node end: [F], [F]\")\n\t})\n}\n\nfunc TestIndirectEdge(t *testing.T) {\n\twf := NewWorkflow[string, map[string]any]()\n\n\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\treturn in + \"_\" + in, nil\n\t})).AddInput(START)\n\n\twf.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, in map[string]string) (output string, err error) {\n\t\treturn in[\"1\"] + \"_\" + in[START], nil\n\t})).AddInput(\"1\", ToField(\"1\")).\n\t\tAddInputWithOptions(START, []*FieldMapping{ToField(START)}, WithNoDirectDependency())\n\n\twf.End().AddInput(\"2\", ToField(\"2\")).\n\t\tAddInputWithOptions(\"1\", []*FieldMapping{ToField(\"1\")}, WithNoDirectDependency())\n\n\tr, err := wf.Compile(context.Background())\n\tassert.NoError(t, err)\n\tout, err := r.Invoke(context.Background(), \"query\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\"1\": \"query_query\", \"2\": \"query_query_query\"}, out)\n}\n\nfunc TestDependencyWithNoInput(t *testing.T) {\n\tt.Run(\"simple case\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, string]()\n\t\twf.AddLambdaNode(\"0\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn \"useless\", nil\n\t\t})).AddInput(START)\n\t\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in + \"_done\", nil\n\t\t})).AddDependency(\"0\").AddInputWithOptions(START, nil, WithNoDirectDependency())\n\t\twf.End().AddInput(\"1\")\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(context.Background(), \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello_done\", out)\n\t})\n\n\tt.Run(\"simple control flow: [Start] --> [Node '0'] --> [End]\", func(t *testing.T) {\n\t\t// [Start] --> [Node \"0\"] --> [End]\n\t\twf := NewWorkflow[map[string]any, map[string]any]()\n\t\twf.AddLambdaNode(\"0\", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) {\n\t\t\treturn map[string]any{\n\t\t\t\t\"result\": \"result from node 0\",\n\t\t\t}, nil\n\t\t})).AddDependency(START)\n\t\twf.End().AddInput(\"0\", ToField(\"final_result\")).\n\t\t\tAddInputWithOptions(START, []*FieldMapping{ToField(\"final_from_start\")}, WithNoDirectDependency())\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tret, err := r.Invoke(context.Background(), map[string]any{\n\t\t\t\"input\": \"hello\",\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\t\"final_result\": map[string]any{\n\t\t\t\t\"result\": \"result from node 0\",\n\t\t\t},\n\t\t\t\"final_from_start\": map[string]any{\n\t\t\t\t\"input\": \"hello\",\n\t\t\t},\n\t\t}, ret)\n\n\t\tsRet, err := r.Stream(context.Background(), map[string]any{\n\t\t\t\"input\": \"hello\",\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tret, err = concatStreamReader(sRet)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\t\"final_result\": map[string]any{\n\t\t\t\t\"result\": \"result from node 0\",\n\t\t\t},\n\t\t\t\"final_from_start\": map[string]any{\n\t\t\t\t\"input\": \"hello\",\n\t\t\t},\n\t\t}, ret)\n\t})\n}\n\nfunc TestStaticValue(t *testing.T) {\n\tt.Run(\"prefill map\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]any]()\n\t\twf.AddLambdaNode(\"0\", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) {\n\t\t\treturn in, nil\n\t\t})).\n\t\t\tAddInput(START, ToField(START)).\n\t\t\tSetStaticValue(FieldPath{\"prefilled\"}, \"yo-ho\")\n\t\twf.End().AddInput(\"0\")\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(context.Background(), \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\"prefilled\": \"yo-ho\", START: \"hello\"}, out)\n\t\tstreamOut, err := r.Stream(context.Background(), \"hello\")\n\t\tassert.NoError(t, err)\n\t\tout = map[string]any{}\n\t\tfor {\n\t\t\tchunk, err := streamOut.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\tfor k, v := range chunk {\n\t\t\t\tout[k] = v\n\t\t\t}\n\t\t}\n\t\tassert.Equal(t, map[string]any{\"prefilled\": \"yo-ho\", START: \"hello\"}, out)\n\t})\n\n\tt.Run(\"static value and to-all mapping conflict\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, map[string]any]()\n\t\twf.AddLambdaNode(\"0\", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) {\n\t\t\treturn in, nil\n\t\t})).\n\t\t\tAddInput(START).\n\t\t\tSetStaticValue(\n\t\t\t\tFieldPath{\"prefilled\"},\n\t\t\t\t\"yo-ho\",\n\t\t\t)\n\t\twf.End().AddInput(\"0\")\n\t\t_, err := wf.Compile(context.Background())\n\t\tassert.ErrorContains(t, err, \"entire output has already been mapped for node: 0\")\n\t})\n\n\tt.Run(\"static value and dynamic mapping conflict\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]any]()\n\t\twf.AddLambdaNode(\"0\", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) {\n\t\t\treturn in, nil\n\t\t})).\n\t\t\tAddInput(START, ToField(\"prefilled\")).\n\t\t\tSetStaticValue(FieldPath{\"prefilled\"}, \"yo-ho\")\n\t\twf.End().AddInput(\"0\")\n\t\t_, err := wf.Compile(context.Background())\n\t\tassert.ErrorContains(t, err, \"two terminal field paths conflict for node 0: [prefilled], [prefilled]\")\n\t})\n\n\tt.Run(\"all inputs are static values\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]any]()\n\t\twf.AddLambdaNode(\"0\", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) {\n\t\t\treturn in, nil\n\t\t})).\n\t\t\tAddDependency(START).\n\t\t\tSetStaticValue(FieldPath{\"a\", \"b\"}, \"a_b\").\n\t\t\tSetStaticValue(FieldPath{\"c\", \"d\"}, \"c_d\").\n\t\t\tSetStaticValue(FieldPath{\"a\", \"d\"}, \"a_d\")\n\t\twf.End().AddInput(\"0\")\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(context.Background(), \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\t\"a\": map[string]any{\n\t\t\t\t\"b\": \"a_b\",\n\t\t\t\t\"d\": \"a_d\",\n\t\t\t},\n\t\t\t\"c\": map[string]any{\n\t\t\t\t\"d\": \"c_d\",\n\t\t\t},\n\t\t}, out)\n\n\t\ttype a struct {\n\t\t\tB string\n\t\t\tD string\n\t\t}\n\n\t\ttype s struct {\n\t\t\tA a\n\t\t\tC map[string]any\n\t\t}\n\n\t\twf1 := NewWorkflow[string, *s]()\n\t\twf1.AddLambdaNode(\"0\", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) {\n\t\t\treturn in, nil\n\t\t})).\n\t\t\tAddDependency(START).\n\t\t\tSetStaticValue(FieldPath{\"A\", \"B\"}, \"a_b\").\n\t\t\tSetStaticValue(FieldPath{\"C\", \"D\"}, \"c_d\").\n\t\t\tSetStaticValue(FieldPath{\"A\", \"D\"}, \"a_d\")\n\t\twf1.End().AddInput(\"0\", MapFieldPaths(FieldPath{\"A\", \"B\"}, FieldPath{\"A\", \"B\"}),\n\t\t\tMapFieldPaths(FieldPath{\"A\", \"D\"}, FieldPath{\"A\", \"D\"}),\n\t\t\tMapFields(\"C\", \"C\"))\n\t\tr1, err := wf1.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tout1, err := r1.Stream(context.Background(), \"hello\")\n\t\tassert.NoError(t, err)\n\t\toutChunk, err := out1.Recv()\n\t\tout1.Close()\n\t\tassert.Equal(t, &s{\n\t\t\tA: a{\n\t\t\t\tB: \"a_b\",\n\t\t\t\tD: \"a_d\",\n\t\t\t},\n\t\t\tC: map[string]any{\n\t\t\t\t\"D\": \"c_d\",\n\t\t\t},\n\t\t}, outChunk)\n\t})\n}\n\nfunc TestBranch(t *testing.T) {\n\tctx := context.Background()\n\tt.Run(\"simple branch: one predecessor, two successor, one of them is END\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]any]()\n\t\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in + \"_\" + in, nil\n\t\t})).AddInputWithOptions(START, nil, WithNoDirectDependency())\n\n\t\twf.AddPassthroughNode(\"branch_1\").AddInput(START, ToField(START))\n\n\t\tbranch := NewGraphBranch(func(ctx context.Context, in map[string]any) (string, error) {\n\t\t\tif in[START] == \"hello\" {\n\t\t\t\treturn \"1\", nil\n\t\t\t}\n\t\t\treturn END, nil\n\t\t}, map[string]bool{\n\t\t\t\"1\": true,\n\t\t\tEND: true,\n\t\t})\n\t\twf.AddBranch(\"branch_1\", branch)\n\t\twf.End().AddInput(\"1\", ToField(\"1\")).AddInputWithOptions(START, []*FieldMapping{ToField(START)}, WithNoDirectDependency())\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\t\"1\":   \"hello_hello\",\n\t\t\tSTART: \"hello\",\n\t\t}, out)\n\t\tout, err = r.Invoke(ctx, \"world\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\tSTART: \"world\",\n\t\t}, out)\n\t})\n\n\tt.Run(\"multiple predecessors\", func(t *testing.T) {\n\t\twf := NewWorkflow[string, map[string]any]()\n\t\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in + \"_\" + in, nil\n\t\t})).AddInput(START)\n\t\twf.AddLambdaNode(\"2\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in + \"_\" + in, nil\n\t\t})).AddInputWithOptions(\"1\", nil, WithNoDirectDependency())\n\t\twf.AddLambdaNode(\"0\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\t\treturn in + \"_\" + in, nil\n\t\t})).AddInput(START)\n\n\t\twf.AddPassthroughNode(\"branch_1\").AddInput(START, ToField(START)).AddInput(\"1\", ToField(\"1\")).AddDependency(\"0\")\n\t\twf.AddBranch(\"branch_1\", NewGraphBranch(func(ctx context.Context, in map[string]any) (string, error) {\n\t\t\tif in[START].(string) == \"hello\" {\n\t\t\t\treturn \"2\", nil\n\t\t\t}\n\t\t\treturn END, nil\n\t\t}, map[string]bool{\n\t\t\t\"2\": true,\n\t\t\tEND: true,\n\t\t}))\n\t\twf.End().AddInput(\"2\", ToField(\"2\")).AddInputWithOptions(START, []*FieldMapping{ToField(START)}, WithNoDirectDependency())\n\t\tr, err := wf.Compile(ctx)\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(ctx, \"hello\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\"2\": \"hello_hello_hello_hello\", START: \"hello\"}, out)\n\t\tout, err = r.Invoke(ctx, \"world\")\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{START: \"world\"}, out)\n\t})\n\tt.Run(\"empty input for node after branch\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, map[string]any]()\n\t\twf.AddLambdaNode(\"start_1\", InvokableLambda(func(ctx context.Context, input map[string]any) (map[string]any, error) {\n\t\t\treturn map[string]any{}, nil\n\t\t})).AddInput(\"start\")\n\t\twf.AddLambdaNode(\"branch_1\", InvokableLambda(func(ctx context.Context, input map[string]any) (map[string]any, error) {\n\t\t\treturn map[string]any{}, nil\n\t\t}))\n\t\twf.AddPassthroughNode(\"my_branch\").AddInput(\"start_1\")\n\t\twf.AddBranch(\"my_branch\", NewGraphBranch(func(ctx context.Context, input map[string]any) (string, error) {\n\t\t\treturn END, nil\n\t\t}, map[string]bool{\n\t\t\t\"branch_1\": true,\n\t\t\tEND:        true,\n\t\t}))\n\t\twf.End().AddInput(\"branch_1\")\n\t\trunner, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tresp, err := runner.Invoke(context.Background(), map[string]any{})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, resp, (map[string]any)(nil))\n\t})\n}\n\ntype goodInterface interface {\n\tGOOD()\n}\ntype goodStruct struct{}\n\nfunc (g *goodStruct) GOOD() {}\n\nfunc TestMayAssignableFieldMapping(t *testing.T) {\n\ttype in struct {\n\t\tA goodInterface\n\t}\n\twf := NewWorkflow[in, *goodStruct]()\n\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input *goodStruct) (output goodInterface, err error) { return input, nil })).\n\t\tAddInput(START, FromField(\"A\"))\n\twf.End().AddInput(\"1\")\n\tctx := context.Background()\n\tr, err := wf.Compile(ctx)\n\tassert.NoError(t, err)\n\tresult, err := r.Invoke(ctx, in{A: &goodStruct{}})\n\tassert.NoError(t, err)\n\tresult.GOOD()\n}\n\nfunc TestNilValue(t *testing.T) {\n\tt.Run(\"from map key with a nil value to map key\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, map[string]any]()\n\t\twf.End().AddInput(START, MapFields(\"a\", \"a\"))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\n\t\tresult, err := r.Invoke(context.Background(), map[string]any{\"a\": nil})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\"a\": nil}, result)\n\t})\n\n\tt.Run(\"from nil struct field to map key\", func(t *testing.T) {\n\t\ttype in struct {\n\t\t\tA *string\n\t\t}\n\t\twf := NewWorkflow[in, map[string]any]()\n\t\twf.End().AddInput(START, MapFields(\"A\", \"A\"))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tresult, err := r.Invoke(context.Background(), in{A: nil})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\"A\": (*string)(nil)}, result)\n\t})\n\n\tt.Run(\"from map key with a nil value to struct field\", func(t *testing.T) {\n\t\ttype out struct {\n\t\t\tA *string\n\t\t}\n\t\twf := NewWorkflow[map[string]any, out]()\n\t\twf.End().AddInput(START, MapFields(\"A\", \"A\"))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tresult, err := r.Invoke(context.Background(), map[string]any{\"A\": nil})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, out{A: (*string)(nil)}, result)\n\t})\n\n\tt.Run(\"from nil struct field to struct field\", func(t *testing.T) {\n\t\ttype inOut struct {\n\t\t\tA *string\n\t\t}\n\t\twf := NewWorkflow[inOut, inOut]()\n\t\twf.End().AddInput(START, MapFields(\"A\", \"A\"))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tresult, err := r.Invoke(context.Background(), inOut{A: nil})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, inOut{A: (*string)(nil)}, result)\n\t})\n\n\tt.Run(\"from nil to a type that can't be nil\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, int]()\n\t\twf.End().AddInput(START, FromField(\"a\"))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), map[string]any{\"a\": nil})\n\t\tassert.ErrorContains(t, err, \"runtime check failed for mapping [from a(field) of start], field[<nil>]-[int] is absolutely not assignable\")\n\t})\n\n\tt.Run(\"from nil to a map other than map[string]any\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, map[string]fmt.Stringer]()\n\t\twf.End().AddInput(START, MapFields(\"a\", \"a\"))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(context.Background(), map[string]any{\"a\": nil})\n\t\tassert.Equal(t, map[string]fmt.Stringer{\n\t\t\t\"a\": nil,\n\t\t}, out)\n\t})\n}\n\nfunc TestStreamFieldMap(t *testing.T) {\n\tt.Run(\"multiple incomplete chunks in source stream\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, map[string]any]()\n\t\twf.End().AddInput(START, MapFields(\"a\", \"a\"), MapFields(\"b\", \"b\"))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\n\t\tsr, sw := schema.Pipe[map[string]any](2)\n\t\tsw.Send(map[string]any{\"a\": 1}, nil)\n\t\tsw.Send(map[string]any{\"b\": 2}, nil)\n\t\tsw.Close()\n\t\toutputS, err := r.Transform(context.Background(), sr)\n\t\tassert.NoError(t, err)\n\t\tresult, err := concatStreamReader(outputS)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]any{\"a\": 1, \"b\": 2}, result)\n\t})\n}\n\nfunc TestRuntimeTypeCheck(t *testing.T) {\n\tg := NewWorkflow[map[string]any, any]()\n\n\t_ = g.\n\t\tAddLambdaNode(\"A\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\t\treturn input, nil\n\t\t})).\n\t\tAddInput(START, FromField(\"A\"))\n\n\t_ = g.AddLambdaNode(\"B\", InvokableLambda(func(ctx context.Context, input string) (output string, err error) {\n\t\treturn input, nil\n\t})).\n\t\tAddInput(START, FromField(\"B\"))\n\n\t_ = g.AddLambdaNode(\"MergeA\", InvokableLambda(func(ctx context.Context, input map[string]any) (output map[string]any, err error) {\n\t\treturn input, nil\n\t})).\n\t\tAddInput(\"A\", ToField(\"a\")).\n\t\tAddInput(\"B\", ToField(\"b\"))\n\n\tg.End().AddInput(\"MergeA\")\n\n\tctx := context.Background()\n\tr, err := g.Compile(ctx)\n\tassert.NoError(t, err)\n\tresult, err := r.Stream(ctx, map[string]any{\"A\": \"1\", \"B\": \"2\"})\n\tassert.NoError(t, err)\n\tchunk, err := result.Recv()\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\"a\": \"1\", \"b\": \"2\"}, chunk)\n\tchunk, err = result.Recv()\n\tassert.True(t, errors.Is(err, io.EOF))\n}\n\nfunc TestIntermediateMappingSource(t *testing.T) {\n\tt.Run(\"intermediate any source is nil\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, any]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"a\", \"b\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), map[string]any{\n\t\t\t\"a\": nil,\n\t\t})\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[a b] is nil for type [interface {}]\")\n\t\toutStream, err := r.Transform(context.Background(), schema.StreamReaderFromArray([]map[string]any{\n\t\t\t{\n\t\t\t\t\"a\": nil,\n\t\t\t},\n\t\t\t{\n\t\t\t\t\"b\": \"ok\",\n\t\t\t},\n\t\t}))\n\t\tassert.NoError(t, err)\n\t\t_, err = outStream.Recv()\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[a b] is nil for type [interface {}]\")\n\t\toutStream.Close()\n\t})\n\n\tt.Run(\"intermediate map source is nil\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, any]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"a\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), nil)\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[a] is nil for map type [map[string]interface {}]\")\n\t\toutStream, err := r.Stream(context.Background(), nil)\n\t\tassert.NoError(t, err)\n\t\t_, err = outStream.Recv()\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[a] is nil for map type [map[string]interface {}]\")\n\t\toutStream.Close()\n\t})\n\n\tt.Run(\"intermediate map ptr source is nil\", func(t *testing.T) {\n\t\twf := NewWorkflow[*map[string]any, any]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"a\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), nil)\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[a] is nil for type [*map[string]interface {}]\")\n\t\toutStream, err := r.Stream(context.Background(), nil)\n\t\tassert.NoError(t, err)\n\t\t_, err = outStream.Recv()\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[a] is nil for type [*map[string]interface {}]\")\n\t\toutStream.Close()\n\t})\n\n\tt.Run(\"intermediate struct ptr source is nil\", func(t *testing.T) {\n\t\ttype inner struct {\n\t\t\tA string\n\t\t}\n\n\t\twf := NewWorkflow[map[string]*inner, string]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"I\", \"A\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), map[string]*inner{\"I\": nil})\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[I A] is nil\")\n\t\toutStream, err := r.Stream(context.Background(), map[string]*inner{\"I\": nil})\n\t\tassert.NoError(t, err)\n\t\t_, err = outStream.Recv()\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[I A] is nil\")\n\t\toutStream.Close()\n\t})\n\n\tt.Run(\"intermediate interface source is nil\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]fmt.Stringer, string]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"a\", \"b\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), map[string]fmt.Stringer{\"a\": nil})\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[a b] is nil for type [fmt.Stringer]\")\n\t\toutStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{\"a\": nil})\n\t\tassert.NoError(t, err)\n\t\t_, err = outStream.Recv()\n\t\tassert.ErrorContains(t, err, \"intermediate source value on path=[a b] is nil for type [fmt.Stringer]\")\n\t\toutStream.Close()\n\t})\n\n\tt.Run(\"intermediate interface source valid\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]fmt.Stringer, string]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"a\", \"A\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tout, err := r.Invoke(context.Background(), map[string]fmt.Stringer{\"a\": &goodStruct2{A: \"hello\"}})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\t\toutStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{\"a\": &goodStruct2{A: \"hello\"}})\n\t\tassert.NoError(t, err)\n\t\tout, err = outStream.Recv()\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"hello\", out)\n\t\toutStream.Close()\n\t})\n\n\tt.Run(\"intermediate interface source, source field not found at request time\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]fmt.Stringer, string]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"a\", \"B\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), map[string]fmt.Stringer{\"a\": &goodStruct2{A: \"hello\"}})\n\t\tassert.ErrorContains(t, err, \"field mapping from a struct field, but field not found. field=B\")\n\t\toutStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{\"a\": &goodStruct2{A: \"hello\"}})\n\t\tassert.NoError(t, err)\n\t\t_, err = outStream.Recv()\n\t\tassert.ErrorContains(t, err, \"field mapping from a struct field, but field not found. field=B\")\n\t\toutStream.Close()\n\t})\n\n\tt.Run(\"intermediate interface source, source field not exported at request time\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]fmt.Stringer, string]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"a\", \"c\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), map[string]fmt.Stringer{\"a\": &goodStruct2{A: \"hello\", c: \"c\"}})\n\t\tassert.ErrorContains(t, err, \"field mapping from a struct field, but field not exported.\")\n\t\toutStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{\"a\": &goodStruct2{A: \"hello\", c: \"c\"}})\n\t\tassert.NoError(t, err)\n\t\t_, err = outStream.Recv()\n\t\tassert.ErrorContains(t, err, \"field mapping from a struct field, but field not exported.\")\n\t\toutStream.Close()\n\t})\n\n\tt.Run(\"intermediate interface source, type mismatch at request time\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]fmt.Stringer, int]()\n\t\twf.End().AddInput(START, FromFieldPath(FieldPath{\"a\", \"A\"}))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\t_, err = r.Invoke(context.Background(), map[string]fmt.Stringer{\"a\": &goodStruct2{A: \"hello\"}})\n\t\tassert.ErrorContains(t, err, \"runtime check failed for mapping [from a\\x1fA(field) of start], field[string]-[int] is absolutely not assignable\")\n\t\toutStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{\"a\": &goodStruct2{A: \"hello\"}})\n\t\tassert.NoError(t, err)\n\t\t_, err = outStream.Recv()\n\t\tassert.ErrorContains(t, err, \"runtime check failed for mapping [from a\\u001FA(field) of start], field[string]-[int] is absolutely not assignable\")\n\t\toutStream.Close()\n\t})\n}\n\ntype goodStruct2 struct {\n\tA string\n\tc string\n}\n\nfunc (g *goodStruct2) String() string {\n\treturn g.A\n}\n\nfunc TestSetFanInMergeConfig_RealStreamNode_Workflow(t *testing.T) {\n\twf := NewWorkflow[int, map[string]int]()\n\n\twf.AddLambdaNode(\"s1\", StreamableLambda(func(ctx context.Context, input int) (*schema.StreamReader[int], error) {\n\t\tsr, sw := schema.Pipe[int](2)\n\t\tsw.Send(input+1, nil)\n\t\tsw.Send(input+2, nil)\n\t\tsw.Close()\n\t\treturn sr, nil\n\t})).AddInput(START)\n\n\twf.AddLambdaNode(\"s2\", StreamableLambda(func(ctx context.Context, input int) (*schema.StreamReader[int], error) {\n\t\tsr, sw := schema.Pipe[int](2)\n\t\tsw.Send(input+10, nil)\n\t\tsw.Send(input+20, nil)\n\t\tsw.Close()\n\t\treturn sr, nil\n\t})).AddInput(START)\n\n\twf.End().AddInput(\"s1\", ToField(\"s1\")).AddInput(\"s2\", ToField(\"s2\"))\n\n\tr, err := wf.Compile(context.Background(),\n\t\tWithFanInMergeConfig(map[string]FanInMergeConfig{END: {StreamMergeWithSourceEOF: true}}))\n\tassert.NoError(t, err)\n\n\tsr, err := r.Stream(context.Background(), 1)\n\tassert.NoError(t, err)\n\n\tmerged := make(map[string]map[int]bool)\n\tvar sourceEOFCount int\n\tsourceNames := make(map[string]bool)\n\tfor {\n\t\tm, e := sr.Recv()\n\t\tif e != nil {\n\t\t\tif name, ok := schema.GetSourceName(e); ok {\n\t\t\t\tsourceEOFCount++\n\t\t\t\tsourceNames[name] = true\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif e == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, e)\n\t\t}\n\n\t\tfor k, v := range m {\n\t\t\tif merged[k] == nil {\n\t\t\t\tmerged[k] = make(map[int]bool)\n\t\t\t}\n\t\t\tmerged[k][v] = true\n\t\t}\n\t}\n\n\tassert.Equal(t, map[string]map[int]bool{\"s1\": {2: true, 3: true}, \"s2\": {11: true, 21: true}}, merged)\n\tassert.Equal(t, 2, sourceEOFCount, \"should receive SourceEOF for each input stream when StreamMergeWithSourceEOF is true\")\n\tassert.True(t, sourceNames[\"s1\"], \"should receive SourceEOF from s1\")\n\tassert.True(t, sourceNames[\"s2\"], \"should receive SourceEOF from s2\")\n}\n\nfunc TestCustomExtractor(t *testing.T) {\n\tt.Run(\"custom extract from array element\", func(t *testing.T) {\n\t\twf := NewWorkflow[[]int, map[string]int]()\n\t\twf.End().AddInput(START, ToField(\"a\", WithCustomExtractor(func(input any) (any, error) {\n\t\t\treturn input.([]int)[0], nil\n\t\t})))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tresult, err := r.Invoke(context.Background(), []int{1, 2})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]int{\"a\": 1}, result)\n\t})\n\n\tt.Run(\"mix custom extract with normal mapping\", func(t *testing.T) {\n\t\twf := NewWorkflow[map[string]any, map[string]int]()\n\t\twf.End().AddInput(START, ToField(\"a\", WithCustomExtractor(func(input any) (any, error) {\n\t\t\treturn input.(map[string]any)[\"a\"].([]any)[0].(map[string]any)[\"c\"], nil\n\t\t})), MapFields(\"b\", \"b\"))\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\t\tresult, err := r.Invoke(context.Background(), map[string]any{\n\t\t\t\"a\": []any{\n\t\t\t\tmap[string]any{\n\t\t\t\t\t\"c\": 1,\n\t\t\t\t},\n\t\t\t},\n\t\t\t\"b\": 2,\n\t\t})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, map[string]int{\"a\": 1, \"b\": 2}, result)\n\t})\n}\n\nfunc TestAddDependency(t *testing.T) {\n\tctx := context.Background()\n\n\twf := NewWorkflow[string, any]()\n\n\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {\n\t\treturn in + \"_\" + in, nil\n\t})).AddDependency(START)\n\n\twf.End().AddDependency(\"1\")\n\n\tr, err := wf.Compile(ctx)\n\tassert.NoError(t, err)\n\tout, err := r.Invoke(ctx, \"input\")\n\tassert.NoError(t, err)\n\tassert.Equal(t, nil, out)\n}\n\nfunc TestIndirectDependencyWithBranch(t *testing.T) {\n\tt.Run(\"data only mapping across branch\", func(t *testing.T) {\n\t\twf := NewWorkflow[[]int, map[string]any]()\n\n\t\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input int) (output int, err error) {\n\t\t\treturn input + 1, nil\n\t\t})).\n\t\t\tAddInputWithOptions(START, []*FieldMapping{ToField(\"\", WithCustomExtractor(func(input any) (any, error) {\n\t\t\t\tinputList := input.([]int)\n\t\t\t\tif len(inputList) == 0 {\n\t\t\t\t\treturn nil, fmt.Errorf(\"input list is empty\")\n\t\t\t\t}\n\t\t\t\treturn input.([]int)[0], nil\n\t\t\t}))}, WithNoDirectDependency())\n\n\t\twf.AddBranch(START, NewGraphBranch(func(ctx context.Context, in []int) (endNode string, err error) {\n\t\t\tif len(in) > 0 {\n\t\t\t\treturn \"1\", nil\n\t\t\t}\n\n\t\t\treturn END, nil\n\t\t}, map[string]bool{\"1\": true, END: true}))\n\n\t\twf.End().\n\t\t\tAddInput(\"1\", ToField(\"output\")).\n\t\t\tSetStaticValue(FieldPath{\"static\"}, 2)\n\n\t\tr, err := wf.Compile(context.Background())\n\t\tassert.NoError(t, err)\n\n\t\t// skip lambda node \"1\"\n\t\tout, err := r.Invoke(context.Background(), nil)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, out, map[string]any{\"static\": 2})\n\n\t\t// choose lambda node \"1\"\n\t\tout, err = r.Invoke(context.Background(), []int{1})\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, out, map[string]any{\"output\": 2, \"static\": 2})\n\t})\n\n\tt.Run(\"data only mapping across branch, with interrupt after branch\", func(t *testing.T) {\n\t\twf := NewWorkflow[[]int, map[string]any]()\n\n\t\twf.AddLambdaNode(\"1\", InvokableLambda(func(ctx context.Context, input int) (output int, err error) {\n\t\t\treturn input + 1, nil\n\t\t})).\n\t\t\tAddInputWithOptions(START, []*FieldMapping{ToField(\"\", WithCustomExtractor(func(input any) (any, error) {\n\t\t\t\tinputList := input.([]int)\n\t\t\t\tif len(inputList) == 0 {\n\t\t\t\t\treturn nil, fmt.Errorf(\"input list is empty\")\n\t\t\t\t}\n\t\t\t\treturn input.([]int)[0], nil\n\t\t\t}))}, WithNoDirectDependency())\n\n\t\twf.AddBranch(START, NewGraphBranch(func(ctx context.Context, in []int) (endNode string, err error) {\n\t\t\tif len(in) > 0 {\n\t\t\t\treturn \"1\", nil\n\t\t\t}\n\n\t\t\treturn END, nil\n\t\t}, map[string]bool{\"1\": true, END: true}))\n\n\t\twf.End().\n\t\t\tAddInput(\"1\", ToField(\"output\")).\n\t\t\tSetStaticValue(FieldPath{\"static\"}, 2)\n\n\t\tr, err := wf.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()),\n\t\t\tWithInterruptBeforeNodes([]string{\"1\"}))\n\t\tassert.NoError(t, err)\n\n\t\t// skip lambda node \"1\"\n\t\tout, err := r.Invoke(context.Background(), nil)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, out, map[string]any{\"static\": 2})\n\n\t\t// choose lambda node \"1\"\n\t\t_, err = r.Invoke(context.Background(), []int{1}, WithCheckPointID(\"123\"))\n\t\t_, ok := ExtractInterruptInfo(err)\n\t\tassert.True(t, ok)\n\n\t\tout, err = r.Invoke(context.Background(), nil, WithCheckPointID(\"123\"))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, out, map[string]any{\"output\": 2, \"static\": 2})\n\t})\n}\n"
  },
  {
    "path": "doc.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     https://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package eino provides building blocks for agent workflows,\n// tools, and composable graph utilities.\npackage eino\n"
  },
  {
    "path": "flow/agent/agent_option.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package agent defines common option types used by agents and multi-agents.\npackage agent\n\nimport \"github.com/cloudwego/eino/compose\"\n\n// AgentOption is the common option type for various agent and multi-agent implementations.\n// For options intended to use with underlying graph or components, use WithComposeOptions to specify.\n// For options intended to use with particular agent/multi-agent implementations, use WrapImplSpecificOptFn to specify.\ntype AgentOption struct {\n\timplSpecificOptFn any\n\tcomposeOptions    []compose.Option\n}\n\n// GetComposeOptions returns all compose options from the given agent options.\nfunc GetComposeOptions(opts ...AgentOption) []compose.Option {\n\tvar result []compose.Option\n\tfor _, opt := range opts {\n\t\tresult = append(result, opt.composeOptions...)\n\t}\n\n\treturn result\n}\n\n// WithComposeOptions returns an agent option that specifies compose options.\nfunc WithComposeOptions(opts ...compose.Option) AgentOption {\n\treturn AgentOption{\n\t\tcomposeOptions: opts,\n\t}\n}\n\n// WrapImplSpecificOptFn returns an agent option that specifies a function to modify the implementation-specific options.\nfunc WrapImplSpecificOptFn[T any](optFn func(*T)) AgentOption {\n\treturn AgentOption{\n\t\timplSpecificOptFn: optFn,\n\t}\n}\n\n// GetImplSpecificOptions returns the implementation-specific options from the given agent options.\nfunc GetImplSpecificOptions[T any](base *T, opts ...AgentOption) *T {\n\tif base == nil {\n\t\tbase = new(T)\n\t}\n\n\tfor i := range opts {\n\t\topt := opts[i]\n\t\tif opt.implSpecificOptFn != nil {\n\t\t\toptFn, ok := opt.implSpecificOptFn.(func(*T))\n\t\t\tif ok {\n\t\t\t\toptFn(base)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn base\n}\n"
  },
  {
    "path": "flow/agent/multiagent/host/callback.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage host\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/flow/agent\"\n\t\"github.com/cloudwego/eino/schema\"\n\ttemplate \"github.com/cloudwego/eino/utils/callbacks\"\n)\n\n// MultiAgentCallback is the callback interface for host multi-agent.\ntype MultiAgentCallback interface {\n\tOnHandOff(ctx context.Context, info *HandOffInfo) context.Context\n}\n\n// HandOffInfo is the info which will be passed to MultiAgentCallback.OnHandOff, representing a hand off event.\ntype HandOffInfo struct {\n\tToAgentName string\n\tArgument    string\n}\n\n// ConvertCallbackHandlers converts []host.MultiAgentCallback to callbacks.Handler.\nfunc ConvertCallbackHandlers(handlers ...MultiAgentCallback) callbacks.Handler {\n\tonChatModelEnd := func(ctx context.Context, info *callbacks.RunInfo, output *model.CallbackOutput) context.Context {\n\t\tmsg := output.Message\n\t\tif msg == nil || msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 {\n\t\t\treturn ctx\n\t\t}\n\n\t\tfor _, cb := range handlers {\n\t\t\tfor _, toolCall := range msg.ToolCalls {\n\t\t\t\tctx = cb.OnHandOff(ctx, &HandOffInfo{\n\t\t\t\t\tToAgentName: toolCall.Function.Name,\n\t\t\t\t\tArgument:    toolCall.Function.Arguments,\n\t\t\t\t})\n\t\t\t}\n\t\t}\n\n\t\treturn ctx\n\t}\n\n\tonChatModelEndWithStreamOutput := func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context {\n\t\tgo func() {\n\t\t\tmsg, err := schema.ConcatMessageStream(schema.StreamReaderWithConvert(output,\n\t\t\t\tfunc(m *model.CallbackOutput) (*schema.Message, error) {\n\t\t\t\t\treturn m.Message, nil\n\t\t\t\t}))\n\t\t\tif err != nil {\n\t\t\t\tfmt.Printf(\"concat message stream for host multi-agent failed: %v\", err)\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\tfor _, cb := range handlers {\n\t\t\t\tfor _, tc := range msg.ToolCalls {\n\t\t\t\t\t_ = cb.OnHandOff(ctx, &HandOffInfo{\n\t\t\t\t\t\tToAgentName: tc.Function.Name,\n\t\t\t\t\t\tArgument:    tc.Function.Arguments,\n\t\t\t\t\t})\n\t\t\t\t}\n\t\t\t}\n\t\t}()\n\n\t\treturn ctx\n\t}\n\n\treturn template.NewHandlerHelper().ChatModel(&template.ModelCallbackHandler{\n\t\tOnEnd:                 onChatModelEnd,\n\t\tOnEndWithStreamOutput: onChatModelEndWithStreamOutput,\n\t}).Handler()\n}\n\n// convertCallbacks reads graph call options, extract host.MultiAgentCallback and convert it to callbacks.Handler.\nfunc convertCallbacks(opts ...agent.AgentOption) callbacks.Handler {\n\tagentOptions := agent.GetImplSpecificOptions(&options{}, opts...)\n\tif len(agentOptions.agentCallbacks) == 0 {\n\t\treturn nil\n\t}\n\n\thandlers := agentOptions.agentCallbacks\n\treturn ConvertCallbackHandlers(handlers...)\n}\n"
  },
  {
    "path": "flow/agent/multiagent/host/compose.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage host\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/flow/agent\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nconst (\n\tdefaultHostNodeKey                 = \"host\" // the key of the host node in the graph\n\tdefaultHostPrompt                  = \"decide which tool is best for the task and call only the best tool.\"\n\tspecialistsAnswersCollectorNodeKey = \"specialist_answers_collect\"\n\tsingleIntentAnswerNodeKey          = \"single_intent_answer\"\n\tmultiIntentSummarizeNodeKey        = \"multi_intents_summarize\"\n\tdefaultSummarizerPrompt            = \"summarize the answers from the specialists into a single answer.\"\n\tmap2ListConverterNodeKey           = \"map_to_list\"\n)\n\ntype state struct {\n\tmsgs              []*schema.Message\n\tisMultipleIntents bool\n}\n\n// NewMultiAgent creates a new host multi-agent system.\n//\n// IMPORTANT!! For models that don't output tool calls in the first streaming chunk (e.g. Claude)\n// the default StreamToolCallChecker may not work properly since it only checks the first chunk for tool calls.\n// In such cases, you need to implement a custom StreamToolCallChecker that can properly detect tool calls.\nfunc NewMultiAgent(ctx context.Context, config *MultiAgentConfig) (*MultiAgent, error) {\n\tif err := config.validate(); err != nil {\n\t\treturn nil, err\n\t}\n\n\thostKeyName := defaultHostNodeKey\n\tif config.HostNodeName != \"\" {\n\t\thostKeyName = config.HostNodeName\n\t}\n\n\tvar (\n\t\thostPrompt      = config.Host.SystemPrompt\n\t\tname            = config.Name\n\t\ttoolCallChecker = config.StreamToolCallChecker\n\t)\n\n\tif len(hostPrompt) == 0 {\n\t\thostPrompt = defaultHostPrompt\n\t}\n\n\tif len(name) == 0 {\n\t\tname = \"host multi agent\"\n\t}\n\n\tif toolCallChecker == nil {\n\t\ttoolCallChecker = firstChunkStreamToolCallChecker\n\t}\n\n\tg := compose.NewGraph[[]*schema.Message, *schema.Message](\n\t\tcompose.WithGenLocalState(func(context.Context) *state { return &state{} }))\n\n\tif err := g.AddPassthroughNode(specialistsAnswersCollectorNodeKey); err != nil {\n\t\treturn nil, err\n\t}\n\n\tagentTools := make([]*schema.ToolInfo, 0, len(config.Specialists))\n\tagentMap := make(map[string]bool, len(config.Specialists)+1)\n\tfor i := range config.Specialists {\n\t\tspecialist := config.Specialists[i]\n\n\t\tagentTools = append(agentTools, &schema.ToolInfo{\n\t\t\tName: specialist.Name,\n\t\t\tDesc: specialist.IntendedUse,\n\t\t\tParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{\n\t\t\t\t\"reason\": {\n\t\t\t\t\tType: schema.String,\n\t\t\t\t\tDesc: \"the reason to call this tool\",\n\t\t\t\t},\n\t\t\t}),\n\t\t})\n\n\t\tif err := addSpecialistAgent(specialist, g); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tagentMap[specialist.Name] = true\n\t}\n\n\tchatModel, err := agent.ChatModelWithTools(config.Host.ChatModel, config.Host.ToolCallingModel, agentTools)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tif err = addHostAgent(chatModel, hostPrompt, g, hostKeyName); err != nil {\n\t\treturn nil, err\n\t}\n\n\tconst convertorName = \"msg2MsgList\"\n\tif err = g.AddLambdaNode(convertorName, compose.ToList[*schema.Message](), compose.WithNodeName(\"converter\")); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif err = addDirectAnswerBranch(convertorName, g, toolCallChecker); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif err = addMultiSpecialistsBranch(convertorName, agentMap, g); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif err = addSingleIntentAnswerNode(g); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif err = addMultiIntentsSummarizeNode(config.Summarizer, g); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif err = addAfterSpecialistsBranch(g); err != nil {\n\t\treturn nil, err\n\t}\n\n\tcompileOpts := []compose.GraphCompileOption{compose.WithNodeTriggerMode(compose.AnyPredecessor), compose.WithGraphName(name)}\n\tr, err := g.Compile(ctx, compileOpts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &MultiAgent{\n\t\trunnable:         r,\n\t\tgraph:            g,\n\t\tgraphAddNodeOpts: []compose.GraphAddNodeOpt{compose.WithGraphCompileOptions(compileOpts...)},\n\t}, nil\n}\n\nfunc addSpecialistAgent(specialist *Specialist, g *compose.Graph[[]*schema.Message, *schema.Message]) error {\n\tif specialist.Invokable != nil || specialist.Streamable != nil {\n\t\tlambda, err := compose.AnyLambda(specialist.Invokable, specialist.Streamable, nil, nil, compose.WithLambdaType(\"Specialist\"))\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tpreHandler := func(_ context.Context, input []*schema.Message, state *state) ([]*schema.Message, error) {\n\t\t\treturn state.msgs, nil // replace the tool call message with input msgs stored in state\n\t\t}\n\t\tif err := g.AddLambdaNode(specialist.Name, lambda, compose.WithStatePreHandler(preHandler),\n\t\t\tcompose.WithNodeName(specialist.Name), compose.WithOutputKey(specialist.Name)); err != nil {\n\t\t\treturn err\n\t\t}\n\t} else if specialist.ChatModel != nil {\n\t\tpreHandler := func(_ context.Context, input []*schema.Message, state *state) ([]*schema.Message, error) {\n\t\t\tif len(specialist.SystemPrompt) > 0 {\n\t\t\t\treturn append([]*schema.Message{{\n\t\t\t\t\tRole:    schema.System,\n\t\t\t\t\tContent: specialist.SystemPrompt,\n\t\t\t\t}}, state.msgs...), nil\n\t\t\t}\n\n\t\t\treturn state.msgs, nil // replace the tool call message with input msgs stored in state\n\t\t}\n\n\t\tif err := g.AddChatModelNode(specialist.Name, specialist.ChatModel, compose.WithStatePreHandler(preHandler), compose.WithNodeName(specialist.Name), compose.WithOutputKey(specialist.Name)); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\treturn g.AddEdge(specialist.Name, specialistsAnswersCollectorNodeKey)\n}\n\nfunc addHostAgent(model model.BaseChatModel, prompt string, g *compose.Graph[[]*schema.Message, *schema.Message], hostNodeName string) error {\n\tpreHandler := func(_ context.Context, input []*schema.Message, state *state) ([]*schema.Message, error) {\n\t\tstate.msgs = input\n\t\tif len(prompt) == 0 {\n\t\t\treturn input, nil\n\t\t}\n\t\treturn append([]*schema.Message{{\n\t\t\tRole:    schema.System,\n\t\t\tContent: prompt,\n\t\t}}, input...), nil\n\t}\n\tif err := g.AddChatModelNode(defaultHostNodeKey, model, compose.WithStatePreHandler(preHandler), compose.WithNodeName(hostNodeName)); err != nil {\n\t\treturn err\n\t}\n\n\treturn g.AddEdge(compose.START, defaultHostNodeKey)\n}\n\nfunc addDirectAnswerBranch(convertorName string, g *compose.Graph[[]*schema.Message, *schema.Message],\n\ttoolCallChecker func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error)) error {\n\t// handles the case where the host agent returns a direct answer, instead of handling off to any specialist\n\tbranch := compose.NewStreamGraphBranch(func(ctx context.Context, sr *schema.StreamReader[*schema.Message]) (endNode string, err error) {\n\t\tisToolCall, err := toolCallChecker(ctx, sr)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tif isToolCall {\n\t\t\treturn convertorName, nil\n\t\t}\n\t\treturn compose.END, nil\n\t}, map[string]bool{convertorName: true, compose.END: true})\n\n\treturn g.AddBranch(defaultHostNodeKey, branch)\n}\n\nfunc addMultiSpecialistsBranch(convertorName string, agentMap map[string]bool, g *compose.Graph[[]*schema.Message, *schema.Message]) error {\n\tbranch := compose.NewGraphMultiBranch(func(ctx context.Context, input []*schema.Message) (map[string]bool, error) {\n\t\tif len(input) != 1 {\n\t\t\treturn nil, fmt.Errorf(\"host agent output %d messages, but expected 1\", len(input))\n\t\t}\n\n\t\tresults := map[string]bool{}\n\t\tfor _, toolCall := range input[0].ToolCalls {\n\t\t\tresults[toolCall.Function.Name] = true\n\t\t}\n\n\t\tif len(results) > 1 {\n\t\t\t_ = compose.ProcessState(ctx, func(_ context.Context, state *state) error {\n\t\t\t\tstate.isMultipleIntents = true\n\t\t\t\treturn nil\n\t\t\t})\n\t\t}\n\n\t\treturn results, nil\n\t}, agentMap)\n\n\treturn g.AddBranch(convertorName, branch)\n}\n\nfunc addSingleIntentAnswerNode(g *compose.Graph[[]*schema.Message, *schema.Message]) error {\n\trc := func(ctx context.Context, input *schema.StreamReader[map[string]any]) (*schema.StreamReader[*schema.Message], error) {\n\t\treturn schema.StreamReaderWithConvert(input, func(msgs map[string]any) (*schema.Message, error) {\n\t\t\tif len(msgs) != 1 {\n\t\t\t\treturn nil, fmt.Errorf(\"host agent output %d messages, but expected 1\", len(msgs))\n\t\t\t}\n\t\t\tfor _, msg := range msgs {\n\t\t\t\treturn msg.(*schema.Message), nil\n\t\t\t}\n\t\t\treturn nil, schema.ErrNoValue\n\t\t}), nil\n\t}\n\n\t_ = g.AddLambdaNode(singleIntentAnswerNodeKey, compose.TransformableLambda(rc))\n\treturn g.AddEdge(singleIntentAnswerNodeKey, compose.END)\n}\n\nfunc addAfterSpecialistsBranch(g *compose.Graph[[]*schema.Message, *schema.Message]) error {\n\tab := func(ctx context.Context, _ *schema.StreamReader[map[string]any]) (string, error) {\n\t\tvar isMultipleIntents bool\n\t\t_ = compose.ProcessState(ctx, func(_ context.Context, state *state) error {\n\t\t\tisMultipleIntents = state.isMultipleIntents\n\t\t\treturn nil\n\t\t})\n\n\t\tif !isMultipleIntents {\n\t\t\treturn singleIntentAnswerNodeKey, nil\n\t\t}\n\n\t\treturn map2ListConverterNodeKey, nil\n\t}\n\n\tb := compose.NewStreamGraphBranch(ab, map[string]bool{\n\t\tsingleIntentAnswerNodeKey: true,\n\t\tmap2ListConverterNodeKey:  true,\n\t})\n\n\treturn g.AddBranch(specialistsAnswersCollectorNodeKey, b)\n}\n\nfunc addMultiIntentsSummarizeNode(summarizer *Summarizer, g *compose.Graph[[]*schema.Message, *schema.Message]) error {\n\tmap2list := func(ctx context.Context, input map[string]any) ([]*schema.Message, error) {\n\t\tvar output []*schema.Message\n\t\tfor k := range input {\n\t\t\toutput = append(output, input[k].(*schema.Message))\n\t\t}\n\t\treturn output, nil\n\t}\n\n\t_ = g.AddLambdaNode(map2ListConverterNodeKey, compose.InvokableLambda(map2list))\n\n\tif summarizer != nil {\n\t\t_ = g.AddChatModelNode(multiIntentSummarizeNodeKey, summarizer.ChatModel,\n\t\t\tcompose.WithStatePreHandler(func(ctx context.Context, in []*schema.Message, state *state) ([]*schema.Message, error) {\n\t\t\t\tvar (\n\t\t\t\t\tout          []*schema.Message\n\t\t\t\t\tsystemPrompt = defaultSummarizerPrompt\n\t\t\t\t)\n\t\t\t\tif summarizer.SystemPrompt != \"\" {\n\t\t\t\t\tsystemPrompt = summarizer.SystemPrompt\n\t\t\t\t}\n\t\t\t\tout = append(out, &schema.Message{\n\t\t\t\t\tRole:    schema.System,\n\t\t\t\t\tContent: systemPrompt,\n\t\t\t\t})\n\n\t\t\t\tout = append(out, state.msgs...)\n\t\t\t\tout = append(out, in...)\n\t\t\t\treturn out, nil\n\t\t\t}))\n\t\t_ = g.AddEdge(map2ListConverterNodeKey, multiIntentSummarizeNodeKey)\n\t\treturn g.AddEdge(multiIntentSummarizeNodeKey, compose.END)\n\t}\n\n\ts := func(ctx context.Context, input []*schema.Message) (*schema.Message, error) {\n\t\toutput := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t}\n\n\t\tfor _, msg := range input {\n\t\t\toutput.Content += msg.Content + \"\\n\"\n\t\t}\n\n\t\treturn output, nil\n\t}\n\n\t_ = g.AddLambdaNode(multiIntentSummarizeNodeKey, compose.InvokableLambda(s))\n\t_ = g.AddEdge(map2ListConverterNodeKey, multiIntentSummarizeNodeKey)\n\treturn g.AddEdge(multiIntentSummarizeNodeKey, compose.END)\n}\n"
  },
  {
    "path": "flow/agent/multiagent/host/compose_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage host\n\nimport (\n\t\"context\"\n\t\"io\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/flow/agent\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestHostMultiAgent(t *testing.T) {\n\tctrl := gomock.NewController(t)\n\tmockHostLLM := model.NewMockToolCallingChatModel(ctrl)\n\tmockSpecialistLLM1 := model.NewMockChatModel(ctrl)\n\n\tspecialist1 := &Specialist{\n\t\tChatModel:    mockSpecialistLLM1,\n\t\tSystemPrompt: \"You are a helpful assistant.\",\n\t\tAgentMeta: AgentMeta{\n\t\t\tName:        \"specialist 1\",\n\t\t\tIntendedUse: \"do stuff that works\",\n\t\t},\n\t}\n\n\tspecialist2Msg1 := &schema.Message{\n\t\tRole:    schema.Assistant,\n\t\tContent: \"specialist2\",\n\t}\n\tspecialist2Msg2 := &schema.Message{\n\t\tRole:    schema.Assistant,\n\t\tContent: \" stream answer\",\n\t}\n\n\tspecialist2 := &Specialist{\n\t\tInvokable: func(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.Message, error) {\n\t\t\treturn &schema.Message{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tContent: \"specialist2 invoke answer\",\n\t\t\t}, nil\n\t\t},\n\t\tStreamable: func(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.StreamReader[*schema.Message], error) {\n\t\t\treturn schema.StreamReaderFromArray([]*schema.Message{specialist2Msg1, specialist2Msg2}), nil\n\t\t},\n\t\tAgentMeta: AgentMeta{\n\t\t\tName:        \"specialist 2\",\n\t\t\tIntendedUse: \"do stuff that works too\",\n\t\t},\n\t}\n\n\tctx := context.Background()\n\n\tmockHostLLM.EXPECT().WithTools(gomock.Any()).Return(mockHostLLM, nil).AnyTimes()\n\n\thostMA, err := NewMultiAgent(ctx, &MultiAgentConfig{\n\t\tHost: Host{\n\t\t\tToolCallingModel: mockHostLLM,\n\t\t},\n\t\tSpecialists: []*Specialist{\n\t\t\tspecialist1,\n\t\t\tspecialist2,\n\t\t},\n\t})\n\n\tassert.NoError(t, err)\n\n\tt.Run(\"generate direct answer from host\", func(t *testing.T) {\n\t\tdirectAnswerMsg := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"direct answer\",\n\t\t}\n\n\t\tmockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(directAnswerMsg, nil).Times(1)\n\n\t\tmockCallback := newMockAgentCallback(0)\n\n\t\tout, err := hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"direct answer\", out.Content)\n\t\tassert.Empty(t, mockCallback.infos)\n\t})\n\n\tt.Run(\"stream direct answer from host\", func(t *testing.T) {\n\t\tdirectAnswerMsg1 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"direct \",\n\t\t}\n\n\t\tdirectAnswerMsg2 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"answer\",\n\t\t}\n\n\t\tsr, sw := schema.Pipe[*schema.Message](0)\n\t\tgo func() {\n\t\t\tsw.Send(directAnswerMsg1, nil)\n\t\t\tsw.Send(directAnswerMsg2, nil)\n\t\t\tsw.Close()\n\t\t}()\n\n\t\tmockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\n\t\tmockCallback := newMockAgentCallback(0)\n\t\toutStream, err := hostMA.Stream(ctx, nil, WithAgentCallbacks(mockCallback))\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, mockCallback.infos)\n\n\t\tvar msgs []*schema.Message\n\t\tfor {\n\t\t\tmsg, err := outStream.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\tmsgs = append(msgs, msg)\n\t\t}\n\n\t\toutStream.Close()\n\n\t\tassert.Equal(t, directAnswerMsg1, msgs[0])\n\t\tassert.Equal(t, directAnswerMsg2, msgs[1])\n\t})\n\n\tt.Run(\"generate hand off\", func(t *testing.T) {\n\t\thandOffMsg := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      specialist1.Name,\n\t\t\t\t\t\tArguments: `{\"reason\": \"specialist 1 is the best\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tspecialistMsg := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"specialist 1 answer\",\n\t\t}\n\n\t\tmockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(1)\n\t\tmockSpecialistLLM1.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(specialistMsg, nil).Times(1)\n\n\t\tmockCallback := newMockAgentCallback(1)\n\n\t\tout, err := hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"specialist 1 answer\", out.Content)\n\t\tmockCallback.wg.Wait()\n\t\tassert.Equal(t, []*HandOffInfo{\n\t\t\t{\n\t\t\t\tToAgentName: specialist1.Name,\n\t\t\t\tArgument:    `{\"reason\": \"specialist 1 is the best\"}`,\n\t\t\t},\n\t\t}, mockCallback.infos)\n\n\t\thandOffMsg.ToolCalls[0].Function.Name = specialist2.Name\n\t\thandOffMsg.ToolCalls[0].Function.Arguments = `{\"reason\": \"specialist 2 is even better\"}`\n\t\tmockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(1)\n\n\t\tmockCallback = newMockAgentCallback(1)\n\n\t\tout, err = hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"specialist2 invoke answer\", out.Content)\n\t\tmockCallback.wg.Wait()\n\t\tassert.Equal(t, []*HandOffInfo{\n\t\t\t{\n\t\t\t\tToAgentName: specialist2.Name,\n\t\t\t\tArgument:    `{\"reason\": \"specialist 2 is even better\"}`,\n\t\t\t},\n\t\t}, mockCallback.infos)\n\t})\n\n\tt.Run(\"stream hand off to chat model\", func(t *testing.T) {\n\t\thandOffMsg1 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"need to call function\",\n\t\t}\n\n\t\thandOffMsg2 := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\thandOffMsg3 := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex:    generic.PtrOf(0),\n\t\t\t\t\tFunction: schema.FunctionCall{},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\thandOffMsg4 := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      specialist1.Name,\n\t\t\t\t\t\tArguments: `{\"reason\": \"specialist 1 is the best\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tsr, sw := schema.Pipe[*schema.Message](0)\n\t\tgo func() {\n\t\t\tsw.Send(handOffMsg1, nil)\n\t\t\tsw.Send(handOffMsg2, nil)\n\t\t\tsw.Send(handOffMsg3, nil)\n\t\t\tsw.Send(handOffMsg4, nil)\n\t\t\tsw.Close()\n\t\t}()\n\n\t\tspecialistMsg1 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"specialist \",\n\t\t}\n\n\t\tspecialistMsg2 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"1 answer\",\n\t\t}\n\n\t\tsr1, sw1 := schema.Pipe[*schema.Message](0)\n\t\tgo func() {\n\t\t\tsw1.Send(specialistMsg1, nil)\n\t\t\tsw1.Send(specialistMsg2, nil)\n\t\t\tsw1.Close()\n\t\t}()\n\n\t\tstreamToolCallChecker := func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error) {\n\t\t\tdefer modelOutput.Close()\n\n\t\t\tfor {\n\t\t\t\tmsg, err := modelOutput.Recv()\n\t\t\t\tif err != nil {\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\treturn false, nil\n\t\t\t\t\t}\n\n\t\t\t\t\treturn false, err\n\t\t\t\t}\n\n\t\t\t\tif len(msg.ToolCalls) == 0 {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\n\t\t\t\tif len(msg.ToolCalls) > 0 {\n\t\t\t\t\treturn true, nil\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\thostMA, err = NewMultiAgent(ctx, &MultiAgentConfig{\n\t\t\tHost: Host{\n\t\t\t\tToolCallingModel: mockHostLLM,\n\t\t\t},\n\t\t\tSpecialists: []*Specialist{\n\t\t\t\tspecialist1,\n\t\t\t\tspecialist2,\n\t\t\t},\n\t\t\tStreamToolCallChecker: streamToolCallChecker,\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tmockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\t\tmockSpecialistLLM1.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr1, nil).Times(1)\n\n\t\tmockCallback := newMockAgentCallback(1)\n\t\toutStream, err := hostMA.Stream(ctx, nil, WithAgentCallbacks(mockCallback))\n\t\tassert.NoError(t, err)\n\n\t\tvar msgs []*schema.Message\n\t\tfor {\n\t\t\tmsg, err := outStream.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\tmsgs = append(msgs, msg)\n\t\t}\n\n\t\toutStream.Close()\n\n\t\tassert.Equal(t, specialistMsg1, msgs[0])\n\t\tassert.Equal(t, specialistMsg2, msgs[1])\n\n\t\tmockCallback.wg.Wait()\n\n\t\tassert.Equal(t, []*HandOffInfo{\n\t\t\t{\n\t\t\t\tToAgentName: specialist1.Name,\n\t\t\t\tArgument:    `{\"reason\": \"specialist 1 is the best\"}`,\n\t\t\t},\n\t\t}, mockCallback.infos)\n\n\t\thandOffMsg4.ToolCalls[0].Function.Name = specialist2.Name\n\t\thandOffMsg4.ToolCalls[0].Function.Arguments = `{\"reason\": \"specialist 2 is even better\"}`\n\t\tsr, sw = schema.Pipe[*schema.Message](0)\n\t\tgo func() {\n\t\t\tsw.Send(handOffMsg1, nil)\n\t\t\tsw.Send(handOffMsg2, nil)\n\t\t\tsw.Send(handOffMsg3, nil)\n\t\t\tsw.Send(handOffMsg4, nil)\n\t\t\tsw.Close()\n\t\t}()\n\n\t\tmockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\n\t\tmockCallback = newMockAgentCallback(1)\n\t\toutStream, err = hostMA.Stream(ctx, nil, WithAgentCallbacks(mockCallback))\n\t\tassert.NoError(t, err)\n\n\t\tmsgs = nil\n\t\tfor {\n\t\t\tmsg, err := outStream.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\tmsgs = append(msgs, msg)\n\t\t}\n\n\t\toutStream.Close()\n\n\t\tassert.Equal(t, specialist2Msg1, msgs[0])\n\t\tassert.Equal(t, specialist2Msg2, msgs[1])\n\n\t\tmockCallback.wg.Wait()\n\n\t\tassert.Equal(t, []*HandOffInfo{\n\t\t\t{\n\t\t\t\tToAgentName: specialist2.Name,\n\t\t\t\tArgument:    `{\"reason\": \"specialist 2 is even better\"}`,\n\t\t\t},\n\t\t}, mockCallback.infos)\n\t})\n\n\tt.Run(\"multi-agent within graph\", func(t *testing.T) {\n\t\thandOffMsg := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      specialist1.Name,\n\t\t\t\t\t\tArguments: `{\"reason\": \"specialist 1 is the best\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tspecialistMsg := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"Beijing\",\n\t\t}\n\n\t\tmockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(1)\n\t\tmockSpecialistLLM1.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(specialistMsg, nil).Times(1)\n\n\t\tmockCallback := newMockAgentCallback(1)\n\n\t\thostMA, err := NewMultiAgent(ctx, &MultiAgentConfig{\n\t\t\tHost: Host{\n\t\t\t\tToolCallingModel: mockHostLLM,\n\t\t\t},\n\t\t\tSpecialists: []*Specialist{\n\t\t\t\tspecialist1,\n\t\t\t\tspecialist2,\n\t\t\t},\n\t\t})\n\n\t\tassert.NoError(t, err)\n\n\t\tmaGraph, opts := hostMA.ExportGraph()\n\n\t\tfullGraph, err := compose.NewChain[map[string]any, *schema.Message]().\n\t\t\tAppendChatTemplate(prompt.FromMessages(schema.FString, schema.UserMessage(\"what's the capital city of {country_name}\"))).\n\t\t\tAppendGraph(maGraph, append(opts, compose.WithNodeKey(\"host_ma_node\"))...).\n\t\t\tCompile(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tout, err := fullGraph.Invoke(ctx, map[string]any{\"country_name\": \"China\"}, compose.WithCallbacks(ConvertCallbackHandlers(mockCallback)).DesignateNodeWithPath(compose.NewNodePath(\"host_ma_node\", hostMA.HostNodeKey())))\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"Beijing\", out.Content)\n\n\t\tmockCallback.wg.Wait()\n\t\tassert.Equal(t, []*HandOffInfo{\n\t\t\t{\n\t\t\t\tToAgentName: specialist1.Name,\n\t\t\t\tArgument:    `{\"reason\": \"specialist 1 is the best\"}`,\n\t\t\t},\n\t\t}, mockCallback.infos)\n\t})\n\n\tt.Run(\"multiple intents\", func(t *testing.T) {\n\t\thandOffMsg1 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"need to call function\",\n\t\t}\n\n\t\thandOffMsg2 := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\thandOffMsg3 := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex:    generic.PtrOf(0),\n\t\t\t\t\tFunction: schema.FunctionCall{},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\thandOffMsg4 := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      specialist1.Name,\n\t\t\t\t\t\tArguments: `{\"reason\": \"specialist 1 is good\"}`,\n\t\t\t\t\t},\n\t\t\t\t}, {\n\t\t\t\t\tIndex: generic.PtrOf(1),\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      specialist2.Name,\n\t\t\t\t\t\tArguments: `{\"reason\": \"specialist 2`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\thandOffMsg5 := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(1),\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tArguments: ` is also good\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tsr := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\thandOffMsg1,\n\t\t\thandOffMsg2,\n\t\t\thandOffMsg3,\n\t\t\thandOffMsg4,\n\t\t\thandOffMsg5,\n\t\t})\n\n\t\tspecialist1Msg1 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"specialist \",\n\t\t}\n\n\t\tspecialist1Msg2 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"1 answer\",\n\t\t}\n\n\t\tsr1 := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\tspecialist1Msg1,\n\t\t\tspecialist1Msg2,\n\t\t})\n\n\t\tstreamToolCallChecker := func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error) {\n\t\t\tdefer modelOutput.Close()\n\n\t\t\tfor {\n\t\t\t\tmsg, err := modelOutput.Recv()\n\t\t\t\tif err != nil {\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\treturn false, nil\n\t\t\t\t\t}\n\n\t\t\t\t\treturn false, err\n\t\t\t\t}\n\n\t\t\t\tif len(msg.ToolCalls) == 0 {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\n\t\t\t\tif len(msg.ToolCalls) > 0 {\n\t\t\t\t\treturn true, nil\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\thostMA, err = NewMultiAgent(ctx, &MultiAgentConfig{\n\t\t\tHost: Host{\n\t\t\t\tToolCallingModel: mockHostLLM,\n\t\t\t},\n\t\t\tSpecialists: []*Specialist{\n\t\t\t\tspecialist1,\n\t\t\t\tspecialist2,\n\t\t\t},\n\t\t\tStreamToolCallChecker: streamToolCallChecker,\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tmockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\t\tmockSpecialistLLM1.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr1, nil).Times(1)\n\n\t\tmockCallback := newMockAgentCallback(2)\n\t\toutStream, err := hostMA.Stream(ctx, nil, WithAgentCallbacks(mockCallback))\n\t\tassert.NoError(t, err)\n\n\t\tvar msgs []*schema.Message\n\t\tfor {\n\t\t\tmsg, err := outStream.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\tmsgs = append(msgs, msg)\n\t\t}\n\n\t\toutStream.Close()\n\n\t\tmsg, err := schema.ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\t\tif msg.Content != \"specialist2 stream answer\\nspecialist 1 answer\\n\" &&\n\t\t\tmsg.Content != \"specialist 1 answer\\nspecialist2 stream answer\\n\" {\n\t\t\tt.Errorf(\"Unexpected message content: %s\", msg.Content)\n\t\t}\n\n\t\tmockCallback.wg.Wait()\n\t\tassert.Equal(t, []*HandOffInfo{\n\t\t\t{\n\t\t\t\tToAgentName: specialist1.Name,\n\t\t\t\tArgument:    `{\"reason\": \"specialist 1 is good\"}`,\n\t\t\t},\n\t\t\t{\n\t\t\t\tToAgentName: specialist2.Name,\n\t\t\t\tArgument:    `{\"reason\": \"specialist 2 is also good\"}`,\n\t\t\t},\n\t\t}, mockCallback.infos)\n\t})\n\n\tt.Run(\"summarize multiple intents\", func(t *testing.T) {\n\n\t\thandOffMsg := &schema.Message{\n\t\t\tRole: schema.Assistant,\n\t\t\tToolCalls: []schema.ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      specialist1.Name,\n\t\t\t\t\t\tArguments: `{\"reason\": \"specialist 1 is good\"}`,\n\t\t\t\t\t},\n\t\t\t\t}, {\n\t\t\t\t\tIndex: generic.PtrOf(1),\n\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\tName:      specialist2.Name,\n\t\t\t\t\t\tArguments: `{\"reason\": \"specialist 2 is also good\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tsr := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\thandOffMsg,\n\t\t})\n\n\t\tspecialist1Msg1 := &schema.Message{\n\t\t\tRole:    schema.Assistant,\n\t\t\tContent: \"specialist 1 answer\",\n\t\t}\n\t\tsr1 := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\tspecialist1Msg1,\n\t\t})\n\n\t\tconst summaryContent = \"summarized answer\"\n\t\tsr2 := schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t{\n\t\t\t\tRole:    schema.Assistant,\n\t\t\t\tContent: summaryContent,\n\t\t\t},\n\t\t})\n\n\t\tmockSumChatModel := model.NewMockChatModel(ctrl)\n\t\thostMA, err = NewMultiAgent(ctx, &MultiAgentConfig{\n\t\t\tHost: Host{\n\t\t\t\tToolCallingModel: mockHostLLM,\n\t\t\t},\n\t\t\tSpecialists: []*Specialist{\n\t\t\t\tspecialist1,\n\t\t\t\tspecialist2,\n\t\t\t},\n\t\t\tSummarizer: &Summarizer{\n\t\t\t\tChatModel: mockSumChatModel,\n\t\t\t},\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tmockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1)\n\t\tmockSpecialistLLM1.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr1, nil).Times(1)\n\t\tmockSumChatModel.EXPECT().Stream(gomock.Any(), gomock.Cond(func(x any) bool {\n\t\t\treturn assert.Equal(t, defaultSummarizerPrompt, func() string {\n\t\t\t\tif input := x.([]*schema.Message); len(input) > 0 {\n\t\t\t\t\treturn input[0].Content\n\t\t\t\t}\n\t\t\t\treturn \"\"\n\t\t\t}())\n\t\t})).Return(sr2, nil).Times(1)\n\n\t\toutStream, err := hostMA.Stream(ctx, nil)\n\t\tassert.NoError(t, err)\n\n\t\tvar msgs []*schema.Message\n\t\tfor {\n\t\t\tmsg, err := outStream.Recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t\tmsgs = append(msgs, msg)\n\t\t}\n\t\toutStream.Close()\n\n\t\tmsg, err := schema.ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\t\tif msg.Content != summaryContent {\n\t\t\tt.Errorf(\"Unexpected message content: %s\", msg.Content)\n\t\t}\n\t},\n\t)\n}\n\ntype mockAgentCallback struct {\n\tinfos []*HandOffInfo\n\twg    sync.WaitGroup\n}\n\nfunc (m *mockAgentCallback) OnHandOff(ctx context.Context, info *HandOffInfo) context.Context {\n\tm.infos = append(m.infos, info)\n\tm.wg.Done()\n\treturn ctx\n}\n\nfunc newMockAgentCallback(expects int) *mockAgentCallback {\n\tm := &mockAgentCallback{\n\t\tinfos: make([]*HandOffInfo, 0),\n\t\twg:    sync.WaitGroup{},\n\t}\n\n\tm.wg.Add(expects)\n\treturn m\n}\n"
  },
  {
    "path": "flow/agent/multiagent/host/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage host\n"
  },
  {
    "path": "flow/agent/multiagent/host/options.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage host\n\nimport \"github.com/cloudwego/eino/flow/agent\"\n\ntype options struct {\n\tagentCallbacks []MultiAgentCallback\n}\n\n// WithAgentCallbacks registers callbacks to be invoked by the host multi-agent.\nfunc WithAgentCallbacks(agentCallbacks ...MultiAgentCallback) agent.AgentOption {\n\treturn agent.WrapImplSpecificOptFn(func(opts *options) {\n\t\topts.agentCallbacks = append(opts.agentCallbacks, agentCallbacks...)\n\t})\n}\n"
  },
  {
    "path": "flow/agent/multiagent/host/types.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package host implements the host pattern for multi-agent system.\npackage host\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/flow/agent\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// MultiAgent is a host multi-agent system.\n// A host agent is responsible for deciding which specialist to 'hand off' the task to.\n// One or more specialist agents are responsible for completing the task.\ntype MultiAgent struct {\n\trunnable         compose.Runnable[[]*schema.Message, *schema.Message]\n\tgraph            *compose.Graph[[]*schema.Message, *schema.Message]\n\tgraphAddNodeOpts []compose.GraphAddNodeOpt\n}\n\n// Generate runs the multi-agent synchronously and returns the final message.\nfunc (ma *MultiAgent) Generate(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.Message, error) {\n\tcomposeOptions := agent.GetComposeOptions(opts...)\n\n\thandler := convertCallbacks(opts...)\n\tif handler != nil {\n\t\tcomposeOptions = append(composeOptions, compose.WithCallbacks(handler).DesignateNode(ma.HostNodeKey()))\n\t}\n\n\treturn ma.runnable.Invoke(ctx, input, composeOptions...)\n}\n\n// Stream runs the multi-agent in streaming mode and returns a message stream.\nfunc (ma *MultiAgent) Stream(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.StreamReader[*schema.Message], error) {\n\tcomposeOptions := agent.GetComposeOptions(opts...)\n\n\thandler := convertCallbacks(opts...)\n\tif handler != nil {\n\t\tcomposeOptions = append(composeOptions, compose.WithCallbacks(handler).DesignateNode(ma.HostNodeKey()))\n\t}\n\n\treturn ma.runnable.Stream(ctx, input, composeOptions...)\n}\n\n// ExportGraph exports the underlying graph from MultiAgent, along with the []compose.GraphAddNodeOpt to be used when adding this graph to another graph.\nfunc (ma *MultiAgent) ExportGraph() (compose.AnyGraph, []compose.GraphAddNodeOpt) {\n\treturn ma.graph, ma.graphAddNodeOpts\n}\n\n// HostNodeKey returns the graph node key used for the host agent.\nfunc (ma *MultiAgent) HostNodeKey() string {\n\treturn defaultHostNodeKey\n}\n\n// MultiAgentConfig is the config for host multi-agent system.\ntype MultiAgentConfig struct {\n\tHost        Host\n\tSpecialists []*Specialist\n\n\tName         string // the name of the host multi-agent\n\tHostNodeName string // the name of the host node in the graph, default is \"host\"\n\t// StreamToolCallChecker is a function to determine whether the model's streaming output contains tool calls.\n\t// Different models have different ways of outputting tool calls in streaming mode:\n\t// - Some models (like OpenAI) output tool calls directly\n\t// - Others (like Claude) output text first, then tool calls\n\t// This handler allows custom logic to check for tool calls in the stream.\n\t// It should return:\n\t// - true if the output contains tool calls and agent should continue processing\n\t// - false if no tool calls and agent should stop\n\t// Note: This field only needs to be configured when using streaming mode\n\t// Note: The handler MUST close the modelOutput stream before returning\n\t// Optional. By default, it checks if the first chunk contains tool calls.\n\t// Note: The default implementation does not work well with Claude, which typically outputs tool calls after text content.\n\t// Note: If your ChatModel doesn't output tool calls first, you can try adding prompts to constrain the model from generating extra text during the tool call.\n\tStreamToolCallChecker func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error)\n\n\t// Summarizer is the summarizer agent that will summarize the outputs of all the chosen specialist agents.\n\t// Only when the Host agent picks multiple Specialist will this be called.\n\t// If you do not provide a summarizer, a default summarizer that simply concatenates all the output messages into one message will be used.\n\t// Note: the default summarizer do not support streaming.\n\tSummarizer *Summarizer\n}\n\nfunc (conf *MultiAgentConfig) validate() error {\n\tif conf == nil {\n\t\treturn errors.New(\"host multi agent config is nil\")\n\t}\n\n\tif conf.Host.ChatModel == nil && conf.Host.ToolCallingModel == nil {\n\t\treturn errors.New(\"host multi agent host ChatModel is nil\")\n\t}\n\n\tif len(conf.Specialists) == 0 {\n\t\treturn errors.New(\"host multi agent specialists are empty\")\n\t}\n\n\tfor _, s := range conf.Specialists {\n\t\tif s.ChatModel == nil && s.Invokable == nil && s.Streamable == nil {\n\t\t\treturn fmt.Errorf(\"specialist %s has no chat model or Invokable or Streamable\", s.Name)\n\t\t}\n\n\t\tif err := s.AgentMeta.validate(); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// AgentMeta is the meta information of an agent within a multi-agent system.\ntype AgentMeta struct {\n\tName        string // the name of the agent, should be unique within multi-agent system\n\tIntendedUse string // the intended use-case of the agent, used as the reason for the multi-agent system to hand over control to this agent\n}\n\nfunc (am AgentMeta) validate() error {\n\tif len(am.Name) == 0 {\n\t\treturn errors.New(\"agent meta name is empty\")\n\t}\n\n\tif len(am.IntendedUse) == 0 {\n\t\treturn errors.New(\"agent meta intended use is empty\")\n\t}\n\n\treturn nil\n}\n\n// Host is the host agent within a multi-agent system.\n// Currently, it can only be a model.ChatModel.\ntype Host struct {\n\tToolCallingModel model.ToolCallingChatModel\n\t// Deprecated: ChatModel is deprecated, please use ToolCallingModel instead.\n\t// This field will be removed in a future release.\n\tChatModel    model.ChatModel\n\tSystemPrompt string\n}\n\n// Specialist is a specialist agent within a host multi-agent system.\n// It can be a model.ChatModel or any Invokable and/or Streamable, such as react.Agent.\n// ChatModel and (Invokable / Streamable) are mutually exclusive, only one should be provided.\n// notice: SystemPrompt only effects when ChatModel has been set.\n// If Invokable is provided but not Streamable, then the Specialist will be 'compose.InvokableLambda'.\n// If Streamable is provided but not Invokable, then the Specialist will be 'compose.StreamableLambda'.\n// if Both Invokable and Streamable is provided, then the Specialist will be 'compose.AnyLambda'.\ntype Specialist struct {\n\tAgentMeta\n\n\tChatModel    model.BaseChatModel\n\tSystemPrompt string\n\n\tInvokable  compose.Invoke[[]*schema.Message, *schema.Message, agent.AgentOption]\n\tStreamable compose.Stream[[]*schema.Message, *schema.Message, agent.AgentOption]\n}\n\n// Summarizer defines a lightweight agent used to summarize\n// conversations or tool outputs using a chat model and prompt.\ntype Summarizer struct {\n\tChatModel    model.BaseChatModel\n\tSystemPrompt string\n}\n\nfunc firstChunkStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) {\n\tdefer sr.Close()\n\n\tfor {\n\t\tmsg, err := sr.Recv()\n\t\tif err == io.EOF {\n\t\t\treturn false, nil\n\t\t}\n\t\tif err != nil {\n\t\t\treturn false, err\n\t\t}\n\n\t\tif len(msg.ToolCalls) > 0 {\n\t\t\treturn true, nil\n\t\t}\n\n\t\tif len(msg.Content) == 0 { // skip empty chunks at the front\n\t\t\tcontinue\n\t\t}\n\n\t\treturn false, nil\n\t}\n}\n"
  },
  {
    "path": "flow/agent/react/callback.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package react provides helpers to build callback handlers for React agents.\npackage react\n\nimport (\n\t\"github.com/cloudwego/eino/callbacks\"\n\ttemplate \"github.com/cloudwego/eino/utils/callbacks\"\n)\n\n// BuildAgentCallback builds a callback handler for agent.\n// e.g.\n//\n//\tcallback := BuildAgentCallback(modelHandler, toolHandler)\n//\tagent, err := react.NewAgent(ctx, &AgentConfig{})\n//\tagent.Generate(ctx, input, agent.WithComposeOptions(compose.WithCallbacks(callback)))\nfunc BuildAgentCallback(modelHandler *template.ModelCallbackHandler, toolHandler *template.ToolCallbackHandler) callbacks.Handler {\n\treturn template.NewHandlerHelper().ChatModel(modelHandler).Tool(toolHandler).Handler()\n}\n"
  },
  {
    "path": "flow/agent/react/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage react\n"
  },
  {
    "path": "flow/agent/react/option.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage react\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/flow/agent\"\n\t\"github.com/cloudwego/eino/internal\"\n\t\"github.com/cloudwego/eino/schema\"\n\tub \"github.com/cloudwego/eino/utils/callbacks\"\n)\n\n// WithToolOptions returns an agent option that specifies tool.Option for the tools in agent.\nfunc WithToolOptions(opts ...tool.Option) agent.AgentOption {\n\treturn agent.WithComposeOptions(compose.WithToolsNodeOption(compose.WithToolOption(opts...)))\n}\n\n// WithChatModelOptions returns an agent option that specifies model.Option for the chat model in agent.\nfunc WithChatModelOptions(opts ...model.Option) agent.AgentOption {\n\treturn agent.WithComposeOptions(compose.WithChatModelOption(opts...))\n}\n\n// WithToolList returns an agent option that specifies compose.ToolsNodeOption for ToolsNode in agent.\n// If you also need to pass ToolInfo to the chat model, use WithTools instead.\n// Deprecated: This changes tool list for ToolsNode ONLY.\nfunc WithToolList(tools ...tool.BaseTool) agent.AgentOption {\n\treturn agent.WithComposeOptions(compose.WithToolsNodeOption(compose.WithToolList(tools...)))\n}\n\n// WithTools is a convenience function that configures a React agent with a list of tools.\n// It performs two essential operations:\n//  1. Extracts tool information for the chat model to understand available tools\n//  2. Registers the actual tool implementations for execution\n//\n// Parameters:\n//   - ctx: The context for the operation, used when calling Info() on each tool\n//   - tools: A variadic list of tools that must implement either InvokableTool or StreamableTool interfaces\n//\n// Returns:\n//   - []agent.AgentOption: A slice containing exactly 2 agent options:\n//   - Option 1: Configures the chat model with tool schemas via model.WithTools(toolInfos)\n//   - Option 2: Registers the tool implementations via compose.WithToolList(tools...)\n//   - error: Returns an error if any tool's Info() method fails\n//\n// Usage Example:\n//\n//\tctx := context.Background()\n//\tagentOptions, err := WithTools(ctx, myTool1, myTool2, myTool3)\n//\tif err != nil {\n//\t    return fmt.Errorf(\"failed to configure tools: %w\", err)\n//\t}\n//\n//\tagent, err := react.NewAgent(ctx, &react.AgentConfig{\n//\t    ToolCallingModel: myModel,\n//\t    // other config...\n//\t})\n//\tif err != nil {\n//\t    return fmt.Errorf(\"failed to create agent: %w\", err)\n//\t}\n//\n//\t// Use the tool options with Generate or Stream methods\n//\tmsg, err := agent.Generate(ctx, messages, agentOptions...)\n//\t// or\n//\tstream, err := agent.Stream(ctx, messages, agentOptions...)\n//\n// Comparison with Related Functions:\n//   - WithToolList: Only registers tool implementations, doesn't configure the chat model\n//   - WithTools: Comprehensive setup that handles both chat model configuration and tool registration\n//\n// Notes:\n//   - The function always returns exactly 2 options when successful\n//   - Both returned options should be applied to the agent for proper tool functionality\nfunc WithTools(ctx context.Context, tools ...tool.BaseTool) ([]agent.AgentOption, error) {\n\ttoolInfos := make([]*schema.ToolInfo, 0, len(tools))\n\tfor _, tl := range tools {\n\t\tinfo, err := tl.Info(ctx)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\ttoolInfos = append(toolInfos, info)\n\t}\n\n\topts := make([]agent.AgentOption, 2)\n\topts[0] = agent.WithComposeOptions(compose.WithChatModelOption(model.WithTools(toolInfos)))\n\topts[1] = agent.WithComposeOptions(compose.WithToolsNodeOption(compose.WithToolList(tools...)))\n\treturn opts, nil\n}\n\n// Iterator provides a lightweight FIFO stream of values and errors\n// produced during agent execution.\ntype Iterator[T any] struct {\n\tch *internal.UnboundedChan[item[T]]\n}\n\n// Next retrieves the next value from the iterator.\n// It returns the zero value and false when the stream is exhausted.\nfunc (iter *Iterator[T]) Next() (T, bool, error) {\n\tch := iter.ch\n\tif ch == nil {\n\t\tvar zero T\n\t\treturn zero, false, nil\n\t}\n\n\ti, ok := ch.Receive()\n\tif !ok {\n\t\tvar zero T\n\t\treturn zero, false, nil\n\t}\n\n\treturn i.v, true, i.err\n}\n\n// MessageFuture exposes asynchronous accessors for messages produced\n// by Generate and Stream calls.\ntype MessageFuture interface {\n\t// GetMessages returns an iterator for retrieving messages generated during \"agent.Generate\" calls.\n\tGetMessages() *Iterator[*schema.Message]\n\n\t// GetMessageStreams returns an iterator for retrieving streaming messages generated during \"agent.Stream\" calls.\n\tGetMessageStreams() *Iterator[*schema.StreamReader[*schema.Message]]\n}\n\n// WithMessageFuture returns an agent option and a MessageFuture interface instance.\n// The option configures the agent to collect messages generated during execution,\n// while the MessageFuture interface allows users to asynchronously retrieve these messages.\nfunc WithMessageFuture() (agent.AgentOption, MessageFuture) {\n\th := &cbHandler{started: make(chan struct{})}\n\n\tcmHandler := &ub.ModelCallbackHandler{\n\t\tOnEnd:                 h.onChatModelEnd,\n\t\tOnEndWithStreamOutput: h.onChatModelEndWithStreamOutput,\n\t}\n\tcreateToolResultSender := func() toolResultSender {\n\t\treturn func(toolName, callID, result string) {\n\t\t\tmsg := schema.ToolMessage(result, callID, schema.WithToolName(toolName))\n\t\t\th.sendMessage(msg)\n\t\t}\n\t}\n\tcreateStreamToolResultSender := func() streamToolResultSender {\n\t\treturn func(toolName, callID string, resultStream *schema.StreamReader[string]) {\n\t\t\tcvt := func(in string) (*schema.Message, error) {\n\t\t\t\treturn schema.ToolMessage(in, callID, schema.WithToolName(toolName)), nil\n\t\t\t}\n\t\t\tmsgStream := schema.StreamReaderWithConvert(resultStream, cvt)\n\t\t\th.sendMessageStream(msgStream)\n\t\t}\n\t}\n\n\tcreateEnhancedToolResultSender := func() enhancedToolResultSender {\n\t\treturn func(toolName, callID string, result *schema.ToolResult) {\n\t\t\tvar err error\n\t\t\tmsg := schema.ToolMessage(\"\", callID, schema.WithToolName(toolName))\n\t\t\tmsg.UserInputMultiContent, err = result.ToMessageInputParts()\n\t\t\tif err != nil {\n\t\t\t\treturn\n\t\t\t}\n\t\t\th.sendMessage(msg)\n\t\t}\n\t}\n\n\tcreateEnhancedStreamToolResultSender := func() enhancedStreamToolResultSender {\n\t\treturn func(toolName, callID string, resultStream *schema.StreamReader[*schema.ToolResult]) {\n\t\t\tcvt := func(result *schema.ToolResult) (*schema.Message, error) {\n\t\t\t\tvar err error\n\t\t\t\tmsg := schema.ToolMessage(\"\", callID, schema.WithToolName(toolName))\n\t\t\t\tmsg.UserInputMultiContent, err = result.ToMessageInputParts()\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\treturn msg, nil\n\t\t\t}\n\t\t\tmsgStream := schema.StreamReaderWithConvert(resultStream, cvt)\n\t\t\th.sendMessageStream(msgStream)\n\t\t}\n\t}\n\n\tgraphHandler := callbacks.NewHandlerBuilder().\n\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\th.onGraphStart(ctx, info, input)\n\t\t\treturn setToolResultSendersToCtx(ctx, &toolResultSenders{\n\t\t\t\tsender:                         createToolResultSender(),\n\t\t\t\tstreamSender:                   createStreamToolResultSender(),\n\t\t\t\tenhancedResultSender:           createEnhancedToolResultSender(),\n\t\t\t\tenhancedStreamToolResultSender: createEnhancedStreamToolResultSender(),\n\t\t\t})\n\t\t}).\n\t\tOnStartWithStreamInputFn(func(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {\n\t\t\th.onGraphStartWithStreamInput(ctx, info, input)\n\t\t\treturn setToolResultSendersToCtx(ctx, &toolResultSenders{\n\t\t\t\tsender:                         createToolResultSender(),\n\t\t\t\tstreamSender:                   createStreamToolResultSender(),\n\t\t\t\tenhancedResultSender:           createEnhancedToolResultSender(),\n\t\t\t\tenhancedStreamToolResultSender: createEnhancedStreamToolResultSender(),\n\t\t\t})\n\n\t\t}).\n\t\tOnEndFn(h.onGraphEnd).\n\t\tOnEndWithStreamOutputFn(h.onGraphEndWithStreamOutput).\n\t\tOnErrorFn(h.onGraphError).Build()\n\tcb := ub.NewHandlerHelper().ChatModel(cmHandler).Graph(graphHandler).Handler()\n\n\toption := agent.WithComposeOptions(compose.WithCallbacks(cb))\n\n\treturn option, h\n}\n\ntype item[T any] struct {\n\tv   T\n\terr error\n}\n\ntype cbHandler struct {\n\tmsgs  *internal.UnboundedChan[item[*schema.Message]]\n\tsMsgs *internal.UnboundedChan[item[*schema.StreamReader[*schema.Message]]]\n\n\tstarted chan struct{}\n}\n\nfunc (h *cbHandler) GetMessages() *Iterator[*schema.Message] {\n\t<-h.started\n\n\treturn &Iterator[*schema.Message]{ch: h.msgs}\n}\n\nfunc (h *cbHandler) GetMessageStreams() *Iterator[*schema.StreamReader[*schema.Message]] {\n\t<-h.started\n\n\treturn &Iterator[*schema.StreamReader[*schema.Message]]{ch: h.sMsgs}\n}\n\nfunc (h *cbHandler) onChatModelEnd(ctx context.Context,\n\t_ *callbacks.RunInfo, input *model.CallbackOutput) context.Context {\n\n\th.sendMessage(input.Message)\n\n\treturn ctx\n}\n\nfunc (h *cbHandler) onChatModelEndWithStreamOutput(ctx context.Context,\n\t_ *callbacks.RunInfo, input *schema.StreamReader[*model.CallbackOutput]) context.Context {\n\n\tc := func(output *model.CallbackOutput) (*schema.Message, error) {\n\t\treturn output.Message, nil\n\t}\n\ts := schema.StreamReaderWithConvert(input, c)\n\n\th.sendMessageStream(s)\n\n\treturn ctx\n}\n\nfunc (h *cbHandler) onGraphError(ctx context.Context,\n\t_ *callbacks.RunInfo, err error) context.Context {\n\n\tif h.msgs != nil {\n\t\th.msgs.Send(item[*schema.Message]{err: err})\n\t} else {\n\t\th.sMsgs.Send(item[*schema.StreamReader[*schema.Message]]{err: err})\n\t}\n\n\treturn ctx\n}\n\nfunc (h *cbHandler) onGraphEnd(ctx context.Context,\n\t_ *callbacks.RunInfo, _ callbacks.CallbackOutput) context.Context {\n\n\th.msgs.Close()\n\n\treturn ctx\n}\n\nfunc (h *cbHandler) onGraphEndWithStreamOutput(ctx context.Context,\n\t_ *callbacks.RunInfo, _ *schema.StreamReader[callbacks.CallbackOutput]) context.Context {\n\n\th.sMsgs.Close()\n\n\treturn ctx\n}\n\nfunc (h *cbHandler) onGraphStart(ctx context.Context,\n\t_ *callbacks.RunInfo, _ callbacks.CallbackInput) context.Context {\n\n\th.msgs = internal.NewUnboundedChan[item[*schema.Message]]()\n\n\tclose(h.started)\n\n\treturn ctx\n}\n\nfunc (h *cbHandler) onGraphStartWithStreamInput(ctx context.Context, _ *callbacks.RunInfo,\n\tinput *schema.StreamReader[callbacks.CallbackInput]) context.Context {\n\tinput.Close()\n\n\th.sMsgs = internal.NewUnboundedChan[item[*schema.StreamReader[*schema.Message]]]()\n\n\tclose(h.started)\n\n\treturn ctx\n}\n\nfunc (h *cbHandler) sendMessage(msg *schema.Message) {\n\tif h.msgs != nil {\n\t\th.msgs.Send(item[*schema.Message]{v: msg})\n\t} else {\n\t\tsMsg := schema.StreamReaderFromArray([]*schema.Message{msg})\n\t\th.sMsgs.Send(item[*schema.StreamReader[*schema.Message]]{v: sMsg})\n\t}\n}\n\nfunc (h *cbHandler) sendMessageStream(sMsg *schema.StreamReader[*schema.Message]) {\n\tif h.sMsgs != nil {\n\t\th.sMsgs.Send(item[*schema.StreamReader[*schema.Message]]{v: sMsg})\n\t} else {\n\t\t// concat\n\t\tmsg, err := schema.ConcatMessageStream(sMsg)\n\n\t\tif err != nil {\n\t\t\th.msgs.Send(item[*schema.Message]{err: err})\n\t\t} else {\n\t\t\th.msgs.Send(item[*schema.Message]{v: msg})\n\t\t}\n\t}\n}\n"
  },
  {
    "path": "flow/agent/react/option_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage react\n\nimport (\n\t\"context\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestWithMessageFuture(t *testing.T) {\n\tctx := context.Background()\n\n\t// Test with tool calls\n\tt.Run(\"test generate with tool calls\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\t\tfakeTool := &fakeToolGreetForTest{}\n\n\t\tinfo, err := fakeTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Mock model response with tool call\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"\",\n\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"name\": \"test user\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"final response\", nil), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create agent with MessageFuture\n\t\toption, future := WithMessageFuture()\n\t\ta, err := NewAgent(ctx, &AgentConfig{\n\t\t\tToolCallingModel: cm,\n\t\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t\tMaxStep: 3,\n\t\t})\n\t\tassert.Nil(t, err)\n\n\t\t// Generate response\n\t\tresponse, err := a.Generate(ctx, []*schema.Message{\n\t\t\tschema.UserMessage(\"use the greet tool\"),\n\t\t}, option)\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, \"final response\", response.Content)\n\n\t\tsIter := future.GetMessageStreams()\n\t\t// Should be no messages\n\t\t_, hasNext, err := sIter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.False(t, hasNext)\n\n\t\titer := future.GetMessages()\n\t\t// First message should be the assistant message for tool calling\n\t\tmsg1, hasNext, err := iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.Equal(t, schema.Assistant, msg1.Role)\n\t\tassert.Equal(t, 1, len(msg1.ToolCalls))\n\n\t\t// Second message should be the tool response\n\t\tmsg2, hasNext, err := iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.Equal(t, schema.Tool, msg2.Role)\n\n\t\t// Third message should be the final response\n\t\tmsg3, hasNext, err := iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.Equal(t, \"final response\", msg3.Content)\n\n\t\t// Should be no more messages\n\t\t_, hasNext, err = iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.False(t, hasNext)\n\t})\n\t// Test with streaming tool calls\n\tt.Run(\"test generate with streaming tool calls\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\t\tfakeTool := &fakeStreamToolGreetForTest{}\n\n\t\tinfo, err := fakeTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Mock model response with tool call\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"\",\n\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"name\": \"test user\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"final response\", nil), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create agent with MessageFuture\n\t\toption, future := WithMessageFuture()\n\t\ta, err := NewAgent(ctx, &AgentConfig{\n\t\t\tToolCallingModel: cm,\n\t\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t\tMaxStep: 3,\n\t\t})\n\t\tassert.Nil(t, err)\n\n\t\t// Generate response\n\t\tresponse, err := a.Generate(ctx, []*schema.Message{\n\t\t\tschema.UserMessage(\"use the greet tool\"),\n\t\t}, option)\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, \"final response\", response.Content)\n\n\t\t// Get messages from future\n\t\titer := future.GetMessages()\n\n\t\t// First message should be the assistant message for tool calling\n\t\tmsg1, hasNext, err := iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.Equal(t, schema.Assistant, msg1.Role)\n\t\tassert.Equal(t, 1, len(msg1.ToolCalls))\n\n\t\t// Second message should be the tool response\n\t\tmsg2, hasNext, err := iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.Equal(t, schema.Tool, msg2.Role)\n\n\t\t// Third message should be the final response\n\t\tmsg3, hasNext, err := iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.Equal(t, \"final response\", msg3.Content)\n\n\t\t// Should be no more messages\n\t\t_, hasNext, err = iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.False(t, hasNext)\n\t})\n\n\t// Test with non-streaming tool but using agent's Stream interface\n\tt.Run(\"test stream with tool calls\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\t\tfakeTool := &fakeToolGreetForTest{}\n\n\t\tinfo, err := fakeTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Mock model response with tool call\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(\"\",\n\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"name\": \"test user\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t})}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(\"final response\", nil)}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create agent with MessageFuture\n\t\toption, future := WithMessageFuture()\n\t\ta, err := NewAgent(ctx, &AgentConfig{\n\t\t\tToolCallingModel: cm,\n\t\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t\tMaxStep: 3,\n\t\t})\n\t\tassert.Nil(t, err)\n\n\t\t// Use Stream interface\n\t\tstream, err := a.Stream(ctx, []*schema.Message{\n\t\t\tschema.UserMessage(\"use the greet tool\"),\n\t\t}, option)\n\t\tassert.Nil(t, err)\n\n\t\t// Collect all chunks from stream\n\t\tfinalResponse, err := schema.ConcatMessageStream(stream)\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, \"final response\", finalResponse.Content)\n\n\t\titer := future.GetMessages()\n\t\t// Should be no messages\n\t\t_, hasNext, err := iter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.False(t, hasNext)\n\n\t\t// Get message streams from future\n\t\tsIter := future.GetMessageStreams()\n\n\t\t// First message should be the assistant message for tool calling\n\t\tstream1, hasNext, err := sIter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.NotNil(t, stream1)\n\t\tmsg1, err := schema.ConcatMessageStream(stream1)\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, schema.Assistant, msg1.Role)\n\t\tassert.Equal(t, 1, len(msg1.ToolCalls))\n\n\t\t// Second message should be the tool response\n\t\tstream2, hasNext, err := sIter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.NotNil(t, stream2)\n\t\tmsg2, err := schema.ConcatMessageStream(stream2)\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, schema.Tool, msg2.Role)\n\n\t\t// Third message should be the final response\n\t\tstream3, hasNext, err := sIter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.True(t, hasNext)\n\t\tassert.NotNil(t, stream3)\n\t\tmsg3, err := schema.ConcatMessageStream(stream3)\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, \"final response\", msg3.Content)\n\n\t\t// Should be no more messages\n\t\t_, hasNext, err = sIter.Next()\n\t\tassert.Nil(t, err)\n\t\tassert.False(t, hasNext)\n\t})\n\n\tt.Run(\"test stream with streaming tool calls and with concurrent goroutines\", func(t *testing.T) {\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\t\tfakeTool := &fakeStreamToolGreetForTest{}\n\n\t\tinfo, err := fakeTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\t// Mock model response with tool call\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(\"\",\n\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"name\": \"test user\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t})}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(\"final response\", nil)}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\t// Create agent with MessageFuture\n\t\toption, future := WithMessageFuture()\n\t\ta, err := NewAgent(ctx, &AgentConfig{\n\t\t\tToolCallingModel: cm,\n\t\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t\t},\n\t\t\tMaxStep: 3,\n\t\t})\n\t\tassert.Nil(t, err)\n\n\t\tvar wg sync.WaitGroup\n\t\twg.Add(1)\n\t\tgo func() {\n\t\t\tdefer wg.Done()\n\n\t\t\t// Get message streams from future\n\t\t\tsIter := future.GetMessageStreams()\n\n\t\t\t// First message should be the assistant message for tool calling\n\t\t\tstream1, hasNext, err_ := sIter.Next()\n\t\t\tassert.Nil(t, err_)\n\t\t\tassert.True(t, hasNext)\n\t\t\tassert.NotNil(t, stream1)\n\t\t\tmsg1, err_ := schema.ConcatMessageStream(stream1)\n\t\t\tassert.Nil(t, err_)\n\t\t\tassert.Equal(t, schema.Assistant, msg1.Role)\n\t\t\tassert.Equal(t, 1, len(msg1.ToolCalls))\n\n\t\t\t// Second message should be the tool response\n\t\t\tstream2, hasNext, err_ := sIter.Next()\n\t\t\tassert.Nil(t, err_)\n\t\t\tassert.True(t, hasNext)\n\t\t\tassert.NotNil(t, stream2)\n\t\t\tmsg2, err_ := schema.ConcatMessageStream(stream2)\n\t\t\tassert.Nil(t, err_)\n\t\t\tassert.Equal(t, schema.Tool, msg2.Role)\n\n\t\t\t// Third message should be the final response\n\t\t\tstream3, hasNext, err_ := sIter.Next()\n\t\t\tassert.Nil(t, err_)\n\t\t\tassert.True(t, hasNext)\n\t\t\tassert.NotNil(t, stream3)\n\t\t\tmsg3, err_ := schema.ConcatMessageStream(stream3)\n\t\t\tassert.Nil(t, err_)\n\t\t\tassert.Equal(t, \"final response\", msg3.Content)\n\n\t\t\t// Should be no more messages\n\t\t\t_, hasNext, err_ = sIter.Next()\n\t\t\tassert.Nil(t, err_)\n\t\t\tassert.False(t, hasNext)\n\t\t}()\n\n\t\t// Use Stream interface\n\t\tstream, err := a.Stream(ctx, []*schema.Message{\n\t\t\tschema.UserMessage(\"use the greet tool\"),\n\t\t}, option)\n\t\tassert.Nil(t, err)\n\n\t\t// Collect all chunks from stream\n\t\tfinalResponse, err := schema.ConcatMessageStream(stream)\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, \"final response\", finalResponse.Content)\n\n\t\twg.Wait()\n\t})\n}\n\nfunc TestWithToolOptions(t *testing.T) {\n\ttype dummyOpt struct{ val string }\n\topt := tool.WrapImplSpecificOptFn(func(o *dummyOpt) { o.val = \"mock\" })\n\tagentOpt := WithToolOptions(opt)\n\tassert.NotNil(t, agentOpt)\n\t// The returned value should be an agent.AgentOption (function)\n\tassert.IsType(t, agentOpt, agentOpt)\n}\n\nfunc TestWithChatModelOptions(t *testing.T) {\n\topt := model.WithModel(\"mock-model\")\n\tagentOpt := WithChatModelOptions(opt)\n\tassert.NotNil(t, agentOpt)\n\tassert.IsType(t, agentOpt, agentOpt)\n}\n\n// dummyBaseTool is a minimal implementation of tool.BaseTool for testing.\ntype dummyBaseTool struct{}\n\nfunc (d *dummyBaseTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"dummy\"}, nil\n}\n\nfunc (d *dummyBaseTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn \"dummy-response\", nil\n}\n\ntype assertTool struct {\n\ttoolOptVal      string\n\treceivedToolOpt bool\n}\ntype toolOpt struct{ val string }\n\nfunc (a *assertTool) Info(ctx context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{Name: \"assert_tool\"}, nil\n}\nfunc (a *assertTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {\n\topt := tool.GetImplSpecificOptions(&toolOpt{}, opts...)\n\tif opt.val == a.toolOptVal {\n\t\ta.receivedToolOpt = true\n\t}\n\treturn \"tool-response\", nil\n}\n\nfunc TestAgentWithAllOptions(t *testing.T) {\n\tctx := context.Background()\n\tctrl := gomock.NewController(t)\n\n\t// Prepare a tool that asserts it receives the tool option\n\ttoolOptVal := \"tool-opt-value\"\n\tto := tool.WrapImplSpecificOptFn(func(o *toolOpt) { o.val = toolOptVal })\n\tat := &assertTool{toolOptVal: toolOptVal}\n\n\t// Prepare a mock chat model that asserts it receives the model option\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\tmodelOpt := model.WithModel(\"test-model\")\n\tmodelOptReceived := false\n\ttimes := 0\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(\n\t\tfunc(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\ttimes++\n\t\t\tif times == 1 {\n\t\t\t\tfor _, o := range opts {\n\t\t\t\t\topt := model.GetCommonOptions(&model.Options{}, o)\n\t\t\t\t\tif opt.Model != nil && *opt.Model == \"test-model\" {\n\t\t\t\t\t\tmodelOptReceived = true\n\t\t\t\t\t}\n\t\t\t\t}\n\n\t\t\t\tinfo, _ := at.Info(ctx)\n\t\t\t\treturn schema.AssistantMessage(\"hello max\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\tArguments: \"\",\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\tnil\n\t\t\t}\n\n\t\t\treturn schema.AssistantMessage(\"ok\", nil), nil\n\t\t},\n\t).AnyTimes()\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tagentOpt := WithToolOptions(to)\n\tagentOpt2 := WithChatModelOptions(modelOpt)\n\tagentOpt3, err := WithTools(context.Background(), at)\n\tassert.NoError(t, err)\n\n\ta, err := NewAgent(ctx, &AgentConfig{\n\t\tToolCallingModel: cm,\n\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{&dummyBaseTool{}},\n\t\t},\n\t\tMaxStep: 20,\n\t})\n\tassert.NoError(t, err)\n\n\t_, err = a.Generate(ctx, []*schema.Message{\n\t\tschema.UserMessage(\"call the tool\"),\n\t}, agentOpt, agentOpt2, agentOpt3[0], agentOpt3[1])\n\tassert.NoError(t, err)\n\tassert.True(t, modelOptReceived, \"model option should be received by chat model\")\n\tassert.True(t, at.receivedToolOpt, \"tool option should be received by tool\")\n}\n\ntype simpleToolForMiddlewareTest struct {\n\tname   string\n\tresult string\n}\n\nfunc (s *simpleToolForMiddlewareTest) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: s.name,\n\t\tDesc: \"simple tool for middleware test\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"input\": {\n\t\t\t\t\tDesc:     \"input\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t\tType:     schema.String,\n\t\t\t\t},\n\t\t\t}),\n\t}, nil\n}\n\nfunc (s *simpleToolForMiddlewareTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn s.result, nil\n}\n\nfunc (s *simpleToolForMiddlewareTest) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) {\n\treturn schema.StreamReaderFromArray([]string{s.result}), nil\n}\n\nfunc TestMessageFuture_ToolResultMiddleware_EmitsFinalResult(t *testing.T) {\n\toriginalResult := \"original_result\"\n\tmodifiedResult := \"modified_by_middleware\"\n\n\tresultModifyingMiddleware := compose.ToolMiddleware{\n\t\tInvokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\t\t\toutput, err := next(ctx, input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\toutput.Result = modifiedResult\n\t\t\t\treturn output, nil\n\t\t\t}\n\t\t},\n\t\tStreamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\t\t\toutput, err := next(ctx, input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\toutput.Result = schema.StreamReaderFromArray([]string{modifiedResult})\n\t\t\t\treturn output, nil\n\t\t\t}\n\t\t},\n\t}\n\n\tt.Run(\"Invoke\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\ttestTool := &simpleToolForMiddlewareTest{name: \"test_tool\", result: originalResult}\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tinfo, err := testTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"\",\n\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"input\": \"test\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.AssistantMessage(\"final response\", nil), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\toption, future := WithMessageFuture()\n\t\ta, err := NewAgent(ctx, &AgentConfig{\n\t\t\tToolCallingModel: cm,\n\t\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools:               []tool.BaseTool{testTool},\n\t\t\t\tToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware},\n\t\t\t},\n\t\t\tMaxStep: 3,\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tresponse, err := a.Generate(ctx, []*schema.Message{\n\t\t\tschema.UserMessage(\"call the tool\"),\n\t\t}, option)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"final response\", response.Content)\n\n\t\titer := future.GetMessages()\n\n\t\tvar allMsgs []*schema.Message\n\t\tfor {\n\t\t\tmsg, hasNext, err := iter.Next()\n\t\t\tif err != nil || !hasNext {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tallMsgs = append(allMsgs, msg)\n\t\t}\n\n\t\tassert.GreaterOrEqual(t, len(allMsgs), 3, \"should have at least 3 messages\")\n\t\tif len(allMsgs) >= 3 {\n\t\t\tassert.Equal(t, schema.Assistant, allMsgs[0].Role)\n\t\t\tassert.Equal(t, 1, len(allMsgs[0].ToolCalls))\n\n\t\t\tassert.Equal(t, schema.Tool, allMsgs[1].Role)\n\t\t\tassert.Equal(t, modifiedResult, allMsgs[1].Content,\n\t\t\t\t\"MessageFuture should receive the middleware-modified tool result\")\n\t\t\tassert.NotEqual(t, originalResult, allMsgs[1].Content,\n\t\t\t\t\"MessageFuture should NOT receive the original tool result\")\n\n\t\t\tassert.Equal(t, \"final response\", allMsgs[2].Content)\n\t\t}\n\t})\n\n\tt.Run(\"Stream\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\ttestTool := &simpleToolForMiddlewareTest{name: \"test_tool_stream\", result: originalResult}\n\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t\tinfo, err := testTool.Info(ctx)\n\t\tassert.NoError(t, err)\n\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"\", []schema.ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tID: \"tool-call-1\",\n\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\tArguments: `{\"input\": \"test\"}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t}),\n\t\t\t}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tReturn(schema.StreamReaderFromArray([]*schema.Message{\n\t\t\t\tschema.AssistantMessage(\"final response\", nil),\n\t\t\t}), nil).\n\t\t\tTimes(1)\n\t\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\t\toption, future := WithMessageFuture()\n\t\ta, err := NewAgent(ctx, &AgentConfig{\n\t\t\tToolCallingModel: cm,\n\t\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools:               []tool.BaseTool{testTool},\n\t\t\t\tToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware},\n\t\t\t},\n\t\t\tMaxStep: 3,\n\t\t})\n\t\tassert.NoError(t, err)\n\n\t\tresponse, err := a.Stream(ctx, []*schema.Message{\n\t\t\tschema.UserMessage(\"call the tool\"),\n\t\t}, option)\n\t\tassert.NoError(t, err)\n\n\t\tvar msgs []*schema.Message\n\t\tfor {\n\t\t\tmsg, err := response.Recv()\n\t\t\tif err != nil {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tmsgs = append(msgs, msg)\n\t\t}\n\t\tfinalMsg, err := schema.ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, \"final response\", finalMsg.Content)\n\n\t\titer := future.GetMessageStreams()\n\n\t\tvar allMsgs []*schema.Message\n\t\tfor {\n\t\t\tmsgStream, hasNext, err := iter.Next()\n\t\t\tif err != nil || !hasNext {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tvar streamMsgs []*schema.Message\n\t\t\tfor {\n\t\t\t\tmsg, err := msgStream.Recv()\n\t\t\t\tif err != nil {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tstreamMsgs = append(streamMsgs, msg)\n\t\t\t}\n\t\t\tif len(streamMsgs) > 0 {\n\t\t\t\tconcated, err := schema.ConcatMessages(streamMsgs)\n\t\t\t\tif err == nil {\n\t\t\t\t\tallMsgs = append(allMsgs, concated)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tassert.GreaterOrEqual(t, len(allMsgs), 3, \"should have at least 3 messages\")\n\t\tif len(allMsgs) >= 3 {\n\t\t\tassert.Equal(t, schema.Assistant, allMsgs[0].Role)\n\t\t\tassert.Equal(t, 1, len(allMsgs[0].ToolCalls))\n\n\t\t\tassert.Equal(t, schema.Tool, allMsgs[1].Role)\n\t\t\tassert.Equal(t, modifiedResult, allMsgs[1].Content,\n\t\t\t\t\"MessageFuture should receive the middleware-modified tool result\")\n\t\t\tassert.NotEqual(t, originalResult, allMsgs[1].Content,\n\t\t\t\t\"MessageFuture should NOT receive the original tool result\")\n\n\t\t\tassert.Equal(t, \"final response\", allMsgs[2].Content)\n\t\t}\n\t})\n}\n"
  },
  {
    "path": "flow/agent/react/react.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage react\n\nimport (\n\t\"context\"\n\t\"io\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/flow/agent\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype toolResultSender func(toolName, callID, result string)\n\ntype enhancedToolResultSender func(toolName, callID string, result *schema.ToolResult)\ntype streamToolResultSender func(toolName, callID string, resultStream *schema.StreamReader[string])\ntype enhancedStreamToolResultSender func(toolName, callID string, resultStream *schema.StreamReader[*schema.ToolResult])\ntype toolResultSenders struct {\n\tsender       toolResultSender\n\tstreamSender streamToolResultSender\n\n\tenhancedResultSender           enhancedToolResultSender\n\tenhancedStreamToolResultSender enhancedStreamToolResultSender\n}\n\ntype toolResultSenderCtxKey struct{}\n\nfunc setToolResultSendersToCtx(ctx context.Context, senders *toolResultSenders) context.Context {\n\treturn context.WithValue(ctx, toolResultSenderCtxKey{}, senders)\n}\n\nfunc getToolResultSendersFromCtx(ctx context.Context) *toolResultSenders {\n\tv := ctx.Value(toolResultSenderCtxKey{})\n\tif v == nil {\n\t\treturn nil\n\t}\n\treturn v.(*toolResultSenders)\n}\n\ntype state struct {\n\tMessages                 []*schema.Message\n\tReturnDirectlyToolCallID string\n}\n\nfunc init() {\n\tschema.RegisterName[*state](\"_eino_react_state\")\n}\n\nfunc newToolResultCollectorMiddleware() compose.ToolMiddleware {\n\treturn compose.ToolMiddleware{\n\t\tInvokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {\n\t\t\t\tsenders := getToolResultSendersFromCtx(ctx)\n\t\t\t\toutput, err := next(ctx, input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tif senders != nil && senders.sender != nil {\n\t\t\t\t\tsenders.sender(input.Name, input.CallID, output.Result)\n\t\t\t\t}\n\t\t\t\treturn output, nil\n\t\t\t}\n\t\t},\n\t\tStreamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {\n\t\t\t\tsenders := getToolResultSendersFromCtx(ctx)\n\t\t\t\toutput, err := next(ctx, input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tif senders != nil && senders.streamSender != nil {\n\t\t\t\t\tstreams := output.Result.Copy(2)\n\t\t\t\t\tsenders.streamSender(input.Name, input.CallID, streams[0])\n\t\t\t\t\toutput.Result = streams[1]\n\t\t\t\t}\n\t\t\t\treturn output, nil\n\t\t\t}\n\t\t},\n\t\tEnhancedInvokable: func(next compose.EnhancedInvokableToolEndpoint) compose.EnhancedInvokableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {\n\t\t\t\tsenders := getToolResultSendersFromCtx(ctx)\n\t\t\t\toutput, err := next(ctx, input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tif senders != nil && senders.enhancedResultSender != nil {\n\t\t\t\t\tsenders.enhancedResultSender(input.Name, input.CallID, output.Result)\n\t\t\t\t}\n\t\t\t\treturn output, nil\n\n\t\t\t}\n\t\t},\n\t\tEnhancedStreamable: func(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint {\n\t\t\treturn func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {\n\t\t\t\tsenders := getToolResultSendersFromCtx(ctx)\n\t\t\t\toutput, err := next(ctx, input)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tif senders != nil && senders.enhancedStreamToolResultSender != nil {\n\t\t\t\t\tstreams := output.Result.Copy(2)\n\t\t\t\t\tsenders.enhancedStreamToolResultSender(input.Name, input.CallID, streams[0])\n\t\t\t\t\toutput.Result = streams[1]\n\t\t\t\t}\n\t\t\t\treturn output, nil\n\t\t\t}\n\t\t},\n\t}\n}\n\nconst (\n\tnodeKeyTools = \"tools\"\n\tnodeKeyModel = \"chat\"\n)\n\n// MessageModifier modify the input messages before the model is called.\ntype MessageModifier func(ctx context.Context, input []*schema.Message) []*schema.Message\n\n// AgentConfig is the config for ReAct agent.\ntype AgentConfig struct {\n\t// ToolCallingModel is the chat model to be used for handling user messages with tool calling capability.\n\t// This is the recommended model field to use.\n\tToolCallingModel model.ToolCallingChatModel\n\n\t// Deprecated: Use ToolCallingModel instead.\n\tModel model.ChatModel\n\n\t// ToolsConfig is the config for tools node.\n\tToolsConfig compose.ToolsNodeConfig\n\n\t// MessageModifier.\n\t// modify the input messages before the model is called, it's useful when you want to add some system prompt or other messages.\n\tMessageModifier MessageModifier\n\n\t// MessageRewriter modifies message in the state, before the ChatModel is called.\n\t// It takes the messages stored accumulated in state, modify them, and put the modified version back into state.\n\t// Useful for compressing message history to fit the model context window,\n\t// or if you want to make changes to messages that take effect across multiple model calls.\n\t// NOTE: if both MessageModifier and MessageRewriter are set, MessageRewriter will be called before MessageModifier.\n\tMessageRewriter MessageModifier\n\n\t// MaxStep.\n\t// default 12 of steps in pregel (node num + 10).\n\tMaxStep int `json:\"max_step\"`\n\n\t// Tools that will make agent return directly when the tool is called.\n\t// When multiple tools are called and more than one tool is in the return directly list, only the first one will be returned.\n\tToolReturnDirectly map[string]struct{}\n\n\t// StreamToolCallChecker is a function to determine whether the model's streaming output contains tool calls.\n\t// Different models have different ways of outputting tool calls in streaming mode:\n\t// - Some models (like OpenAI) output tool calls directly\n\t// - Others (like Claude) output text first, then tool calls\n\t// This handler allows custom logic to check for tool calls in the stream.\n\t// It should return:\n\t// - true if the output contains tool calls and agent should continue processing\n\t// - false if no tool calls and agent should stop\n\t// Note: This field only needs to be configured when using streaming mode\n\t// Note: The handler MUST close the modelOutput stream before returning\n\t// Optional. By default, it checks if the first chunk contains tool calls.\n\t// Note: The default implementation does not work well with Claude, which typically outputs tool calls after text content.\n\t// Note: If your ChatModel doesn't output tool calls first, you can try adding prompts to constrain the model from generating extra text during the tool call.\n\tStreamToolCallChecker func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error)\n\n\t// GraphName is the graph name of the ReAct Agent.\n\t// Optional. Default `ReActAgent`.\n\tGraphName string\n\t// ModelNodeName is the node name of the model node in the ReAct Agent graph.\n\t// Optional. Default `ChatModel`.\n\tModelNodeName string\n\t// ToolsNodeName is the node name of the tools node in the ReAct Agent graph.\n\t// Optional. Default `Tools`.\n\tToolsNodeName string\n}\n\n// NewPersonaModifier returns a MessageModifier that adds a persona message to the input.\n// example:\n//\n//\tpersona := \"You are an expert in golang.\"\n//\tconfig := AgentConfig{\n//\t\tModel: model,\n//\t\tMessageModifier: NewPersonaModifier(persona),\n//\t}\n//\tagent, err := NewAgent(ctx, config)\n//\tif err != nil {return}\n//\tmsg, err := agent.Generate(ctx, []*schema.Message{{Role: schema.User, Content: \"how to build agent with eino\"}})\n//\tif err != nil {return}\n//\tprintln(msg.Content)\n//\n// Deprecated: Prefer directly including the persona message in the\n// input when calling Generate or Stream to avoid extra copying.\nfunc NewPersonaModifier(persona string) MessageModifier {\n\treturn func(ctx context.Context, input []*schema.Message) []*schema.Message {\n\t\tres := make([]*schema.Message, 0, len(input)+1)\n\n\t\tres = append(res, schema.SystemMessage(persona))\n\t\tres = append(res, input...)\n\t\treturn res\n\t}\n}\n\nfunc firstChunkStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) {\n\tdefer sr.Close()\n\n\tfor {\n\t\tmsg, err := sr.Recv()\n\t\tif err == io.EOF {\n\t\t\treturn false, nil\n\t\t}\n\t\tif err != nil {\n\t\t\treturn false, err\n\t\t}\n\n\t\tif len(msg.ToolCalls) > 0 {\n\t\t\treturn true, nil\n\t\t}\n\n\t\tif len(msg.Content) == 0 { // skip empty chunks at the front\n\t\t\tcontinue\n\t\t}\n\n\t\treturn false, nil\n\t}\n}\n\n// Default graph and node names for the ReAct agent.\nconst (\n\tGraphName     = \"ReActAgent\"\n\tModelNodeName = \"ChatModel\"\n\tToolsNodeName = \"Tools\"\n)\n\n// SetReturnDirectly is a helper function that can be called within a tool's execution.\n// It signals the ReAct agent to stop further processing and return the result of the current tool call directly.\n// This is useful when the tool's output is the final answer and no more steps are needed.\n// Note: If multiple tools call this function in the same step, only the last call will take effect.\n// This setting has a higher priority than the AgentConfig.ToolReturnDirectly.\nfunc SetReturnDirectly(ctx context.Context) error {\n\treturn compose.ProcessState(ctx, func(ctx context.Context, s *state) error {\n\t\ts.ReturnDirectlyToolCallID = compose.GetToolCallID(ctx)\n\t\treturn nil\n\t})\n}\n\n// Agent is the ReAct agent.\n// ReAct agent is a simple agent that handles user messages with a chat model and tools.\n// ReAct will call the chat model, if the message contains tool calls, it will call the tools.\n// if the tool is configured to return directly, ReAct will return directly.\n// otherwise, ReAct will continue to call the chat model until the message contains no tool calls.\n// e.g.\n//\n//\tagent, err := ReAct.NewAgent(ctx, &react.AgentConfig{})\n//\tif err != nil {...}\n//\tmsg, err := agent.Generate(ctx, []*schema.Message{{Role: schema.User, Content: \"how to build agent with eino\"}})\n//\tif err != nil {...}\n//\tprintln(msg.Content)\ntype Agent struct {\n\trunnable         compose.Runnable[[]*schema.Message, *schema.Message]\n\tgraph            *compose.Graph[[]*schema.Message, *schema.Message]\n\tgraphAddNodeOpts []compose.GraphAddNodeOpt\n}\n\n// NewAgent creates a ReAct agent that feeds tool response into next round of Chat Model generation.\n//\n// IMPORTANT!! For models that don't output tool calls in the first streaming chunk (e.g. Claude)\n// the default StreamToolCallChecker may not work properly since it only checks the first chunk for tool calls.\n// In such cases, you need to implement a custom StreamToolCallChecker that can properly detect tool calls.\nfunc NewAgent(ctx context.Context, config *AgentConfig) (_ *Agent, err error) {\n\tvar (\n\t\tchatModel       model.BaseChatModel\n\t\ttoolsNode       *compose.ToolsNode\n\t\ttoolInfos       []*schema.ToolInfo\n\t\ttoolCallChecker = config.StreamToolCallChecker\n\t\tmessageModifier = config.MessageModifier\n\t)\n\n\tgraphName := GraphName\n\tif config.GraphName != \"\" {\n\t\tgraphName = config.GraphName\n\t}\n\n\tmodelNodeName := ModelNodeName\n\tif config.ModelNodeName != \"\" {\n\t\tmodelNodeName = config.ModelNodeName\n\t}\n\n\ttoolsNodeName := ToolsNodeName\n\tif config.ToolsNodeName != \"\" {\n\t\ttoolsNodeName = config.ToolsNodeName\n\t}\n\n\tif toolCallChecker == nil {\n\t\ttoolCallChecker = firstChunkStreamToolCallChecker\n\t}\n\n\tif toolInfos, err = genToolInfos(ctx, config.ToolsConfig); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif chatModel, err = agent.ChatModelWithTools(config.Model, config.ToolCallingModel, toolInfos); err != nil {\n\t\treturn nil, err\n\t}\n\n\tconfig.ToolsConfig.ToolCallMiddlewares = append(\n\t\t[]compose.ToolMiddleware{newToolResultCollectorMiddleware()},\n\t\tconfig.ToolsConfig.ToolCallMiddlewares...,\n\t)\n\n\tif toolsNode, err = compose.NewToolNode(ctx, &config.ToolsConfig); err != nil {\n\t\treturn nil, err\n\t}\n\n\tgraph := compose.NewGraph[[]*schema.Message, *schema.Message](compose.WithGenLocalState(func(ctx context.Context) *state {\n\t\treturn &state{Messages: make([]*schema.Message, 0, config.MaxStep+1)}\n\t}))\n\n\tmodelPreHandle := func(ctx context.Context, input []*schema.Message, state *state) ([]*schema.Message, error) {\n\t\tstate.Messages = append(state.Messages, input...)\n\n\t\tif config.MessageRewriter != nil {\n\t\t\tstate.Messages = config.MessageRewriter(ctx, state.Messages)\n\t\t}\n\n\t\tif messageModifier == nil {\n\t\t\treturn state.Messages, nil\n\t\t}\n\n\t\tmodifiedInput := make([]*schema.Message, len(state.Messages))\n\t\tcopy(modifiedInput, state.Messages)\n\t\treturn messageModifier(ctx, modifiedInput), nil\n\t}\n\n\tif err = graph.AddChatModelNode(nodeKeyModel, chatModel, compose.WithStatePreHandler(modelPreHandle), compose.WithNodeName(modelNodeName)); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif err = graph.AddEdge(compose.START, nodeKeyModel); err != nil {\n\t\treturn nil, err\n\t}\n\n\ttoolsNodePreHandle := func(ctx context.Context, input *schema.Message, state *state) (*schema.Message, error) {\n\t\tif input == nil {\n\t\t\treturn state.Messages[len(state.Messages)-1], nil // used for rerun interrupt resume\n\t\t}\n\t\tstate.Messages = append(state.Messages, input)\n\t\tstate.ReturnDirectlyToolCallID = getReturnDirectlyToolCallID(input, config.ToolReturnDirectly)\n\t\treturn input, nil\n\t}\n\tif err = graph.AddToolsNode(nodeKeyTools, toolsNode, compose.WithStatePreHandler(toolsNodePreHandle), compose.WithNodeName(toolsNodeName)); err != nil {\n\t\treturn nil, err\n\t}\n\n\tmodelPostBranchCondition := func(ctx context.Context, sr *schema.StreamReader[*schema.Message]) (endNode string, err error) {\n\t\tif isToolCall, err := toolCallChecker(ctx, sr); err != nil {\n\t\t\treturn \"\", err\n\t\t} else if isToolCall {\n\t\t\treturn nodeKeyTools, nil\n\t\t}\n\t\treturn compose.END, nil\n\t}\n\n\tif err = graph.AddBranch(nodeKeyModel, compose.NewStreamGraphBranch(modelPostBranchCondition, map[string]bool{nodeKeyTools: true, compose.END: true})); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif err = buildReturnDirectly(graph); err != nil {\n\t\treturn nil, err\n\t}\n\n\tcompileOpts := []compose.GraphCompileOption{compose.WithMaxRunSteps(config.MaxStep), compose.WithNodeTriggerMode(compose.AnyPredecessor), compose.WithGraphName(graphName)}\n\trunnable, err := graph.Compile(ctx, compileOpts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &Agent{\n\t\trunnable:         runnable,\n\t\tgraph:            graph,\n\t\tgraphAddNodeOpts: []compose.GraphAddNodeOpt{compose.WithGraphCompileOptions(compileOpts...)},\n\t}, nil\n}\n\nfunc buildReturnDirectly(graph *compose.Graph[[]*schema.Message, *schema.Message]) (err error) {\n\tdirectReturn := func(ctx context.Context, msgs *schema.StreamReader[[]*schema.Message]) (*schema.StreamReader[*schema.Message], error) {\n\t\treturn schema.StreamReaderWithConvert(msgs, func(msgs []*schema.Message) (*schema.Message, error) {\n\t\t\tvar msg *schema.Message\n\t\t\terr = compose.ProcessState[*state](ctx, func(_ context.Context, state *state) error {\n\t\t\t\tfor i := range msgs {\n\t\t\t\t\tif msgs[i] != nil && msgs[i].ToolCallID == state.ReturnDirectlyToolCallID {\n\t\t\t\t\t\tmsg = msgs[i]\n\t\t\t\t\t\treturn nil\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn nil\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tif msg == nil {\n\t\t\t\treturn nil, schema.ErrNoValue\n\t\t\t}\n\t\t\treturn msg, nil\n\t\t}), nil\n\t}\n\n\tnodeKeyDirectReturn := \"direct_return\"\n\tif err = graph.AddLambdaNode(nodeKeyDirectReturn, compose.TransformableLambda(directReturn)); err != nil {\n\t\treturn err\n\t}\n\n\t// this branch checks if the tool called should return directly. It either leads to END or back to ChatModel\n\terr = graph.AddBranch(nodeKeyTools, compose.NewStreamGraphBranch(func(ctx context.Context, msgsStream *schema.StreamReader[[]*schema.Message]) (endNode string, err error) {\n\t\tmsgsStream.Close()\n\n\t\terr = compose.ProcessState[*state](ctx, func(_ context.Context, state *state) error {\n\t\t\tif len(state.ReturnDirectlyToolCallID) > 0 {\n\t\t\t\tendNode = nodeKeyDirectReturn\n\t\t\t} else {\n\t\t\t\tendNode = nodeKeyModel\n\t\t\t}\n\t\t\treturn nil\n\t\t})\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn endNode, nil\n\t}, map[string]bool{nodeKeyModel: true, nodeKeyDirectReturn: true}))\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treturn graph.AddEdge(nodeKeyDirectReturn, compose.END)\n}\n\nfunc genToolInfos(ctx context.Context, config compose.ToolsNodeConfig) ([]*schema.ToolInfo, error) {\n\ttoolInfos := make([]*schema.ToolInfo, 0, len(config.Tools))\n\tfor _, t := range config.Tools {\n\t\ttl, err := t.Info(ctx)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\ttoolInfos = append(toolInfos, tl)\n\t}\n\n\treturn toolInfos, nil\n}\n\nfunc getReturnDirectlyToolCallID(input *schema.Message, toolReturnDirectly map[string]struct{}) string {\n\tif len(toolReturnDirectly) == 0 {\n\t\treturn \"\"\n\t}\n\n\tfor _, toolCall := range input.ToolCalls {\n\t\tif _, ok := toolReturnDirectly[toolCall.Function.Name]; ok {\n\t\t\treturn toolCall.ID\n\t\t}\n\t}\n\n\treturn \"\"\n}\n\n// Generate generates a response from the agent.\nfunc (r *Agent) Generate(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.Message, error) {\n\treturn r.runnable.Invoke(ctx, input, agent.GetComposeOptions(opts...)...)\n}\n\n// Stream calls the agent and returns a stream response.\nfunc (r *Agent) Stream(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (output *schema.StreamReader[*schema.Message], err error) {\n\treturn r.runnable.Stream(ctx, input, agent.GetComposeOptions(opts...)...)\n}\n\n// ExportGraph exports the underlying graph from Agent, along with the []compose.GraphAddNodeOpt to be used when adding this graph to another graph.\nfunc (r *Agent) ExportGraph() (compose.AnyGraph, []compose.GraphAddNodeOpt) {\n\treturn r.graph, r.graphAddNodeOpts\n}\n"
  },
  {
    "path": "flow/agent/react/react_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage react\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"math/rand\"\n\t\"testing\"\n\n\t\"github.com/bytedance/sonic\"\n\t\"github.com/stretchr/testify/assert\"\n\t\"go.uber.org/mock/gomock\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/flow/agent\"\n\tmockModel \"github.com/cloudwego/eino/internal/mock/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n\ttemplate \"github.com/cloudwego/eino/utils/callbacks\"\n)\n\nfunc TestReact(t *testing.T) {\n\tctx := context.Background()\n\n\tfakeTool := &fakeToolGreetForTest{\n\t\ttarCount: 3,\n\t}\n\n\tinfo, err := fakeTool.Info(ctx)\n\tassert.NoError(t, err)\n\n\tctrl := gomock.NewController(t)\n\tcm := mockModel.NewMockChatModel(ctrl)\n\n\ttimes := 0\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\ttimes++\n\t\t\tif times <= 2 {\n\t\t\t\tinfo, _ := fakeTool.Info(ctx)\n\n\t\t\t\treturn schema.AssistantMessage(\"hello max\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"123\"}`, randStr()),\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\tnil\n\t\t\t}\n\n\t\t\treturn schema.AssistantMessage(\"bye\", nil), nil\n\t\t}).AnyTimes()\n\tcm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()\n\n\ta, err := NewAgent(ctx, &AgentConfig{\n\t\tModel: cm,\n\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t},\n\t\tMessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message {\n\t\t\tassert.Equal(t, len(input), times*2+1)\n\t\t\treturn input\n\t\t},\n\t\tMaxStep: 40,\n\t})\n\tassert.Nil(t, err)\n\n\tout, err := a.Generate(ctx, []*schema.Message{\n\t\t{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!\",\n\t\t},\n\t}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))\n\tassert.Nil(t, err)\n\n\tif out != nil {\n\t\tt.Log(out.Content)\n\t}\n\n\t// test return directly\n\ttimes = 0\n\ta, err = NewAgent(ctx, &AgentConfig{\n\t\tModel: cm,\n\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t},\n\t\tMessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message {\n\t\t\tassert.Equal(t, len(input), times*2+1)\n\t\t\treturn input\n\t\t},\n\t\tMaxStep:            40,\n\t\tToolReturnDirectly: map[string]struct{}{info.Name: {}},\n\t})\n\tassert.Nil(t, err)\n\n\tout, err = a.Generate(ctx, []*schema.Message{\n\t\t{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!\",\n\t\t},\n\t}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))\n\tassert.Nil(t, err)\n\n\tif out != nil {\n\t\tt.Log(out.Content)\n\t}\n}\n\nfunc TestReactWithMessageRewriterAndModifier(t *testing.T) {\n\tctx := context.Background()\n\n\tctrl := gomock.NewController(t)\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\t// This test simulates a single Generate call with a long history.\n\t// The MessageRewriter should shorten the history.\n\t// The MessageModifier should add a system prompt.\n\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\t// Check messages passed to the model.\n\t\t\t// Expected: [system prompt, user: \"message 2\", assistant: \"response 2\"]\n\t\t\tassert.Len(t, input, 3)\n\t\t\tassert.Equal(t, schema.System, input[0].Role)\n\t\t\tassert.Equal(t, \"system prompt\", input[0].Content)\n\t\t\tassert.Equal(t, schema.User, input[1].Role)\n\t\t\tassert.Equal(t, \"message 2\", input[1].Content)\n\t\t\tassert.Equal(t, schema.Assistant, input[2].Role)\n\t\t\tassert.Equal(t, \"response 2\", input[2].Content)\n\t\t\treturn schema.AssistantMessage(\"final response\", nil), nil\n\t\t}).Times(1)\n\tcm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()\n\n\tra, err := NewAgent(ctx, &AgentConfig{\n\t\tToolCallingModel: cm,\n\t\tMessageRewriter: func(ctx context.Context, messages []*schema.Message) []*schema.Message {\n\t\t\t// Keep only the last 2 messages if history is longer.\n\t\t\tassert.Len(t, messages, 4) // user1, assistant1, user2, assistant2\n\t\t\tif len(messages) > 2 {\n\t\t\t\treturn messages[len(messages)-2:]\n\t\t\t}\n\t\t\treturn messages\n\t\t},\n\t\tMessageModifier: func(ctx context.Context, messages []*schema.Message) []*schema.Message {\n\t\t\t// messages should be the result from rewriter\n\t\t\tassert.Len(t, messages, 2) // user2, assistant2\n\n\t\t\t// Add a system prompt\n\t\t\tres := make([]*schema.Message, 0, len(messages)+1)\n\t\t\tres = append(res, schema.SystemMessage(\"system prompt\"))\n\t\t\tres = append(res, messages...)\n\t\t\treturn res\n\t\t},\n\t})\n\tassert.NoError(t, err)\n\n\t// Simulate a conversation history\n\thistory := []*schema.Message{\n\t\tschema.UserMessage(\"message 1\"),\n\t\tschema.AssistantMessage(\"response 1\", nil),\n\t\tschema.UserMessage(\"message 2\"),\n\t\tschema.AssistantMessage(\"response 2\", nil),\n\t}\n\n\t// Run the react agent\n\tfinalMsg, err := ra.Generate(ctx, history)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"final response\", finalMsg.Content)\n}\n\nfunc TestReactStream(t *testing.T) {\n\tctx := context.Background()\n\n\tfakeTool := &fakeToolGreetForTest{\n\t\ttarCount: 20,\n\t}\n\n\tfakeStreamTool := &fakeStreamToolGreetForTest{\n\t\ttarCount: 20,\n\t}\n\n\tctrl := gomock.NewController(t)\n\tcm := mockModel.NewMockChatModel(ctrl)\n\n\ttimes := 0\n\tcm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()\n\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (\n\t\t\t*schema.StreamReader[*schema.Message], error) {\n\t\t\tsr, sw := schema.Pipe[*schema.Message](1)\n\t\t\tdefer sw.Close()\n\n\t\t\tinfo, _ := fakeTool.Info(ctx)\n\t\t\tstreamInfo, _ := fakeStreamTool.Info(ctx)\n\n\t\t\ttimes++\n\t\t\tif times <= 2 {\n\t\t\t\tsw.Send(schema.AssistantMessage(\"hello max\",\n\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"tool\"}`, randStr()),\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t}),\n\t\t\t\t\tnil)\n\t\t\t\treturn sr, nil\n\t\t\t} else if times == 3 {\n\t\t\t\tsw.Send(schema.AssistantMessage(\"hello max\",\n\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\tName:      streamInfo.Name,\n\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"stream tool\"}`, randStr()),\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t}),\n\t\t\t\t\tnil)\n\t\t\t\treturn sr, nil\n\t\t\t} else if times == 4 { // parallel tool call\n\t\t\t\tsw.Send(schema.AssistantMessage(\"hello max\",\n\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"tool\"}`, randStr()),\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\tName:      streamInfo.Name,\n\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"stream tool\"}`, randStr()),\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t}),\n\t\t\t\t\tnil)\n\t\t\t\treturn sr, nil\n\t\t\t}\n\n\t\t\tsw.Send(schema.AssistantMessage(\"bye\", nil), nil)\n\t\t\treturn sr, nil\n\t\t}).AnyTimes()\n\n\ta, err := NewAgent(ctx, &AgentConfig{\n\t\tModel: cm,\n\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{fakeTool, fakeStreamTool},\n\t\t},\n\n\t\tMaxStep: 40,\n\t})\n\tassert.Nil(t, err)\n\n\tout, err := a.Stream(ctx, []*schema.Message{\n\t\t{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!\",\n\t\t},\n\t}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tdefer out.Close()\n\n\tmsgs := make([]*schema.Message, 0)\n\tfor {\n\t\tmsg, err := out.Recv()\n\t\tif err != nil {\n\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tmsgs = append(msgs, msg)\n\t}\n\n\tassert.Equal(t, 1, len(msgs))\n\n\tmsg, err := schema.ConcatMessages(msgs)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tt.Log(msg.Content)\n\n\tinfo, err := fakeStreamTool.Info(ctx)\n\tassert.NoError(t, err)\n\n\t// test return directly\n\ta, err = NewAgent(ctx, &AgentConfig{\n\t\tModel: cm,\n\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{fakeTool, fakeStreamTool},\n\t\t},\n\n\t\tMaxStep:            40,\n\t\tToolReturnDirectly: map[string]struct{}{info.Name: {}}, // one of the two tools is return directly\n\t})\n\tassert.Nil(t, err)\n\n\ttimes = 0\n\tout, err = a.Stream(ctx, []*schema.Message{\n\t\t{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!\",\n\t\t},\n\t}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tdefer out.Close()\n\n\tmsgs = make([]*schema.Message, 0)\n\tfor {\n\t\tmsg, err := out.Recv()\n\t\tif err != nil {\n\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tmsgs = append(msgs, msg)\n\t}\n\n\tassert.Equal(t, 1, len(msgs))\n\n\tmsg, err = schema.ConcatMessages(msgs)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tt.Log(msg.Content)\n\n\t// return directly tool call within parallel tool calls\n\tout, err = a.Stream(ctx, []*schema.Message{\n\t\t{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!\",\n\t\t},\n\t}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))\n\tassert.NoError(t, err)\n\n\tdefer out.Close()\n\n\tmsgs = make([]*schema.Message, 0)\n\tfor {\n\t\tmsg, err := out.Recv()\n\t\tif err != nil {\n\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t}\n\n\t\tmsgs = append(msgs, msg)\n\t}\n\n\tassert.Equal(t, 1, len(msgs))\n\n\tmsg, err = schema.ConcatMessages(msgs)\n\tassert.NoError(t, err)\n\n\tt.Log(\"parallel tool call with return directly: \", msg.Content)\n}\n\nfunc TestReactWithModifier(t *testing.T) {\n\tctx := context.Background()\n\n\tfakeTool := &fakeToolGreetForTest{}\n\tctrl := gomock.NewController(t)\n\tcm := mockModel.NewMockChatModel(ctrl)\n\n\ttimes := 0\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\ttimes++\n\t\t\tif times <= 2 {\n\t\t\t\tinfo, _ := fakeTool.Info(ctx)\n\n\t\t\t\treturn schema.AssistantMessage(\"hello max\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"123\"}`, randStr()),\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\tnil\n\t\t\t}\n\n\t\t\treturn schema.AssistantMessage(\"bye\", nil), nil\n\t\t}).AnyTimes()\n\tcm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()\n\n\ta, err := NewAgent(ctx, &AgentConfig{\n\t\tModel: cm,\n\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\tTools: []tool.BaseTool{fakeTool},\n\t\t},\n\t\tMessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message {\n\t\t\tres := make([]*schema.Message, 0, len(input)+1)\n\n\t\t\tres = append(res, schema.SystemMessage(\"you are a helpful assistant\"))\n\t\t\tres = append(res, input...)\n\t\t\treturn res\n\t\t},\n\n\t\tMaxStep: 40,\n\t})\n\n\tassert.Nil(t, err)\n\n\tout, err := a.Generate(ctx, []*schema.Message{\n\t\t{\n\t\t\tRole:    schema.User,\n\t\t\tContent: \"hello\",\n\t\t},\n\t}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tif out != nil {\n\t\tt.Log(out.Content)\n\t}\n}\n\nfunc TestAgentInGraph(t *testing.T) {\n\tt.Run(\"agent generate in chain\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tfakeTool := &fakeToolGreetForTest{}\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockChatModel(ctrl)\n\n\t\ttimes := 0\n\t\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\n\t\t\t\ttimes += 1\n\t\t\t\tif times <= 2 {\n\t\t\t\t\tinfo, _ := fakeTool.Info(ctx)\n\n\t\t\t\t\treturn schema.AssistantMessage(\"hello max\",\n\t\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"123\"}`, randStr()),\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil\n\t\t\t\t}\n\n\t\t\t\treturn schema.AssistantMessage(\"bye\", nil), nil\n\n\t\t\t}).Times(3)\n\t\tcm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()\n\n\t\ta, err := NewAgent(ctx, &AgentConfig{\n\t\t\tModel: cm,\n\t\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{fakeTool, &fakeStreamToolGreetForTest{}},\n\t\t\t},\n\n\t\t\tMaxStep: 40,\n\t\t})\n\t\tassert.Nil(t, err)\n\n\t\tchain := compose.NewChain[[]*schema.Message, string]()\n\t\tagentLambda, err := compose.AnyLambda(a.Generate, a.Stream, nil, nil)\n\t\tassert.Nil(t, err)\n\n\t\tchain.\n\t\t\tAppendLambda(agentLambda).\n\t\t\tAppendLambda(compose.InvokableLambda(func(ctx context.Context, input *schema.Message) (string, error) {\n\t\t\t\tt.Log(\"got agent response: \", input.Content)\n\t\t\t\treturn input.Content, nil\n\t\t\t}))\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.Nil(t, err)\n\n\t\tres, err := r.Invoke(ctx, []*schema.Message{{Role: schema.User, Content: \"hello\"}},\n\t\t\tcompose.WithCallbacks(callbackForTest))\n\t\tassert.Nil(t, err)\n\n\t\tt.Log(res)\n\t})\n\n\tt.Run(\"agent stream in chain\", func(t *testing.T) {\n\n\t\tfakeStreamTool := &fakeStreamToolGreetForTest{}\n\t\tctx := context.Background()\n\t\tctrl := gomock.NewController(t)\n\t\tcm := mockModel.NewMockChatModel(ctrl)\n\n\t\ttimes := 0\n\t\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (\n\t\t\t\t*schema.StreamReader[*schema.Message], error) {\n\t\t\t\tsr, sw := schema.Pipe[*schema.Message](1)\n\t\t\t\tdefer sw.Close()\n\n\t\t\t\ttimes += 1\n\t\t\t\tif times <= 2 {\n\t\t\t\t\tinfo, _ := fakeStreamTool.Info(ctx)\n\t\t\t\t\tsw.Send(schema.AssistantMessage(\"hello max\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\tArguments: fmt.Sprintf(`{\"name\": \"%s\", \"hh\": \"123\"}`, randStr()),\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\t\tnil)\n\t\t\t\t\treturn sr, nil\n\t\t\t\t}\n\n\t\t\t\tsw.Send(schema.AssistantMessage(\"bye\", nil), nil)\n\t\t\t\treturn sr, nil\n\t\t\t}).Times(3)\n\t\tcm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()\n\n\t\ta, err := NewAgent(ctx, &AgentConfig{\n\t\t\tModel: cm,\n\t\t\tToolsConfig: compose.ToolsNodeConfig{\n\t\t\t\tTools: []tool.BaseTool{&fakeToolGreetForTest{}, fakeStreamTool},\n\t\t\t},\n\n\t\t\tMaxStep: 40,\n\t\t})\n\t\tassert.Nil(t, err)\n\n\t\tchain := compose.NewChain[[]*schema.Message, string]()\n\t\tagentGraph, opts := a.ExportGraph()\n\t\tassert.Nil(t, err)\n\n\t\tchain.\n\t\t\tAppendGraph(agentGraph, opts...).\n\t\t\tAppendLambda(compose.InvokableLambda(func(ctx context.Context, input *schema.Message) (string, error) {\n\t\t\t\tt.Log(\"got agent response: \", input.Content)\n\t\t\t\treturn input.Content, nil\n\t\t\t}))\n\t\tr, err := chain.Compile(ctx)\n\t\tassert.Nil(t, err)\n\n\t\toutStream, err := r.Stream(ctx, []*schema.Message{{Role: schema.User, Content: \"hello\"}},\n\t\t\tcompose.WithCallbacks(callbackForTest))\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tdefer outStream.Close()\n\n\t\tmsg := \"\"\n\t\tfor {\n\t\t\tmsgItem, err := outStream.Recv()\n\t\t\tif err != nil {\n\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\n\t\t\tmsg += msgItem\n\t\t}\n\n\t\tt.Log(msg)\n\t})\n\n}\n\nfunc TestWithTools(t *testing.T) {\n\tctx := context.Background()\n\n\tfakeTool := &fakeToolGreetForTest{\n\t\ttarCount: 2,\n\t}\n\tfakeStreamTool := &fakeStreamToolGreetForTest{\n\t\ttarCount: 2,\n\t}\n\n\tctrl := gomock.NewController(t)\n\tcm := mockModel.NewMockToolCallingChatModel(ctrl)\n\n\ttimes := 0\n\tcm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\t\t\ttimes++\n\t\t\tif times <= 1 {\n\t\t\t\tinfo, _ := fakeTool.Info(ctx)\n\t\t\t\treturn schema.AssistantMessage(\"calling tool\",\n\t\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\t\tArguments: `{\"name\": \"test\"}`,\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t}),\n\t\t\t\t\tnil\n\t\t\t}\n\t\t\treturn schema.AssistantMessage(\"done\", nil), nil\n\t\t}).AnyTimes()\n\n\tcm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).\n\t\tDoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (\n\t\t\t*schema.StreamReader[*schema.Message], error) {\n\t\t\tsr, sw := schema.Pipe[*schema.Message](1)\n\t\t\tdefer sw.Close()\n\n\t\t\ttimes++\n\t\t\tif times <= 2 {\n\t\t\t\tinfo, _ := fakeStreamTool.Info(ctx)\n\t\t\t\tsw.Send(schema.AssistantMessage(\"calling stream tool\",\n\t\t\t\t\t[]schema.ToolCall{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\tID: randStr(),\n\t\t\t\t\t\t\tFunction: schema.FunctionCall{\n\t\t\t\t\t\t\t\tName:      info.Name,\n\t\t\t\t\t\t\t\tArguments: `{\"name\": \"test\"}`,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t}),\n\t\t\t\t\tnil)\n\t\t\t\treturn sr, nil\n\t\t\t}\n\n\t\t\tsw.Send(schema.AssistantMessage(\"stream done\", nil), nil)\n\t\t\treturn sr, nil\n\t\t}).AnyTimes()\n\n\t// Test WithTools function\n\ttoolOptions, err := WithTools(ctx, fakeTool, fakeStreamTool)\n\tassert.NoError(t, err)\n\tassert.Len(t, toolOptions, 2, \"WithTools should return exactly 2 options\")\n\n\t// Create agent without tools in config\n\ta, err := NewAgent(ctx, &AgentConfig{\n\t\tToolCallingModel: cm,\n\t\tMaxStep:          10,\n\t})\n\tassert.NoError(t, err)\n\n\t// Test Generate with WithTools options\n\ttimes = 0\n\tmsg, err := a.Generate(ctx, []*schema.Message{\n\t\tschema.UserMessage(\"test generate with tools\"),\n\t}, toolOptions...)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"done\", msg.Content)\n\n\t// Test Stream with WithTools options\n\ttimes = 0\n\tstream, err := a.Stream(ctx, []*schema.Message{\n\t\tschema.UserMessage(\"test stream with tools\"),\n\t}, toolOptions...)\n\tassert.NoError(t, err)\n\n\tdefer stream.Close()\n\tmsgs := make([]*schema.Message, 0)\n\tfor {\n\t\tmsg, err := stream.Recv()\n\t\tif err != nil {\n\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tassert.NoError(t, err)\n\t\t}\n\t\tmsgs = append(msgs, msg)\n\t}\n\n\tassert.Len(t, msgs, 1)\n\tconcatMsg, err := schema.ConcatMessages(msgs)\n\tassert.NoError(t, err)\n\tassert.Equal(t, \"stream done\", concatMsg.Content)\n\n\t// Test error case - tool Info() returns error\n\terrorTool := &errorToolForTest{}\n\t_, err = WithTools(ctx, errorTool)\n\tassert.Error(t, err)\n\tassert.Contains(t, err.Error(), \"info error\")\n}\n\n// Helper tool for testing error cases\ntype errorToolForTest struct{}\n\nfunc (t *errorToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn nil, errors.New(\"info error\")\n}\n\nfunc (t *errorToolForTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {\n\treturn \"\", nil\n}\n\ntype fakeStreamToolGreetForTest struct {\n\ttarCount int\n\tcurCount int\n}\n\nfunc (t *fakeStreamToolGreetForTest) StreamableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (\n\t*schema.StreamReader[string], error) {\n\tp := &fakeToolInput{}\n\terr := sonic.UnmarshalString(argumentsInJSON, p)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tif t.curCount >= t.tarCount {\n\t\ts := schema.StreamReaderFromArray([]string{`{\"say\": \"bye\"}`})\n\t\treturn s, nil\n\t}\n\tt.curCount++\n\ts := schema.StreamReaderFromArray([]string{fmt.Sprintf(`{\"say\": \"hello %v\"}`, p.Name)})\n\treturn s, nil\n}\n\ntype fakeToolGreetForTest struct {\n\ttarCount int\n\tcurCount int\n}\n\nfunc (t *fakeToolGreetForTest) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"greet\",\n\t\tDesc: \"greet with name\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"name\": {\n\t\t\t\t\tDesc:     \"user name who to greet\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t\tType:     schema.String,\n\t\t\t\t},\n\t\t\t}),\n\t}, nil\n}\n\nfunc (t *fakeStreamToolGreetForTest) Info(_ context.Context) (*schema.ToolInfo, error) {\n\treturn &schema.ToolInfo{\n\t\tName: \"greet in stream\",\n\t\tDesc: \"greet with name in stream\",\n\t\tParamsOneOf: schema.NewParamsOneOfByParams(\n\t\t\tmap[string]*schema.ParameterInfo{\n\t\t\t\t\"name\": {\n\t\t\t\t\tDesc:     \"user name who to greet\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t\tType:     schema.String,\n\t\t\t\t},\n\t\t\t}),\n\t}, nil\n}\n\nfunc (t *fakeToolGreetForTest) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {\n\tp := &fakeToolInput{}\n\terr := sonic.UnmarshalString(argumentsInJSON, p)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tif t.curCount >= t.tarCount {\n\t\treturn `{\"say\": \"bye\"}`, nil\n\t}\n\n\tt.curCount++\n\treturn fmt.Sprintf(`{\"say\": \"hello %v\"}`, p.Name), nil\n}\n\ntype fakeToolInput struct {\n\tName string `json:\"name\"`\n}\n\nfunc randStr() string {\n\tseeds := []rune(\"this is a seed\")\n\tb := make([]rune, 8)\n\tfor i := range b {\n\t\tb[i] = seeds[rand.Intn(len(seeds))]\n\t}\n\treturn string(b)\n}\n\nvar callbackForTest = BuildAgentCallback(&template.ModelCallbackHandler{}, &template.ToolCallbackHandler{})\n"
  },
  {
    "path": "flow/agent/utils.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage agent\n\nimport (\n\t\"errors\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// ChatModelWithTools returns a chat model configured with tool schemas.\n// If a ToolCallingChatModel is provided, it is used directly (and optionally\n// configured with tools). Otherwise, a plain ChatModel is bound with tools.\nfunc ChatModelWithTools(cm model.ChatModel, toolCallingModel model.ToolCallingChatModel, toolInfos []*schema.ToolInfo) (\n\tmodel.BaseChatModel, error) {\n\n\tif toolCallingModel != nil {\n\t\tif len(toolInfos) == 0 {\n\t\t\treturn toolCallingModel, nil\n\t\t}\n\t\treturn toolCallingModel.WithTools(toolInfos)\n\t}\n\n\tif cm != nil {\n\t\tif len(toolInfos) == 0 {\n\t\t\treturn cm, nil\n\t\t}\n\t\terr := cm.BindTools(toolInfos)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn cm, nil\n\t}\n\n\treturn nil, errors.New(\"no chat model provided\")\n}\n"
  },
  {
    "path": "flow/indexer/parent/parent.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package parent provides an indexer that assigns stable IDs to sub-documents\n// and preserves relationships to their original parent document.\npackage parent\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Config configures the parent indexer that assigns IDs to sub-documents.\ntype Config struct {\n\t// Indexer is the underlying indexer implementation that handles the actual document indexing.\n\t// For example: a vector database indexer like Milvus, or a full-text search indexer like Elasticsearch.\n\tIndexer indexer.Indexer\n\n\t// Transformer processes documents before indexing, typically splitting them into smaller chunks.\n\t// Each sub-document generated by the transformer must retain its parent document's ID.\n\t// For example: if a document with ID \"doc_1\" is split into 3 chunks, all chunks will initially\n\t// have ID \"doc_1\". These IDs will later be modified by the SubIDGenerator.\n\t//\n\t// Example transformations:\n\t// - A text splitter that breaks down large documents into paragraphs\n\t// - A code splitter that separates code files into functions\n\tTransformer document.Transformer\n\n\t// ParentIDKey specifies the metadata key used to store the original document's ID in each sub-document.\n\t// For example: if ParentIDKey is \"parent_id\", each sub-document will have metadata like:\n\t// {\"parent_id\": \"original_doc_123\"}\n\tParentIDKey string\n\n\t// SubIDGenerator generates unique IDs for sub-documents based on their parent document ID.\n\t// For example: if parent ID is \"doc_1\" and we need 3 sub-document IDs, it might generate:\n\t// [\"doc_1_chunk_1\", \"doc_1_chunk_2\", \"doc_1_chunk_3\"]\n\t//\n\t// Parameters:\n\t//   - ctx: context for the operation\n\t//   - parentID: the ID of the parent document\n\t//   - num: number of sub-document IDs needed\n\t// Returns:\n\t//   - []string: slice of generated sub-document IDs\n\t//   - error: any error encountered during ID generation\n\tSubIDGenerator func(ctx context.Context, parentID string, num int) ([]string, error)\n}\n\n// NewIndexer creates a new parent indexer that handles document splitting and sub-document management.\n//\n// Parameters:\n//   - ctx: context for the operation\n//   - config: configuration for the parent indexer\n//\n// Example usage:\n//\n//\tindexer, err := NewIndexer(ctx, &Config{\n//\t    Indexer: milvusIndexer,\n//\t    Transformer: textSplitter,\n//\t    ParentIDKey: \"source_doc_id\",\n//\t    SubIDGenerator: func(ctx context.Context, parentID string, num int) ([]string, error) {\n//\t        ids := make([]string, num)\n//\t        for i := 0; i < num; i++ {\n//\t            ids[i] = fmt.Sprintf(\"%s_chunk_%d\", parentID, i+1)\n//\t        }\n//\t        return ids, nil\n//\t    },\n//\t})\n//\n// Returns:\n//   - indexer.Indexer: the created parent indexer\n//   - error: any error encountered during creation\nfunc NewIndexer(ctx context.Context, config *Config) (indexer.Indexer, error) {\n\tif config.Indexer == nil {\n\t\treturn nil, fmt.Errorf(\"indexer is empty\")\n\t}\n\tif config.Transformer == nil {\n\t\treturn nil, fmt.Errorf(\"transformer is empty\")\n\t}\n\tif config.SubIDGenerator == nil {\n\t\treturn nil, fmt.Errorf(\"sub id generator is empty\")\n\t}\n\n\treturn &parentIndexer{\n\t\tindexer:        config.Indexer,\n\t\ttransformer:    config.Transformer,\n\t\tparentIDKey:    config.ParentIDKey,\n\t\tsubIDGenerator: config.SubIDGenerator,\n\t}, nil\n}\n\ntype parentIndexer struct {\n\tindexer        indexer.Indexer\n\ttransformer    document.Transformer\n\tparentIDKey    string\n\tsubIDGenerator func(ctx context.Context, parentID string, num int) ([]string, error)\n}\n\nfunc (p *parentIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) ([]string, error) {\n\tsubDocs, err := p.transformer.Transform(ctx, docs)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"transform docs fail: %w\", err)\n\t}\n\tif len(subDocs) == 0 {\n\t\treturn nil, fmt.Errorf(\"doc transformer returned no documents\")\n\t}\n\tcurrentID := subDocs[0].ID\n\tstartIdx := 0\n\tfor i, subDoc := range subDocs {\n\t\tif subDoc.MetaData == nil {\n\t\t\tsubDoc.MetaData = make(map[string]any)\n\t\t}\n\t\tsubDoc.MetaData[p.parentIDKey] = subDoc.ID\n\n\t\tif subDoc.ID == currentID {\n\t\t\tcontinue\n\t\t}\n\n\t\t// generate new doc id\n\t\tsubIDs, err_ := p.subIDGenerator(ctx, subDocs[startIdx].ID, i-startIdx)\n\t\tif err_ != nil {\n\t\t\treturn nil, err_\n\t\t}\n\t\tif len(subIDs) != i-startIdx {\n\t\t\treturn nil, fmt.Errorf(\"generated sub IDs' num is unexpected\")\n\t\t}\n\t\tfor j := startIdx; j < i; j++ {\n\t\t\tsubDocs[j].ID = subIDs[j-startIdx]\n\t\t}\n\t\tstartIdx = i\n\t\tcurrentID = subDoc.ID\n\t}\n\t// generate new doc id\n\tsubIDs, err := p.subIDGenerator(ctx, subDocs[startIdx].ID, len(subDocs)-startIdx)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif len(subIDs) != len(subDocs)-startIdx {\n\t\treturn nil, fmt.Errorf(\"generated sub IDs' num is unexpected\")\n\t}\n\tfor j := startIdx; j < len(subDocs); j++ {\n\t\tsubDocs[j].ID = subIDs[j-startIdx]\n\t}\n\n\treturn p.indexer.Store(ctx, subDocs, opts...)\n}\n"
  },
  {
    "path": "flow/indexer/parent/parent_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parent\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"reflect\"\n\t\"strconv\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype testIndexer struct{}\n\nfunc (t *testIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) (ids []string, err error) {\n\tret := make([]string, len(docs))\n\tfor i, d := range docs {\n\t\tret[i] = d.ID\n\t\tif !strings.HasPrefix(d.ID, d.MetaData[\"parent\"].(string)) {\n\t\t\treturn nil, fmt.Errorf(\"invalid parent key\")\n\t\t}\n\t}\n\treturn ret, nil\n}\n\ntype testTransformer struct {\n}\n\nfunc (t *testTransformer) Transform(ctx context.Context, src []*schema.Document, opts ...document.TransformerOption) ([]*schema.Document, error) {\n\tvar ret []*schema.Document\n\tfor _, d := range src {\n\t\tret = append(ret, &schema.Document{\n\t\t\tID:       d.ID,\n\t\t\tContent:  d.Content[:len(d.Content)/2],\n\t\t\tMetaData: deepCopyMap(d.MetaData),\n\t\t}, &schema.Document{\n\t\t\tID:       d.ID,\n\t\t\tContent:  d.Content[len(d.Content)/2:],\n\t\t\tMetaData: deepCopyMap(d.MetaData),\n\t\t})\n\t}\n\treturn ret, nil\n}\n\nfunc TestParentIndexer(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\tconfig *Config\n\t\tinput  []*schema.Document\n\t\twant   []string\n\t}{\n\t\t{\n\t\t\tname: \"success\",\n\t\t\tconfig: &Config{\n\t\t\t\tIndexer:     &testIndexer{},\n\t\t\t\tTransformer: &testTransformer{},\n\t\t\t\tParentIDKey: \"parent\",\n\t\t\t\tSubIDGenerator: func(ctx context.Context, parentID string, num int) ([]string, error) {\n\t\t\t\t\tret := make([]string, num)\n\t\t\t\t\tfor i := range ret {\n\t\t\t\t\t\tret[i] = parentID + strconv.Itoa(i)\n\t\t\t\t\t}\n\t\t\t\t\treturn ret, nil\n\t\t\t\t},\n\t\t\t},\n\t\t\tinput: []*schema.Document{{\n\t\t\t\tID:       \"id\",\n\t\t\t\tContent:  \"1234567890\",\n\t\t\t\tMetaData: map[string]interface{}{},\n\t\t\t}, {\n\t\t\t\tID:       \"ID\",\n\t\t\t\tContent:  \"0987654321\",\n\t\t\t\tMetaData: map[string]interface{}{},\n\t\t\t}},\n\t\t\twant: []string{\"id0\", \"id1\", \"ID0\", \"ID1\"},\n\t\t},\n\t}\n\tctx := context.Background()\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tindex, err := NewIndexer(ctx, tt.config)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t\tret, err := index.Store(ctx, tt.input)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t\tif !reflect.DeepEqual(ret, tt.want) {\n\t\t\t\tt.Errorf(\"NewHeaderSplitter() got = %v, want %v\", ret, tt.want)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc deepCopyMap(in map[string]interface{}) map[string]interface{} {\n\tout := make(map[string]interface{})\n\tfor k, v := range in {\n\t\tout[k] = v\n\t}\n\treturn out\n}\n"
  },
  {
    "path": "flow/retriever/multiquery/multi_query.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package multiquery implements a query-rewriting retriever that expands\n// user queries into multiple variants to improve recall.\npackage multiquery\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/flow/retriever/utils\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nconst (\n\tdefaultRewritePrompt = `You are an helpful assistant.\n\tYour role is to create three different versions of the user query to retrieve relevant documents from store.\n    Your goal is to improve the performance of similarity search by generating text from different perspectives based on the user query.\n\tOnly provide the generated queries and separate them by newlines. \n\tuser query: {{query}}`\n\tdefaultQueryVariable = \"query\"\n\tdefaultMaxQueriesNum = 5\n)\n\nvar deduplicateFusion = func(ctx context.Context, docs [][]*schema.Document) ([]*schema.Document, error) {\n\tm := map[string]bool{}\n\tvar ret []*schema.Document\n\tfor i := range docs {\n\t\tfor j := range docs[i] {\n\t\t\tif _, ok := m[docs[i][j].ID]; !ok {\n\t\t\t\tm[docs[i][j].ID] = true\n\t\t\t\tret = append(ret, docs[i][j])\n\t\t\t}\n\t\t}\n\t}\n\treturn ret, nil\n}\n\n// NewRetriever creates a multi-query retriever.\n// multi-query retriever is useful when you want to retrieve documents from multiple retrievers with different queries.\n// e.g.\n//\n//\tmultiRetriever := multiquery.NewRetriever(ctx, &multiquery.Config{})\n//\tdocs, err := multiRetriever.Retrieve(ctx, \"how to build agent with eino\")\n//\tif err != nil {\n//\t\t...\n//\t}\n//\tprintln(docs)\nfunc NewRetriever(ctx context.Context, config *Config) (retriever.Retriever, error) {\n\tvar err error\n\n\t// config validate\n\tif config.OrigRetriever == nil {\n\t\treturn nil, fmt.Errorf(\"OrigRetriever is required\")\n\t}\n\tif config.RewriteHandler == nil && config.RewriteLLM == nil {\n\t\treturn nil, fmt.Errorf(\"at least one of RewriteHandler and RewriteLLM must not be empty\")\n\t}\n\n\t// construct rewrite chain\n\trewriteChain := compose.NewChain[string, []string]()\n\tif config.RewriteHandler != nil {\n\t\trewriteChain.AppendLambda(compose.InvokableLambda(config.RewriteHandler), compose.WithNodeName(\"CustomQueryRewriter\"))\n\t} else {\n\t\ttpl := config.RewriteTemplate\n\t\tvariable := config.QueryVar\n\t\tparser := config.LLMOutputParser\n\t\tif tpl == nil {\n\t\t\ttpl = prompt.FromMessages(schema.Jinja2, schema.UserMessage(defaultRewritePrompt))\n\t\t\tvariable = defaultQueryVariable\n\t\t}\n\t\tif parser == nil {\n\t\t\tparser = func(ctx context.Context, message *schema.Message) ([]string, error) {\n\t\t\t\treturn strings.Split(message.Content, \"\\n\"), nil\n\t\t\t}\n\t\t}\n\n\t\trewriteChain.\n\t\t\tAppendLambda(compose.InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) {\n\t\t\t\treturn map[string]any{variable: input}, nil\n\t\t\t}), compose.WithNodeName(\"Converter\")).\n\t\t\tAppendChatTemplate(tpl).\n\t\t\tAppendChatModel(config.RewriteLLM).\n\t\t\tAppendLambda(compose.InvokableLambda(parser), compose.WithNodeName(\"OutputParser\"))\n\t}\n\trewriteRunner, err := rewriteChain.Compile(ctx, compose.WithGraphName(\"QueryRewrite\"))\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tmaxQueriesNum := config.MaxQueriesNum\n\tif maxQueriesNum == 0 {\n\t\tmaxQueriesNum = defaultMaxQueriesNum\n\t}\n\n\tfusionFunc := config.FusionFunc\n\tif fusionFunc == nil {\n\t\tfusionFunc = deduplicateFusion\n\t}\n\n\treturn &multiQueryRetriever{\n\t\tqueryRunner:   rewriteRunner,\n\t\tmaxQueriesNum: maxQueriesNum,\n\t\torigRetriever: config.OrigRetriever,\n\t\tfusionFunc:    fusionFunc,\n\t}, nil\n}\n\n// Config is the config for multi-query retriever.\ntype Config struct {\n\t// Rewrite\n\t// 1. set the following fields to use llm to generate multi queries\n\t// \ta. chat model, required\n\tRewriteLLM model.ChatModel\n\t//\tb. prompt llm to generate multi queries, we provide default template so you can leave this field blank\n\tRewriteTemplate prompt.ChatTemplate\n\t//\tc. origin query variable of your custom template, it can be empty if you use default template\n\tQueryVar string\n\t//\td. parser llm output to queries, split content using \"\\n\" by default\n\tLLMOutputParser func(context.Context, *schema.Message) ([]string, error)\n\t// 2. set RewriteHandler to provide custom query generation logic, possibly without a ChatModel. If this field is set, it takes precedence over other configurations above\n\tRewriteHandler func(ctx context.Context, query string) ([]string, error)\n\t// limit max queries num that Rewrite generates, and excess queries will be truncated, 5 by default\n\tMaxQueriesNum int\n\n\t// Origin Retriever\n\tOrigRetriever retriever.Retriever\n\n\t// fusion docs recalled from multi retrievers, remove dup based on document id by default\n\tFusionFunc func(ctx context.Context, docs [][]*schema.Document) ([]*schema.Document, error)\n}\n\ntype multiQueryRetriever struct {\n\tqueryRunner   compose.Runnable[string, []string]\n\tmaxQueriesNum int\n\torigRetriever retriever.Retriever\n\tfusionFunc    func(ctx context.Context, docs [][]*schema.Document) ([]*schema.Document, error)\n}\n\n// Retrieve retrieves documents from the multi-query retriever.\nfunc (m *multiQueryRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n\t// generate queries\n\tqueries, err := m.queryRunner.Invoke(ctx, query)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif len(queries) > m.maxQueriesNum {\n\t\tqueries = queries[:m.maxQueriesNum]\n\t}\n\n\t// retrieve\n\ttasks := make([]*utils.RetrieveTask, len(queries))\n\tfor i := range queries {\n\t\ttasks[i] = &utils.RetrieveTask{Retriever: m.origRetriever, Query: queries[i]}\n\t}\n\tutils.ConcurrentRetrieveWithCallback(ctx, tasks)\n\tresult := make([][]*schema.Document, len(queries))\n\tfor i, task := range tasks {\n\t\tif task.Err != nil {\n\t\t\treturn nil, task.Err\n\t\t}\n\t\tresult[i] = task.Result\n\t}\n\n\t// fusion\n\tctx = ctxWithFusionRunInfo(ctx)\n\tctx = callbacks.OnStart(ctx, result)\n\tfusionDocs, err := m.fusionFunc(ctx, result)\n\tif err != nil {\n\t\tcallbacks.OnError(ctx, err)\n\t\treturn nil, err\n\t}\n\tcallbacks.OnEnd(ctx, fusionDocs)\n\treturn fusionDocs, nil\n}\n\n// GetType returns the type of the retriever (MultiQuery).\nfunc (m *multiQueryRetriever) GetType() string {\n\treturn \"MultiQuery\"\n}\n\nfunc ctxWithFusionRunInfo(ctx context.Context) context.Context {\n\trunInfo := &callbacks.RunInfo{\n\t\tComponent: compose.ComponentOfLambda,\n\t\tType:      \"FusionFunc\",\n\t}\n\n\trunInfo.Name = runInfo.Type + string(runInfo.Component)\n\n\treturn callbacks.ReuseHandlers(ctx, runInfo)\n}\n"
  },
  {
    "path": "flow/retriever/multiquery/multi_query_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage multiquery\n\nimport (\n\t\"context\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype mockRetriever struct {\n}\n\nfunc (m *mockRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n\tvar ret []*schema.Document\n\tif strings.Contains(query, \"1\") {\n\t\tret = append(ret, &schema.Document{ID: \"1\"})\n\t}\n\tif strings.Contains(query, \"2\") {\n\t\tret = append(ret, &schema.Document{ID: \"2\"})\n\t}\n\tif strings.Contains(query, \"3\") {\n\t\tret = append(ret, &schema.Document{ID: \"3\"})\n\t}\n\tif strings.Contains(query, \"4\") {\n\t\tret = append(ret, &schema.Document{ID: \"4\"})\n\t}\n\tif strings.Contains(query, \"5\") {\n\t\tret = append(ret, &schema.Document{ID: \"5\"})\n\t}\n\treturn ret, nil\n}\n\ntype mockModel struct {\n}\n\nfunc (m *mockModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\treturn &schema.Message{\n\t\tContent: \"12\\n23\\n34\\n14\\n23\\n45\",\n\t}, nil\n}\n\nfunc (m *mockModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tpanic(\"implement me\")\n}\n\nfunc (m *mockModel) BindTools(tools []*schema.ToolInfo) error {\n\tpanic(\"implement me\")\n}\n\nfunc TestMultiQueryRetriever(t *testing.T) {\n\tctx := context.Background()\n\n\t// use default llm\n\tmqr, err := NewRetriever(ctx, &Config{\n\t\tRewriteLLM:    &mockModel{},\n\t\tOrigRetriever: &mockRetriever{},\n\t})\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tc := compose.NewChain[string, []*schema.Document]()\n\tcr, err := c.AppendRetriever(mqr).Compile(ctx)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tresult, err := cr.Invoke(ctx, \"query\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif len(result) != 4 {\n\t\tt.Fatal(\"default llm retrieve result is unexpected\")\n\t}\n\n\t// use custom\n\tmqr, err = NewRetriever(ctx, &Config{\n\t\tRewriteHandler: func(ctx context.Context, query string) ([]string, error) {\n\t\t\treturn []string{\"1\", \"3\", \"5\"}, nil\n\t\t},\n\t\tOrigRetriever: &mockRetriever{},\n\t})\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tc = compose.NewChain[string, []*schema.Document]()\n\tcr, err = c.AppendRetriever(mqr).Compile(ctx)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tresult, err = cr.Invoke(ctx, \"query\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif len(result) != 3 {\n\t\tt.Fatal(\"default llm retrieve result is unexpected\")\n\t}\n}\n"
  },
  {
    "path": "flow/retriever/parent/doc.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     https://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package parent provides a retriever that maps sub-document results\n// back to their original parent documents.\npackage parent\n"
  },
  {
    "path": "flow/retriever/parent/parent.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parent\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// Config configures the parent retriever.\ntype Config struct {\n\t// Retriever specifies the original retriever used to retrieve documents.\n\t// For example: a vector database retriever like Milvus, or a full-text search retriever like Elasticsearch.\n\tRetriever retriever.Retriever\n\t// ParentIDKey specifies the key used in the sub-document metadata to store the parent document ID.\n\t// Documents without this key will be removed from the recall results.\n\t// For example: if ParentIDKey is \"parent_id\", it will look for metadata like:\n\t// {\"parent_id\": \"original_doc_123\"}\n\tParentIDKey string\n\t// OrigDocGetter specifies the method for getting original documents by ids from the sub-document metadata.\n\t// Parameters:\n\t//   - ctx: context for the operation\n\t//   - ids: slice of parent document IDs to retrieve\n\t// Returns:\n\t//   - []*schema.Document: slice of retrieved parent documents\n\t//   - error: any error encountered during retrieval\n\t//\n\t// For example: if sub-documents with parent IDs [\"doc_1\", \"doc_2\"] are retrieved,\n\t// OrigDocGetter will be called to fetch the original documents with these IDs.\n\tOrigDocGetter func(ctx context.Context, ids []string) ([]*schema.Document, error)\n}\n\n// NewRetriever creates a new parent retriever that handles retrieving original documents\n// based on sub-document search results.\n//\n// Parameters:\n//   - ctx: context for the operation\n//   - config: configuration for the parent retriever\n//\n// Example usage:\n//\n//\tretriever, err := NewRetriever(ctx, &Config{\n//\t    Retriever: milvusRetriever,\n//\t    ParentIDKey: \"source_doc_id\",\n//\t    OrigDocGetter: func(ctx context.Context, ids []string) ([]*schema.Document, error) {\n//\t        return documentStore.GetByIDs(ctx, ids)\n//\t    },\n//\t})\n//\n// Returns:\n//   - retriever.Retriever: the created parent retriever\n//   - error: any error encountered during creation\nfunc NewRetriever(ctx context.Context, config *Config) (retriever.Retriever, error) {\n\tif config.Retriever == nil {\n\t\treturn nil, fmt.Errorf(\"retriever is required\")\n\t}\n\tif config.OrigDocGetter == nil {\n\t\treturn nil, fmt.Errorf(\"orig doc getter is required\")\n\t}\n\treturn &parentRetriever{\n\t\tretriever:     config.Retriever,\n\t\tparentIDKey:   config.ParentIDKey,\n\t\torigDocGetter: config.OrigDocGetter,\n\t}, nil\n}\n\ntype parentRetriever struct {\n\tretriever     retriever.Retriever\n\tparentIDKey   string\n\torigDocGetter func(ctx context.Context, ids []string) ([]*schema.Document, error)\n}\n\nfunc (p *parentRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n\tsubDocs, err := p.retriever.Retrieve(ctx, query, opts...)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tids := make([]string, 0, len(subDocs))\n\tfor _, subDoc := range subDocs {\n\t\tif k, ok := subDoc.MetaData[p.parentIDKey]; ok {\n\t\t\tif s, okk := k.(string); okk && !inList(s, ids) {\n\t\t\t\tids = append(ids, s)\n\t\t\t}\n\t\t}\n\t}\n\treturn p.origDocGetter(ctx, ids)\n}\n\nfunc inList(elem string, list []string) bool {\n\tfor _, v := range list {\n\t\tif v == elem {\n\t\t\treturn true\n\t\t}\n\t}\n\treturn false\n}\n"
  },
  {
    "path": "flow/retriever/parent/parent_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage parent\n\nimport (\n\t\"context\"\n\t\"reflect\"\n\t\"testing\"\n\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype testRetriever struct{}\n\nfunc (t *testRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n\tret := make([]*schema.Document, 0)\n\tfor i := range query {\n\t\tret = append(ret, &schema.Document{\n\t\t\tID:      \"\",\n\t\t\tContent: \"\",\n\t\t\tMetaData: map[string]interface{}{\n\t\t\t\t\"parent\": query[i : i+1],\n\t\t\t},\n\t\t})\n\t}\n\treturn ret, nil\n}\n\nfunc TestParentRetriever(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\tconfig *Config\n\t\tinput  string\n\t\twant   []*schema.Document\n\t}{\n\t\t{\n\t\t\tname: \"success\",\n\t\t\tconfig: &Config{\n\t\t\t\tRetriever:   &testRetriever{},\n\t\t\t\tParentIDKey: \"parent\",\n\t\t\t\tOrigDocGetter: func(ctx context.Context, ids []string) ([]*schema.Document, error) {\n\t\t\t\t\tvar ret []*schema.Document\n\t\t\t\t\tfor i := range ids {\n\t\t\t\t\t\tret = append(ret, &schema.Document{ID: ids[i]})\n\t\t\t\t\t}\n\t\t\t\t\treturn ret, nil\n\t\t\t\t},\n\t\t\t},\n\t\t\tinput: \"123233\",\n\t\t\twant: []*schema.Document{\n\t\t\t\t{ID: \"1\"},\n\t\t\t\t{ID: \"2\"},\n\t\t\t\t{ID: \"3\"},\n\t\t\t},\n\t\t},\n\t}\n\tctx := context.Background()\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tr, err := NewRetriever(ctx, tt.config)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t\tret, err := r.Retrieve(ctx, tt.input)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t\tif !reflect.DeepEqual(ret, tt.want) {\n\t\t\t\tt.Errorf(\"got %v, want %v\", ret, tt.want)\n\t\t\t}\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "flow/retriever/router/router.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package router provides retrieval routing helpers that merge results\n// from multiple retrievers and apply ranking strategies.\npackage router\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"sort\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/flow/retriever/utils\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nvar rrf = func(ctx context.Context, result map[string][]*schema.Document) ([]*schema.Document, error) {\n\tif len(result) < 1 {\n\t\treturn nil, fmt.Errorf(\"no documents\")\n\t}\n\tif len(result) == 1 {\n\t\tfor _, docs := range result {\n\t\t\treturn docs, nil\n\t\t}\n\t}\n\n\tdocRankMap := make(map[string]float64)\n\tdocMap := make(map[string]*schema.Document)\n\tfor _, v := range result {\n\t\tfor i := range v {\n\t\t\tdocMap[v[i].ID] = v[i]\n\t\t\tif _, ok := docRankMap[v[i].ID]; !ok {\n\t\t\t\tdocRankMap[v[i].ID] = 1.0 / float64(i+60)\n\t\t\t} else {\n\t\t\t\tdocRankMap[v[i].ID] += 1.0 / float64(i+60)\n\t\t\t}\n\t\t}\n\t}\n\tdocList := make([]*schema.Document, 0, len(docMap))\n\tfor id := range docMap {\n\t\tdocList = append(docList, docMap[id])\n\t}\n\n\tsort.Slice(docList, func(i, j int) bool {\n\t\treturn docRankMap[docList[i].ID] > docRankMap[docList[j].ID]\n\t})\n\n\treturn docList, nil\n}\n\n// NewRetriever creates a router retriever.\n// router retriever is useful when you want to retrieve documents from multiple retrievers with different queries.\n// eg.\n//\n//\trouterRetriever := router.NewRetriever(ctx, &router.Config{})\n//\tdocs, err := routerRetriever.Retrieve(ctx, \"how to build agent with eino\")\n//\tif err != nil {\n//\t\t...\n//\t}\n//\tprintln(docs)\nfunc NewRetriever(ctx context.Context, config *Config) (retriever.Retriever, error) {\n\tif len(config.Retrievers) == 0 {\n\t\treturn nil, fmt.Errorf(\"retrievers is empty\")\n\t}\n\n\trouter := config.Router\n\tif router == nil {\n\t\tvar retrieverSet []string\n\t\tfor k := range config.Retrievers {\n\t\t\tretrieverSet = append(retrieverSet, k)\n\t\t}\n\t\trouter = func(ctx context.Context, query string) ([]string, error) {\n\t\t\treturn retrieverSet, nil\n\t\t}\n\t}\n\n\tfusion := config.FusionFunc\n\tif fusion == nil {\n\t\tfusion = rrf\n\t}\n\n\treturn &routerRetriever{\n\t\tretrievers: config.Retrievers,\n\t\trouter:     config.Router,\n\t\tfusionFunc: fusion,\n\t}, nil\n}\n\n// Config is the config for router retriever.\ntype Config struct {\n\t// Retrievers is the retrievers to be used.\n\tRetrievers map[string]retriever.Retriever\n\t// Router is the function to route the query to the retrievers.\n\tRouter func(ctx context.Context, query string) ([]string, error)\n\t// FusionFunc is the function to fuse the documents from the retrievers.\n\tFusionFunc func(ctx context.Context, result map[string][]*schema.Document) ([]*schema.Document, error)\n}\n\ntype routerRetriever struct {\n\tretrievers map[string]retriever.Retriever\n\trouter     func(ctx context.Context, query string) ([]string, error)\n\tfusionFunc func(ctx context.Context, result map[string][]*schema.Document) ([]*schema.Document, error)\n}\n\n// Retrieve retrieves documents from the router retriever.\nfunc (e *routerRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n\trouteCtx := ctxWithRouterRunInfo(ctx)\n\trouteCtx = callbacks.OnStart(routeCtx, query)\n\tretrieverNames, err := e.router(routeCtx, query)\n\tif err != nil {\n\t\tcallbacks.OnError(routeCtx, err)\n\t\treturn nil, err\n\t}\n\tif len(retrieverNames) == 0 {\n\t\terr = fmt.Errorf(\"no retriever has been selected\")\n\t\tcallbacks.OnError(routeCtx, err)\n\t\treturn nil, err\n\t}\n\tcallbacks.OnEnd(routeCtx, retrieverNames)\n\n\t// retrieve\n\ttasks := make([]*utils.RetrieveTask, len(retrieverNames))\n\tfor i := range retrieverNames {\n\t\tr, ok := e.retrievers[retrieverNames[i]]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"router output[%s] has not registered\", retrieverNames[i])\n\t\t}\n\t\ttasks[i] = &utils.RetrieveTask{\n\t\t\tName:            retrieverNames[i],\n\t\t\tRetriever:       r,\n\t\t\tQuery:           query,\n\t\t\tRetrieveOptions: opts,\n\t\t}\n\t}\n\tutils.ConcurrentRetrieveWithCallback(ctx, tasks)\n\tresult := make(map[string][]*schema.Document)\n\tfor i := range tasks {\n\t\tif tasks[i].Err != nil {\n\t\t\treturn nil, tasks[i].Err\n\t\t}\n\t\tresult[tasks[i].Name] = tasks[i].Result\n\t}\n\n\t// fusion\n\tfusionCtx := ctxWithFusionRunInfo(ctx)\n\tfusionCtx = callbacks.OnStart(fusionCtx, result)\n\tfusionDocs, err := e.fusionFunc(fusionCtx, result)\n\tif err != nil {\n\t\tcallbacks.OnError(fusionCtx, err)\n\t\treturn nil, err\n\t}\n\tcallbacks.OnEnd(fusionCtx, fusionDocs)\n\treturn fusionDocs, nil\n}\n\n// GetType returns the type of the retriever (Router).\nfunc (e *routerRetriever) GetType() string { return \"Router\" }\n\nfunc ctxWithRouterRunInfo(ctx context.Context) context.Context {\n\trunInfo := &callbacks.RunInfo{\n\t\tComponent: compose.ComponentOfLambda,\n\t\tType:      \"Router\",\n\t}\n\n\trunInfo.Name = runInfo.Type + string(runInfo.Component)\n\n\treturn callbacks.ReuseHandlers(ctx, runInfo)\n}\n\nfunc ctxWithFusionRunInfo(ctx context.Context) context.Context {\n\trunInfo := &callbacks.RunInfo{\n\t\tComponent: compose.ComponentOfLambda,\n\t\tType:      \"FusionFunc\",\n\t}\n\n\trunInfo.Name = runInfo.Type + string(runInfo.Component)\n\n\treturn callbacks.ReuseHandlers(ctx, runInfo)\n}\n"
  },
  {
    "path": "flow/retriever/router/router_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage router\n\nimport (\n\t\"context\"\n\t\"reflect\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype mockRetriever struct {\n}\n\nfunc (m *mockRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n\tvar ret []*schema.Document\n\tif strings.Contains(query, \"1\") {\n\t\tret = append(ret, &schema.Document{ID: \"1\"})\n\t}\n\tif strings.Contains(query, \"2\") {\n\t\tret = append(ret, &schema.Document{ID: \"2\"})\n\t}\n\tif strings.Contains(query, \"3\") {\n\t\tret = append(ret, &schema.Document{ID: \"3\"})\n\t}\n\tif strings.Contains(query, \"4\") {\n\t\tret = append(ret, &schema.Document{ID: \"4\"})\n\t}\n\tif strings.Contains(query, \"5\") {\n\t\tret = append(ret, &schema.Document{ID: \"5\"})\n\t}\n\treturn ret, nil\n}\n\nfunc (m *mockRetriever) GetType() string {\n\treturn \"Mock\"\n}\n\nfunc TestRouterRetriever(t *testing.T) {\n\tctx := context.Background()\n\tr, err := NewRetriever(ctx, &Config{\n\t\tRetrievers: map[string]retriever.Retriever{\n\t\t\t\"1\": &mockRetriever{},\n\t\t\t\"2\": &mockRetriever{},\n\t\t\t\"3\": &mockRetriever{},\n\t\t},\n\t\tRouter: func(ctx context.Context, query string) ([]string, error) {\n\t\t\treturn []string{\"2\", \"3\"}, nil\n\t\t},\n\t\tFusionFunc: func(ctx context.Context, result map[string][]*schema.Document) ([]*schema.Document, error) {\n\t\t\tvar ret []*schema.Document\n\t\t\tfor _, v := range result {\n\t\t\t\tret = append(ret, v...)\n\t\t\t}\n\t\t\treturn ret, nil\n\t\t},\n\t})\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\thandler := callbacks.NewHandlerBuilder().\n\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\tswitch info.Name {\n\t\t\tcase \"FusionFuncLambda\":\n\t\t\t\tif _, ok := output.([]*schema.Document); !ok {\n\t\t\t\t\tt.Fatal(\"FusionFuncLambda output is not a []*schema.Document\")\n\t\t\t\t}\n\t\t\tcase \"RouterLambda\":\n\t\t\t\tif _, ok := output.([]string); !ok {\n\t\t\t\t\tt.Fatal(\"RouterLambda output is not a []string\")\n\t\t\t\t}\n\t\t\tcase \"MockRetriever\":\n\t\t\t\tif _, ok := output.([]*schema.Document); !ok {\n\t\t\t\t\tt.Fatal(\"MockRetriever output is not a []string\")\n\t\t\t\t}\n\t\t\tdefault:\n\t\t\t\tt.Fatalf(\"unknown name: %s\", info.Name)\n\t\t\t}\n\t\t\treturn ctx\n\t\t}).\n\t\tOnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\tt.Fatal(err)\n\t\t\treturn ctx\n\t\t}).Build()\n\tctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{}, handler)\n\tresult, err := r.Retrieve(ctx, \"3\")\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif len(result) != 2 {\n\t\tt.Fatal(\"expected 2 results\")\n\t}\n}\n\nfunc TestRRF(t *testing.T) {\n\tdoc1 := &schema.Document{ID: \"1\"}\n\tdoc2 := &schema.Document{ID: \"2\"}\n\tdoc3 := &schema.Document{ID: \"3\"}\n\tdoc4 := &schema.Document{ID: \"4\"}\n\tdoc5 := &schema.Document{ID: \"5\"}\n\n\tinput := map[string][]*schema.Document{\n\t\t\"1\": {doc1, doc2, doc3, doc4, doc5},\n\t\t\"2\": {doc2, doc3, doc4, doc5, doc1},\n\t\t\"3\": {doc3, doc4, doc5, doc1, doc2},\n\t}\n\n\tresult, err := rrf(context.Background(), input)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif !reflect.DeepEqual(result, []*schema.Document{doc3, doc2, doc4, doc1, doc5}) {\n\t\tt.Fatal(\"rrf fail\")\n\t}\n}\n"
  },
  {
    "path": "flow/retriever/utils/utils.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package utils provides helper utilities for retriever flows, including\n// concurrent retrieval with callback instrumentation.\npackage utils\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"sync\"\n\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// RetrieveTask is a task for retrieving documents.\n// RetrieveTask represents a single retrieval job with its result or error.\ntype RetrieveTask struct {\n\tName            string\n\tRetriever       retriever.Retriever\n\tQuery           string\n\tRetrieveOptions []retriever.Option\n\tResult          []*schema.Document\n\tErr             error\n}\n\n// ConcurrentRetrieveWithCallback concurrently retrieves documents with callback.\nfunc ConcurrentRetrieveWithCallback(ctx context.Context, tasks []*RetrieveTask) {\n\twg := sync.WaitGroup{}\n\tfor i := range tasks {\n\t\twg.Add(1)\n\t\tgo func(ctx context.Context, t *RetrieveTask) {\n\t\t\tctx = ctxWithRetrieverRunInfo(ctx, t.Retriever)\n\n\t\t\tdefer func() {\n\t\t\t\tif e := recover(); e != nil {\n\t\t\t\t\tt.Err = fmt.Errorf(\"retrieve panic, query: %s, error: %v\", t.Query, e)\n\t\t\t\t\tctx = callbacks.OnError(ctx, t.Err)\n\t\t\t\t}\n\t\t\t\twg.Done()\n\t\t\t}()\n\n\t\t\tctx = callbacks.OnStart(ctx, t.Query)\n\t\t\tdocs, err := t.Retriever.Retrieve(ctx, t.Query, t.RetrieveOptions...)\n\t\t\tif err != nil {\n\t\t\t\tcallbacks.OnError(ctx, err)\n\t\t\t\tt.Err = err\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\tcallbacks.OnEnd(ctx, docs)\n\t\t\tt.Result = docs\n\t\t}(ctx, tasks[i])\n\t}\n\twg.Wait()\n}\n\nfunc ctxWithRetrieverRunInfo(ctx context.Context, r retriever.Retriever) context.Context {\n\trunInfo := &callbacks.RunInfo{\n\t\tComponent: components.ComponentOfRetriever,\n\t}\n\n\tif typ, okk := components.GetType(r); okk {\n\t\trunInfo.Type = typ\n\t}\n\n\trunInfo.Name = runInfo.Type + string(runInfo.Component)\n\n\treturn callbacks.ReuseHandlers(ctx, runInfo)\n}\n"
  },
  {
    "path": "go.mod",
    "content": "module github.com/cloudwego/eino\n\ngo 1.18\n\nrequire (\n\tgithub.com/bmatcuk/doublestar/v4 v4.10.0\n\tgithub.com/bytedance/sonic v1.15.0\n\tgithub.com/eino-contrib/jsonschema v1.0.3\n\tgithub.com/google/uuid v1.6.0\n\tgithub.com/nikolalohinski/gonja v1.5.3\n\tgithub.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f\n\tgithub.com/smartystreets/goconvey v1.8.1\n\tgithub.com/stretchr/testify v1.10.0\n\tgithub.com/wk8/go-ordered-map/v2 v2.1.8\n\tgo.uber.org/mock v0.4.0\n\tgopkg.in/yaml.v3 v3.0.1\n)\n\nrequire (\n\tgithub.com/bahlo/generic-list-go v0.2.0 // indirect\n\tgithub.com/buger/jsonparser v1.1.1 // indirect\n\tgithub.com/bytedance/gopkg v0.1.3 // indirect\n\tgithub.com/bytedance/sonic/loader v0.5.0 // indirect\n\tgithub.com/cloudwego/base64x v0.1.6 // indirect\n\tgithub.com/davecgh/go-spew v1.1.1 // indirect\n\tgithub.com/dustin/go-humanize v1.0.1 // indirect\n\tgithub.com/goph/emperror v0.17.2 // indirect\n\tgithub.com/gopherjs/gopherjs v1.17.2 // indirect\n\tgithub.com/json-iterator/go v1.1.12 // indirect\n\tgithub.com/jtolds/gls v4.20.0+incompatible // indirect\n\tgithub.com/klauspost/cpuid/v2 v2.2.9 // indirect\n\tgithub.com/mailru/easyjson v0.7.7 // indirect\n\tgithub.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect\n\tgithub.com/modern-go/reflect2 v1.0.2 // indirect\n\tgithub.com/pelletier/go-toml/v2 v2.0.9 // indirect\n\tgithub.com/pkg/errors v0.9.1 // indirect\n\tgithub.com/pmezard/go-difflib v1.0.0 // indirect\n\tgithub.com/sirupsen/logrus v1.9.3 // indirect\n\tgithub.com/smarty/assertions v1.15.0 // indirect\n\tgithub.com/twitchyliquid64/golang-asm v0.15.1 // indirect\n\tgithub.com/yargevad/filepathx v1.0.0 // indirect\n\tgolang.org/x/arch v0.11.0 // indirect\n\tgolang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect\n\tgolang.org/x/sys v0.26.0 // indirect\n\tgopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect\n)\n"
  },
  {
    "path": "go.sum",
    "content": "github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o=\ngithub.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=\ngithub.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=\ngithub.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA=\ngithub.com/bmatcuk/doublestar/v4 v4.10.0 h1:zU9WiOla1YA122oLM6i4EXvGW62DvKZVxIe6TYWexEs=\ngithub.com/bmatcuk/doublestar/v4 v4.10.0/go.mod h1:xBQ8jztBU6kakFMg+8WGxn0c6z1fTSPVIjEY1Wr7jzc=\ngithub.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=\ngithub.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=\ngithub.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=\ngithub.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8=\ngithub.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE=\ngithub.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M=\ngithub.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM=\ngithub.com/bytedance/sonic v1.15.0 h1:/PXeWFaR5ElNcVE84U0dOHjiMHQOwNIx3K4ymzh/uSE=\ngithub.com/bytedance/sonic v1.15.0/go.mod h1:tFkWrPz0/CUCLEF4ri4UkHekCIcdnkqXw9VduqpJh0k=\ngithub.com/bytedance/sonic/loader v0.5.0 h1:gXH3KVnatgY7loH5/TkeVyXPfESoqSBSBEiDd5VjlgE=\ngithub.com/bytedance/sonic/loader v0.5.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo=\ngithub.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4=\ngithub.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M=\ngithub.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU=\ngithub.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=\ngithub.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=\ngithub.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=\ngithub.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=\ngithub.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=\ngithub.com/eino-contrib/jsonschema v1.0.3 h1:2Kfsm1xlMV0ssY2nuxshS4AwbLFuqmPmzIjLVJ1Fsp0=\ngithub.com/eino-contrib/jsonschema v1.0.3/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4=\ngithub.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=\ngithub.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ=\ngithub.com/go-check/check v0.0.0-20180628173108-788fd7840127 h1:0gkP6mzaMqkmpcJYCFOLkIBwI7xFExG03bbkOkCvUPI=\ngithub.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=\ngithub.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=\ngithub.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=\ngithub.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=\ngithub.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=\ngithub.com/goph/emperror v0.17.2 h1:yLapQcmEsO0ipe9p5TaN22djm3OFV/TfM/fcYP0/J18=\ngithub.com/goph/emperror v0.17.2/go.mod h1:+ZbQ+fUNO/6FNiUo0ujtMjhgad9Xa6fQL9KhH4LNHic=\ngithub.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g=\ngithub.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k=\ngithub.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=\ngithub.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=\ngithub.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=\ngithub.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=\ngithub.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo=\ngithub.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=\ngithub.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8=\ngithub.com/klauspost/cpuid/v2 v2.2.9 h1:66ze0taIn2H33fBvCkXuv9BmCwDfafmiIVpKV9kKGuY=\ngithub.com/klauspost/cpuid/v2 v2.2.9/go.mod h1:rqkxqrZ1EhYM9G+hXH7YdowN5R5RGN6NK4QwQ3WMXF8=\ngithub.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=\ngithub.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=\ngithub.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=\ngithub.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=\ngithub.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=\ngithub.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=\ngithub.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=\ngithub.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=\ngithub.com/mattn/go-colorable v0.1.2 h1:/bC9yWikZXAL9uJdulbSfyVNIR3n3trXl+v8+1sx8mU=\ngithub.com/mattn/go-isatty v0.0.8 h1:HLtExJ+uU2HOZ+wI0Tt5DtUDrx8yhUqDcp7fYERX4CE=\ngithub.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4=\ngithub.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=\ngithub.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=\ngithub.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=\ngithub.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=\ngithub.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=\ngithub.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c=\ngithub.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4=\ngithub.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=\ngithub.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=\ngithub.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=\ngithub.com/pelletier/go-toml/v2 v2.0.9 h1:uH2qQXheeefCCkuBBSLi7jCiSmj3VRh2+Goq2N7Xxu0=\ngithub.com/pelletier/go-toml/v2 v2.0.9/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc=\ngithub.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=\ngithub.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=\ngithub.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=\ngithub.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=\ngithub.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=\ngithub.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtISeXco0L5PKQ=\ngithub.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=\ngithub.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=\ngithub.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=\ngithub.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI=\ngithub.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg=\ngithub.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY=\ngithub.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec=\ngithub.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY=\ngithub.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60=\ngithub.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=\ngithub.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=\ngithub.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=\ngithub.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=\ngithub.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=\ngithub.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=\ngithub.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=\ngithub.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=\ngithub.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=\ngithub.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=\ngithub.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=\ngithub.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=\ngithub.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=\ngithub.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=\ngithub.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=\ngithub.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=\ngithub.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=\ngithub.com/x-cray/logrus-prefixed-formatter v0.5.2 h1:00txxvfBM9muc0jiLIEAkAcIMJzfthRT6usrui8uGmg=\ngithub.com/yargevad/filepathx v1.0.0 h1:SYcT+N3tYGi+NvazubCNlvgIPbzAk7i7y2dwg3I5FYc=\ngithub.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA=\ngo.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=\ngo.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=\ngolang.org/x/arch v0.11.0 h1:KXV8WWKCXm6tRpLirl2szsO5j/oOODwZf4hATmGVNs4=\ngolang.org/x/arch v0.11.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys=\ngolang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=\ngolang.org/x/crypto v0.11.0 h1:6Ewdq3tDic1mg5xRO4milcWCfMVQhI4NkqWWvqejpuA=\ngolang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 h1:MGwJjxBy0HJshjDNfLsYO8xppfqWlA5ZT9OhtUUhTNw=\ngolang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc=\ngolang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=\ngolang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=\ngolang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=\ngolang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=\ngolang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=\ngolang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo=\ngolang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=\ngolang.org/x/term v0.10.0 h1:3R7pNqamzBraeqj/Tj8qt1aQ2HpmlC+Cx/qL/7hn4/c=\ngolang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=\ngopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=\ngopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=\ngopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=\ngopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=\ngopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=\ngopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=\ngopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=\ngopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=\ngopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=\n"
  },
  {
    "path": "internal/callbacks/inject.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc InitCallbacks(ctx context.Context, info *RunInfo, handlers ...Handler) context.Context {\n\tmgr, ok := newManager(info, handlers...)\n\tif ok {\n\t\treturn ctxWithManager(ctx, mgr)\n\t}\n\n\treturn ctxWithManager(ctx, nil)\n}\n\nfunc EnsureRunInfo(ctx context.Context, typ string, comp components.Component) context.Context {\n\tcbm, ok := managerFromCtx(ctx)\n\tif !ok {\n\t\treturn InitCallbacks(ctx, &RunInfo{\n\t\t\tType:      typ,\n\t\t\tComponent: comp,\n\t\t})\n\t}\n\tif cbm.runInfo == nil {\n\t\treturn ReuseHandlers(ctx, &RunInfo{\n\t\t\tType:      typ,\n\t\t\tComponent: comp,\n\t\t})\n\t}\n\treturn ctx\n}\n\nfunc ReuseHandlers(ctx context.Context, info *RunInfo) context.Context {\n\tcbm, ok := managerFromCtx(ctx)\n\tif !ok {\n\t\treturn InitCallbacks(ctx, info)\n\t}\n\treturn ctxWithManager(ctx, cbm.withRunInfo(info))\n}\n\nfunc AppendHandlers(ctx context.Context, info *RunInfo, handlers ...Handler) context.Context {\n\tcbm, ok := managerFromCtx(ctx)\n\tif !ok {\n\t\treturn InitCallbacks(ctx, info, handlers...)\n\t}\n\tnh := make([]Handler, len(cbm.handlers)+len(handlers))\n\tcopy(nh[:len(cbm.handlers)], cbm.handlers)\n\tcopy(nh[len(cbm.handlers):], handlers)\n\treturn InitCallbacks(ctx, info, nh...)\n}\n\ntype Handle[T any] func(context.Context, T, *RunInfo, []Handler) (context.Context, T)\n\nfunc On[T any](ctx context.Context, inOut T, handle Handle[T], timing CallbackTiming, start bool) (context.Context, T) {\n\tmgr, ok := managerFromCtx(ctx)\n\tif !ok {\n\t\treturn ctx, inOut\n\t}\n\tnMgr := *mgr\n\n\tvar info *RunInfo\n\tif start {\n\t\tinfo = nMgr.runInfo\n\t\tnMgr.runInfo = nil\n\t\tctx = context.WithValue(ctx, CtxRunInfoKey{}, info)\n\t} else {\n\t\tif nMgr.runInfo != nil {\n\t\t\tinfo = nMgr.runInfo\n\t\t} else {\n\t\t\tinfo, _ = ctx.Value(CtxRunInfoKey{}).(*RunInfo)\n\t\t}\n\t}\n\n\ths := make([]Handler, 0, len(nMgr.handlers)+len(nMgr.globalHandlers))\n\tfor _, handler := range append(nMgr.handlers, nMgr.globalHandlers...) {\n\t\ttimingChecker, ok_ := handler.(TimingChecker)\n\t\tif !ok_ || timingChecker.Needed(ctx, info, timing) {\n\t\t\ths = append(hs, handler)\n\t\t}\n\t}\n\n\tvar out T\n\tctx, out = handle(ctx, inOut, info, hs)\n\treturn ctxWithManager(ctx, &nMgr), out\n}\n\nfunc OnStartHandle[T any](ctx context.Context, input T,\n\trunInfo *RunInfo, handlers []Handler) (context.Context, T) {\n\n\tfor i := len(handlers) - 1; i >= 0; i-- {\n\t\tctx = handlers[i].OnStart(ctx, runInfo, input)\n\t}\n\n\treturn ctx, input\n}\n\nfunc OnEndHandle[T any](ctx context.Context, output T,\n\trunInfo *RunInfo, handlers []Handler) (context.Context, T) {\n\n\tfor _, handler := range handlers {\n\t\tctx = handler.OnEnd(ctx, runInfo, output)\n\t}\n\n\treturn ctx, output\n}\n\nfunc BuildOnEndHandleWithCopy[T any](copyFn func(T, int) []T) Handle[T] {\n\treturn func(ctx context.Context, output T, runInfo *RunInfo, handlers []Handler) (context.Context, T) {\n\t\tif len(handlers) == 0 {\n\t\t\treturn ctx, output\n\t\t}\n\n\t\tcopies := copyFn(output, len(handlers))\n\n\t\tfor i, handler := range handlers {\n\t\t\tctx = handler.OnEnd(ctx, runInfo, copies[i])\n\t\t}\n\n\t\treturn ctx, output\n\t}\n}\n\nfunc OnWithStreamHandle[S any](\n\tctx context.Context,\n\tinOut S,\n\thandlers []Handler,\n\tcpy func(int) []S,\n\thandle func(context.Context, Handler, S) context.Context) (context.Context, S) {\n\n\tif len(handlers) == 0 {\n\t\treturn ctx, inOut\n\t}\n\n\tinOuts := cpy(len(handlers) + 1)\n\n\tfor i, handler := range handlers {\n\t\tctx = handle(ctx, handler, inOuts[i])\n\t}\n\n\treturn ctx, inOuts[len(inOuts)-1]\n}\n\nfunc OnStartWithStreamInputHandle[T any](ctx context.Context, input *schema.StreamReader[T],\n\trunInfo *RunInfo, handlers []Handler) (context.Context, *schema.StreamReader[T]) {\n\n\thandlers = generic.Reverse(handlers)\n\n\tcpy := input.Copy\n\n\thandle := func(ctx context.Context, handler Handler, in *schema.StreamReader[T]) context.Context {\n\t\tin_ := schema.StreamReaderWithConvert(in, func(i T) (CallbackInput, error) {\n\t\t\treturn i, nil\n\t\t})\n\t\treturn handler.OnStartWithStreamInput(ctx, runInfo, in_)\n\t}\n\n\treturn OnWithStreamHandle(ctx, input, handlers, cpy, handle)\n}\n\nfunc OnEndWithStreamOutputHandle[T any](ctx context.Context, output *schema.StreamReader[T],\n\trunInfo *RunInfo, handlers []Handler) (context.Context, *schema.StreamReader[T]) {\n\n\tcpy := output.Copy\n\n\thandle := func(ctx context.Context, handler Handler, out *schema.StreamReader[T]) context.Context {\n\t\tout_ := schema.StreamReaderWithConvert(out, func(i T) (CallbackOutput, error) {\n\t\t\treturn i, nil\n\t\t})\n\t\treturn handler.OnEndWithStreamOutput(ctx, runInfo, out_)\n\t}\n\n\treturn OnWithStreamHandle(ctx, output, handlers, cpy, handle)\n}\n\nfunc OnErrorHandle(ctx context.Context, err error,\n\trunInfo *RunInfo, handlers []Handler) (context.Context, error) {\n\n\tfor _, handler := range handlers {\n\t\tctx = handler.OnError(ctx, runInfo, err)\n\t}\n\n\treturn ctx, err\n}\n"
  },
  {
    "path": "internal/callbacks/interface.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\ntype RunInfo struct {\n\t// Name is the graph node name for display purposes, not unique.\n\t// Passed from compose.WithNodeName().\n\tName      string\n\tType      string\n\tComponent components.Component\n}\n\ntype CallbackInput any\n\ntype CallbackOutput any\n\ntype Handler interface {\n\tOnStart(ctx context.Context, info *RunInfo, input CallbackInput) context.Context\n\tOnEnd(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context\n\n\tOnError(ctx context.Context, info *RunInfo, err error) context.Context\n\n\tOnStartWithStreamInput(ctx context.Context, info *RunInfo,\n\t\tinput *schema.StreamReader[CallbackInput]) context.Context\n\tOnEndWithStreamOutput(ctx context.Context, info *RunInfo,\n\t\toutput *schema.StreamReader[CallbackOutput]) context.Context\n}\n\ntype CallbackTiming uint8\n\ntype TimingChecker interface {\n\tNeeded(ctx context.Context, info *RunInfo, timing CallbackTiming) bool\n}\n"
  },
  {
    "path": "internal/callbacks/manager.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport \"context\"\n\ntype CtxManagerKey struct{}\ntype CtxRunInfoKey struct{}\n\ntype manager struct {\n\tglobalHandlers []Handler\n\thandlers       []Handler\n\trunInfo        *RunInfo\n}\n\nvar GlobalHandlers []Handler\n\nfunc newManager(runInfo *RunInfo, handlers ...Handler) (*manager, bool) {\n\tif len(handlers)+len(GlobalHandlers) == 0 {\n\t\treturn nil, false\n\t}\n\n\ths := make([]Handler, len(GlobalHandlers))\n\tcopy(hs, GlobalHandlers)\n\n\treturn &manager{\n\t\tglobalHandlers: hs,\n\t\thandlers:       handlers,\n\t\trunInfo:        runInfo,\n\t}, true\n}\n\nfunc ctxWithManager(ctx context.Context, manager *manager) context.Context {\n\treturn context.WithValue(ctx, CtxManagerKey{}, manager)\n}\n\nfunc (m *manager) withRunInfo(runInfo *RunInfo) *manager {\n\tif m == nil {\n\t\treturn nil\n\t}\n\n\tn := *m\n\tn.runInfo = runInfo\n\treturn &n\n}\n\nfunc managerFromCtx(ctx context.Context) (*manager, bool) {\n\tv := ctx.Value(CtxManagerKey{})\n\tm, ok := v.(*manager)\n\tif ok && m != nil {\n\t\tn := *m\n\t\treturn &n, true\n\t}\n\n\treturn nil, false\n}\n"
  },
  {
    "path": "internal/channel.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport \"sync\"\n\n// UnboundedChan represents a channel with unlimited capacity\ntype UnboundedChan[T any] struct {\n\tbuffer   []T        // Internal buffer to store data\n\tmutex    sync.Mutex // Mutex to protect buffer access\n\tnotEmpty *sync.Cond // Condition variable to wait for data\n\tclosed   bool       // Indicates if the channel has been closed\n}\n\n// NewUnboundedChan initializes and returns an UnboundedChan\nfunc NewUnboundedChan[T any]() *UnboundedChan[T] {\n\tch := &UnboundedChan[T]{}\n\tch.notEmpty = sync.NewCond(&ch.mutex)\n\treturn ch\n}\n\n// Send puts an item into the channel\nfunc (ch *UnboundedChan[T]) Send(value T) {\n\tch.mutex.Lock()\n\tdefer ch.mutex.Unlock()\n\n\tif ch.closed {\n\t\tpanic(\"send on closed channel\")\n\t}\n\n\tch.buffer = append(ch.buffer, value)\n\tch.notEmpty.Signal() // Wake up one goroutine waiting to receive\n}\n\n// Receive gets an item from the channel (blocks if empty)\nfunc (ch *UnboundedChan[T]) Receive() (T, bool) {\n\tch.mutex.Lock()\n\tdefer ch.mutex.Unlock()\n\n\tfor len(ch.buffer) == 0 && !ch.closed {\n\t\tch.notEmpty.Wait() // Wait until data is available\n\t}\n\n\tif len(ch.buffer) == 0 {\n\t\t// Channel is closed and empty\n\t\tvar zero T\n\t\treturn zero, false\n\t}\n\n\tval := ch.buffer[0]\n\tch.buffer = ch.buffer[1:]\n\treturn val, true\n}\n\n// Close marks the channel as closed\nfunc (ch *UnboundedChan[T]) Close() {\n\tch.mutex.Lock()\n\tdefer ch.mutex.Unlock()\n\n\tif !ch.closed {\n\t\tch.closed = true\n\t\tch.notEmpty.Broadcast() // Wake up all waiting goroutines\n\t}\n}\n"
  },
  {
    "path": "internal/channel_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport (\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n)\n\nfunc TestUnboundedChan_Send(t *testing.T) {\n\tch := NewUnboundedChan[string]()\n\n\t// Test sending a value\n\tch.Send(\"test\")\n\tif len(ch.buffer) != 1 {\n\t\tt.Errorf(\"buffer length should be 1, got %d\", len(ch.buffer))\n\t}\n\tif ch.buffer[0] != \"test\" {\n\t\tt.Errorf(\"expected 'test', got '%s'\", ch.buffer[0])\n\t}\n\n\t// Test sending multiple values\n\tch.Send(\"test2\")\n\tch.Send(\"test3\")\n\tif len(ch.buffer) != 3 {\n\t\tt.Errorf(\"buffer length should be 3, got %d\", len(ch.buffer))\n\t}\n}\n\nfunc TestUnboundedChan_SendPanic(t *testing.T) {\n\tch := NewUnboundedChan[int]()\n\tch.Close()\n\n\t// Test sending to closed channel should panic\n\tdefer func() {\n\t\tif r := recover(); r == nil {\n\t\t\tt.Error(\"sending to closed channel should panic\")\n\t\t}\n\t}()\n\n\tch.Send(1)\n}\n\nfunc TestUnboundedChan_Receive(t *testing.T) {\n\tch := NewUnboundedChan[int]()\n\n\t// Send values\n\tch.Send(1)\n\tch.Send(2)\n\n\t// Test receiving values\n\tval, ok := ch.Receive()\n\tif !ok {\n\t\tt.Error(\"receive should succeed\")\n\t}\n\tif val != 1 {\n\t\tt.Errorf(\"expected 1, got %d\", val)\n\t}\n\n\tval, ok = ch.Receive()\n\tif !ok {\n\t\tt.Error(\"receive should succeed\")\n\t}\n\tif val != 2 {\n\t\tt.Errorf(\"expected 2, got %d\", val)\n\t}\n}\n\nfunc TestUnboundedChan_ReceiveFromClosed(t *testing.T) {\n\tch := NewUnboundedChan[int]()\n\tch.Close()\n\n\t// Test receiving from closed, empty channel\n\tval, ok := ch.Receive()\n\tif ok {\n\t\tt.Error(\"receive from closed, empty channel should return ok=false\")\n\t}\n\tif val != 0 {\n\t\tt.Errorf(\"expected zero value, got %d\", val)\n\t}\n\n\t// Test receiving from closed channel with values\n\tch = NewUnboundedChan[int]()\n\tch.Send(42)\n\tch.Close()\n\n\tval, ok = ch.Receive()\n\tif !ok {\n\t\tt.Error(\"receive should succeed\")\n\t}\n\tif val != 42 {\n\t\tt.Errorf(\"expected 42, got %d\", val)\n\t}\n\n\t// After consuming all values\n\tval, ok = ch.Receive()\n\tif ok {\n\t\tt.Error(\"receive from closed, empty channel should return ok=false\")\n\t}\n}\n\nfunc TestUnboundedChan_Close(t *testing.T) {\n\tch := NewUnboundedChan[int]()\n\n\t// Test closing\n\tch.Close()\n\tif !ch.closed {\n\t\tt.Error(\"channel should be marked as closed\")\n\t}\n\n\t// Test double closing (should not panic)\n\tch.Close()\n}\n\nfunc TestUnboundedChan_Concurrency(t *testing.T) {\n\tch := NewUnboundedChan[int]()\n\tconst numSenders = 5\n\tconst numReceivers = 3\n\tconst messagesPerSender = 100\n\n\tvar rwg, swg sync.WaitGroup\n\trwg.Add(numReceivers)\n\tswg.Add(numSenders)\n\n\t// Start senders\n\tfor i := 0; i < numSenders; i++ {\n\t\tgo func(id int) {\n\t\t\tdefer swg.Done()\n\t\t\tfor j := 0; j < messagesPerSender; j++ {\n\t\t\t\tch.Send(id*messagesPerSender + j)\n\t\t\t\ttime.Sleep(time.Microsecond) // Small delay to increase concurrency chance\n\t\t\t}\n\t\t}(i)\n\t}\n\n\t// Start receivers\n\treceived := make([]int, 0, numSenders*messagesPerSender)\n\tvar mu sync.Mutex\n\n\tfor i := 0; i < numReceivers; i++ {\n\t\tgo func() {\n\t\t\tdefer rwg.Done()\n\t\t\tfor {\n\t\t\t\tval, ok := ch.Receive()\n\t\t\t\tif !ok {\n\t\t\t\t\treturn\n\t\t\t\t}\n\t\t\t\tmu.Lock()\n\t\t\t\treceived = append(received, val)\n\t\t\t\tmu.Unlock()\n\t\t\t}\n\t\t}()\n\t}\n\n\t// Wait for senders to finish\n\tswg.Wait()\n\tch.Close()\n\n\t// Wait for all goroutines to finish\n\trwg.Wait()\n\n\t// Verify we received all messages\n\tif len(received) != numSenders*messagesPerSender {\n\t\tt.Errorf(\"expected %d messages, got %d\", numSenders*messagesPerSender, len(received))\n\t}\n\n\t// Create a map to check for duplicates and missing values\n\treceivedMap := make(map[int]bool)\n\tfor _, val := range received {\n\t\treceivedMap[val] = true\n\t}\n\n\tif len(receivedMap) != numSenders*messagesPerSender {\n\t\tt.Error(\"duplicate or missing messages detected\")\n\t}\n}\n\nfunc TestUnboundedChan_BlockingReceive(t *testing.T) {\n\tch := NewUnboundedChan[int]()\n\n\t// Test that Receive blocks when channel is empty\n\treceiveDone := make(chan bool)\n\tgo func() {\n\t\tch.Receive()\n\t\treceiveDone <- true\n\t}()\n\n\t// Check that receive is blocked\n\tselect {\n\tcase <-receiveDone:\n\t\tt.Error(\"Receive should block on empty channel\")\n\tcase <-time.After(50 * time.Millisecond):\n\t\t// This is expected\n\t}\n\n\t// Send a value to unblock\n\tch.Send(1)\n\n\t// Now receive should complete\n\tselect {\n\tcase <-receiveDone:\n\t\t// This is expected\n\tcase <-time.After(50 * time.Millisecond):\n\t\tt.Error(\"Receive should have unblocked\")\n\t}\n}\n"
  },
  {
    "path": "internal/concat.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport (\n\t\"fmt\"\n\t\"reflect\"\n\t\"strings\"\n\t\"time\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\nvar (\n\tconcatFuncs = map[reflect.Type]any{\n\t\tgeneric.TypeOf[string]():        concatStrings,\n\t\tgeneric.TypeOf[int8]():          useLast[int8],\n\t\tgeneric.TypeOf[int16]():         useLast[int16],\n\t\tgeneric.TypeOf[int32]():         useLast[int32],\n\t\tgeneric.TypeOf[int64]():         useLast[int64],\n\t\tgeneric.TypeOf[int]():           useLast[int],\n\t\tgeneric.TypeOf[uint8]():         useLast[uint8],\n\t\tgeneric.TypeOf[uint16]():        useLast[uint16],\n\t\tgeneric.TypeOf[uint32]():        useLast[uint32],\n\t\tgeneric.TypeOf[uint64]():        useLast[uint64],\n\t\tgeneric.TypeOf[uint]():          useLast[uint],\n\t\tgeneric.TypeOf[bool]():          useLast[bool],\n\t\tgeneric.TypeOf[float32]():       useLast[float32],\n\t\tgeneric.TypeOf[float64]():       useLast[float64],\n\t\tgeneric.TypeOf[time.Time]():     useLast[time.Time],\n\t\tgeneric.TypeOf[time.Duration](): useLast[time.Duration],\n\t}\n)\n\nfunc useLast[T any](s []T) (T, error) {\n\treturn s[len(s)-1], nil\n}\n\nfunc concatStrings(ss []string) (string, error) {\n\tvar n int\n\tfor _, s := range ss {\n\t\tn += len(s)\n\t}\n\n\tvar b strings.Builder\n\tb.Grow(n)\n\tfor _, s := range ss {\n\t\t_, err := b.WriteString(s)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t}\n\n\treturn b.String(), nil\n}\n\nfunc RegisterStreamChunkConcatFunc[T any](fn func([]T) (T, error)) {\n\tconcatFuncs[generic.TypeOf[T]()] = fn\n}\n\nfunc GetConcatFunc(typ reflect.Type) func(reflect.Value) (reflect.Value, error) {\n\tif fn, ok := concatFuncs[typ]; ok {\n\t\treturn func(a reflect.Value) (reflect.Value, error) {\n\t\t\trvs := reflect.ValueOf(fn).Call([]reflect.Value{a})\n\t\t\tvar err error\n\t\t\tif !rvs[1].IsNil() {\n\t\t\t\terr = rvs[1].Interface().(error)\n\t\t\t}\n\t\t\treturn rvs[0], err\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// ConcatItems the caller should ensure len(items) > 1\nfunc ConcatItems[T any](items []T) (T, error) {\n\ttyp := generic.TypeOf[T]()\n\tv := reflect.ValueOf(items)\n\n\tvar cv reflect.Value\n\tvar err error\n\n\t// handle map kind\n\tif typ.Kind() == reflect.Map {\n\t\tcv, err = concatMaps(v)\n\t} else {\n\t\tcv, err = concatSliceValue(v)\n\t}\n\n\tif err != nil {\n\t\tvar t T\n\t\treturn t, err\n\t}\n\n\treturn cv.Interface().(T), nil\n}\n\nfunc concatMaps(ms reflect.Value) (reflect.Value, error) {\n\ttyp := ms.Type().Elem()\n\n\trms := reflect.MakeMap(reflect.MapOf(typ.Key(), generic.TypeOf[[]any]()))\n\tret := reflect.MakeMap(typ)\n\n\tn := ms.Len()\n\tfor i := 0; i < n; i++ {\n\t\tm := ms.Index(i)\n\n\t\tfor _, key := range m.MapKeys() {\n\t\t\tvals := rms.MapIndex(key)\n\t\t\tif !vals.IsValid() {\n\t\t\t\tvar s []any\n\t\t\t\tvals = reflect.ValueOf(s)\n\t\t\t}\n\n\t\t\tval := m.MapIndex(key)\n\t\t\tvals = reflect.Append(vals, val)\n\t\t\trms.SetMapIndex(key, vals)\n\t\t}\n\t}\n\n\tfor _, key := range rms.MapKeys() {\n\t\tvals := rms.MapIndex(key)\n\n\t\tanyVals := vals.Interface().([]any)\n\t\tif len(anyVals) == 1 {\n\t\t\tele := anyVals[0]\n\t\t\tif ele == nil { // we cannot SetMapIndex with nil because it will delete the key\n\t\t\t\tret.SetMapIndex(key, reflect.Zero(typ.Elem()))\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tret.SetMapIndex(key, reflect.ValueOf(ele))\n\t\t\tcontinue\n\t\t}\n\n\t\tv, err := toSliceValue(anyVals)\n\t\tif err != nil {\n\t\t\treturn reflect.Value{}, err\n\t\t}\n\n\t\tvar cv reflect.Value\n\n\t\tif v.Type().Elem().Kind() == reflect.Map {\n\t\t\tcv, err = concatMaps(v)\n\t\t} else {\n\t\t\tcv, err = concatSliceValue(v)\n\t\t}\n\n\t\tif err != nil {\n\t\t\treturn reflect.Value{}, err\n\t\t}\n\n\t\tret.SetMapIndex(key, cv)\n\t}\n\n\treturn ret, nil\n}\n\nfunc concatSliceValue(val reflect.Value) (reflect.Value, error) {\n\telmType := val.Type().Elem()\n\n\tif val.Len() == 1 {\n\t\treturn val.Index(0), nil\n\t}\n\n\tf := GetConcatFunc(elmType)\n\tif f != nil {\n\t\treturn f(val)\n\t}\n\n\t// if all elements in the slice are empty, return an empty value\n\t// if there is exactly one non-empty element in the slice, return that non-empty element\n\t// otherwise, throw an error.\n\tvar filtered reflect.Value\n\tfor i := 0; i < val.Len(); i++ {\n\t\toneVal := val.Index(i)\n\t\tif !oneVal.IsZero() {\n\t\t\tif filtered.IsValid() {\n\t\t\t\treturn reflect.Value{}, fmt.Errorf(\"cannot concat multiple non-zero value of type %s\", elmType)\n\t\t\t}\n\n\t\t\tfiltered = oneVal\n\t\t}\n\t}\n\tif !filtered.IsValid() {\n\t\tfiltered = reflect.New(elmType).Elem()\n\t}\n\n\treturn filtered, nil\n}\n\nfunc toSliceValue(vs []any) (reflect.Value, error) {\n\ttyp := reflect.TypeOf(vs[0])\n\n\tret := reflect.MakeSlice(reflect.SliceOf(typ), len(vs), len(vs))\n\tret.Index(0).Set(reflect.ValueOf(vs[0]))\n\n\tfor i := 1; i < len(vs); i++ {\n\t\tv := vs[i]\n\t\tvt := reflect.TypeOf(v)\n\t\tif typ != vt {\n\t\t\treturn reflect.Value{}, fmt.Errorf(\"unexpected slice element type. Got %v, expected %v\", typ, vt)\n\t\t}\n\n\t\tret.Index(i).Set(reflect.ValueOf(v))\n\t}\n\n\treturn ret, nil\n}\n"
  },
  {
    "path": "internal/concat_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestConcat(t *testing.T) {\n\tt.Run(\"concat map chunks with nil value\", func(t *testing.T) {\n\t\tc1 := map[string]any{\n\t\t\t\"a\": map[string]any{\n\t\t\t\t\"b\": map[string]any{\n\t\t\t\t\t\"c1\": nil,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tc2 := map[string]any{\n\t\t\t\"a\": map[string]any{\n\t\t\t\t\"b\": map[string]any{\n\t\t\t\t\t\"c2\": \"c2\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tm, err := ConcatItems([]map[string]any{c1, c2})\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, map[string]any{\n\t\t\t\"a\": map[string]any{\n\t\t\t\t\"b\": map[string]any{\n\t\t\t\t\t\"c1\": nil,\n\t\t\t\t\t\"c2\": \"c2\",\n\t\t\t\t},\n\t\t\t},\n\t\t}, m)\n\t})\n}\n"
  },
  {
    "path": "internal/core/address.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage core\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\t\"sync\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\n// AddressSegmentType defines the type of a segment in an execution address.\ntype AddressSegmentType string\n\n// Address represents a full, hierarchical address to a point in the execution structure.\ntype Address []AddressSegment\n\n// String converts an Address into its unique string representation.\nfunc (p Address) String() string {\n\tif p == nil {\n\t\treturn \"\"\n\t}\n\tvar sb strings.Builder\n\tfor i, s := range p {\n\t\tsb.WriteString(string(s.Type))\n\t\tsb.WriteString(\":\")\n\t\tsb.WriteString(s.ID)\n\t\tif s.SubID != \"\" {\n\t\t\tsb.WriteString(\":\")\n\t\t\tsb.WriteString(s.SubID)\n\t\t}\n\t\tif i != len(p)-1 {\n\t\t\tsb.WriteString(\";\")\n\t\t}\n\t}\n\treturn sb.String()\n}\n\nfunc (p Address) Equals(other Address) bool {\n\tif len(p) != len(other) {\n\t\treturn false\n\t}\n\tfor i := range p {\n\t\tif p[i].Type != other[i].Type || p[i].ID != other[i].ID || p[i].SubID != other[i].SubID {\n\t\t\treturn false\n\t\t}\n\t}\n\treturn true\n}\n\n// AddressSegment represents a single segment in the hierarchical address of an execution point.\n// A sequence of AddressSegments uniquely identifies a location within a potentially nested structure.\ntype AddressSegment struct {\n\t// ID is the unique identifier for this segment, e.g., the node's key or the tool's name.\n\tID string\n\t// Type indicates whether this address segment is a graph node, a tool call, an agent, etc.\n\tType AddressSegmentType\n\t// In some cases, ID alone are not unique enough, we need this SubID to guarantee uniqueness.\n\t// e.g. parallel tool calls with the same name but different tool call IDs.\n\tSubID string\n}\n\ntype addrCtxKey struct{}\n\ntype addrCtx struct {\n\taddr           Address\n\tinterruptState *InterruptState\n\tisResumeTarget bool\n\tresumeData     any\n}\n\ntype globalResumeInfoKey struct{}\n\ntype globalResumeInfo struct {\n\tmu                sync.Mutex\n\tid2ResumeData     map[string]any\n\tid2ResumeDataUsed map[string]bool\n\tid2State          map[string]InterruptState\n\tid2StateUsed      map[string]bool\n\tid2Addr           map[string]Address\n}\n\n// GetCurrentAddress returns the hierarchical address of the currently executing component.\n// The address is a sequence of segments, each identifying a structural part of the execution\n// like an agent, a graph node, or a tool call. This can be useful for logging or debugging.\nfunc GetCurrentAddress(ctx context.Context) Address {\n\tif p, ok := ctx.Value(addrCtxKey{}).(*addrCtx); ok {\n\t\treturn p.addr\n\t}\n\n\treturn nil\n}\n\n// AppendAddressSegment creates a new execution context for a sub-component (e.g., a graph node or a tool call).\n//\n// It extends the current context's address with a new segment and populates the new context with the\n// appropriate interrupt state and resume data for that specific sub-address.\n//\n//   - ctx: The parent context, typically the one passed into the component's Invoke/Stream method.\n//   - segType: The type of the new address segment (e.g., \"node\", \"tool\").\n//   - segID: The unique ID for the new address segment.\nfunc AppendAddressSegment(ctx context.Context, segType AddressSegmentType, segID string,\n\tsubID string) context.Context {\n\t// get current address\n\tcurrentAddress := GetCurrentAddress(ctx)\n\tif len(currentAddress) == 0 {\n\t\tcurrentAddress = []AddressSegment{\n\t\t\t{\n\t\t\t\tType:  segType,\n\t\t\t\tID:    segID,\n\t\t\t\tSubID: subID,\n\t\t\t},\n\t\t}\n\t} else {\n\t\tnewAddress := make([]AddressSegment, len(currentAddress)+1)\n\t\tcopy(newAddress, currentAddress)\n\t\tnewAddress[len(newAddress)-1] = AddressSegment{\n\t\t\tType:  segType,\n\t\t\tID:    segID,\n\t\t\tSubID: subID,\n\t\t}\n\t\tcurrentAddress = newAddress\n\t}\n\n\trunCtx := &addrCtx{\n\t\taddr: currentAddress,\n\t}\n\n\trInfo, hasRInfo := getResumeInfo(ctx)\n\tif !hasRInfo {\n\t\treturn context.WithValue(ctx, addrCtxKey{}, runCtx)\n\t}\n\n\tvar id string\n\tfor id_, addr := range rInfo.id2Addr {\n\t\tif addr.Equals(currentAddress) {\n\t\t\trInfo.mu.Lock()\n\t\t\tif used, ok := rInfo.id2StateUsed[id_]; !ok || !used {\n\t\t\t\trunCtx.interruptState = generic.PtrOf(rInfo.id2State[id_])\n\t\t\t\trInfo.id2StateUsed[id_] = true\n\t\t\t\tid = id_\n\t\t\t\trInfo.mu.Unlock()\n\t\t\t\tbreak\n\t\t\t}\n\t\t\trInfo.mu.Unlock()\n\t\t}\n\t}\n\n\t// take from globalResumeInfo the data for the new address if there is any\n\trInfo.mu.Lock()\n\tdefer rInfo.mu.Unlock()\n\tused := rInfo.id2ResumeDataUsed[id]\n\tif !used {\n\t\trData, existed := rInfo.id2ResumeData[id]\n\t\tif existed {\n\t\t\trInfo.id2ResumeDataUsed[id] = true\n\t\t\trunCtx.resumeData = rData\n\t\t\trunCtx.isResumeTarget = true\n\t\t}\n\t}\n\n\t// Also mark as resume target if any descendant address is a resume target.\n\t// This allows composite components (e.g., a tool containing a nested graph) to know\n\t// they should execute their children to reach the actual resume target.\n\t// We only consider descendants whose resume data has not yet been consumed.\n\tif !runCtx.isResumeTarget {\n\t\tfor id_, addr := range rInfo.id2Addr {\n\t\t\tif len(addr) > len(currentAddress) && addr[:len(currentAddress)].Equals(currentAddress) {\n\t\t\t\tif !rInfo.id2ResumeDataUsed[id_] {\n\t\t\t\t\trunCtx.isResumeTarget = true\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\treturn context.WithValue(ctx, addrCtxKey{}, runCtx)\n}\n\n// GetNextResumptionPoints finds the immediate child resumption points for a given parent address.\nfunc GetNextResumptionPoints(ctx context.Context) (map[string]bool, error) {\n\tparentAddr := GetCurrentAddress(ctx)\n\n\trInfo, exists := getResumeInfo(ctx)\n\tif !exists {\n\t\treturn nil, fmt.Errorf(\"GetNextResumptionPoints: failed to get resume info from context\")\n\t}\n\n\tnextPoints := make(map[string]bool)\n\tparentAddrLen := len(parentAddr)\n\n\tfor _, addr := range rInfo.id2Addr {\n\t\t// Check if addr is a potential child (must be longer than parent)\n\t\tif len(addr) <= parentAddrLen {\n\t\t\tcontinue\n\t\t}\n\n\t\t// Check if it has the parent address as a prefix\n\t\tvar isPrefix bool\n\t\tif parentAddrLen == 0 {\n\t\t\tisPrefix = true\n\t\t} else {\n\t\t\tisPrefix = addr[:parentAddrLen].Equals(parentAddr)\n\t\t}\n\n\t\tif !isPrefix {\n\t\t\tcontinue\n\t\t}\n\n\t\t// We are looking for immediate children.\n\t\t// The address of an immediate child should be one segment longer.\n\t\tchildAddr := addr[parentAddrLen : parentAddrLen+1]\n\t\tchildID := childAddr[0].ID\n\n\t\t// Avoid adding duplicates.\n\t\tif _, ok := nextPoints[childID]; !ok {\n\t\t\tnextPoints[childID] = true\n\t\t}\n\t}\n\n\treturn nextPoints, nil\n}\n\n// BatchResumeWithData is the core function for preparing a resume context. It injects a map\n// of resume targets and their corresponding data into the context.\n//\n// The `resumeData` map should contain the interrupt IDs (which are the string form of addresses) of the\n// components to be resumed as keys. The value can be the resume data for that component, or `nil`\n// if no data is needed (equivalent to using `Resume`).\n//\n// This function is the foundation for the \"Explicit Targeted Resume\" strategy. Components whose interrupt IDs\n// are present as keys in the map will receive `isResumeFlow = true` when they call `GetResumeContext`.\nfunc BatchResumeWithData(ctx context.Context, resumeData map[string]any) context.Context {\n\trInfo, ok := ctx.Value(globalResumeInfoKey{}).(*globalResumeInfo)\n\tif !ok {\n\t\t// Create a new globalResumeInfo and copy the map to prevent external mutation.\n\t\tnewMap := make(map[string]any, len(resumeData))\n\t\tfor k, v := range resumeData {\n\t\t\tnewMap[k] = v\n\t\t}\n\t\treturn context.WithValue(ctx, globalResumeInfoKey{}, &globalResumeInfo{\n\t\t\tid2ResumeData:     newMap,\n\t\t\tid2ResumeDataUsed: make(map[string]bool),\n\t\t\tid2StateUsed:      make(map[string]bool),\n\t\t})\n\t}\n\n\trInfo.mu.Lock()\n\tdefer rInfo.mu.Unlock()\n\tif rInfo.id2ResumeData == nil {\n\t\trInfo.id2ResumeData = make(map[string]any)\n\t}\n\tfor id, data := range resumeData {\n\t\trInfo.id2ResumeData[id] = data\n\t}\n\treturn ctx\n}\n\nfunc PopulateInterruptState(ctx context.Context, id2Addr map[string]Address,\n\tid2State map[string]InterruptState) context.Context {\n\trInfo, ok := ctx.Value(globalResumeInfoKey{}).(*globalResumeInfo)\n\tif ok {\n\t\tif rInfo.id2Addr == nil {\n\t\t\trInfo.id2Addr = make(map[string]Address)\n\t\t}\n\t\tfor id, addr := range id2Addr {\n\t\t\trInfo.id2Addr[id] = addr\n\t\t}\n\t\trInfo.id2State = id2State\n\t} else {\n\t\trInfo = &globalResumeInfo{\n\t\t\tid2Addr:           id2Addr,\n\t\t\tid2State:          id2State,\n\t\t\tid2StateUsed:      make(map[string]bool),\n\t\t\tid2ResumeDataUsed: make(map[string]bool),\n\t\t}\n\t\tctx = context.WithValue(ctx, globalResumeInfoKey{}, rInfo)\n\t}\n\n\trunCtx, ok := getRunCtx(ctx)\n\tif ok {\n\t\tfor id_, addr := range id2Addr {\n\t\t\tif addr.Equals(runCtx.addr) {\n\t\t\t\tif used, ok := rInfo.id2StateUsed[id_]; !ok || !used {\n\t\t\t\t\trunCtx.interruptState = generic.PtrOf(rInfo.id2State[id_])\n\t\t\t\t\trInfo.mu.Lock()\n\t\t\t\t\trInfo.id2StateUsed[id_] = true\n\t\t\t\t\trInfo.mu.Unlock()\n\t\t\t\t}\n\n\t\t\t\tif used, ok := rInfo.id2ResumeDataUsed[id_]; !ok || !used {\n\t\t\t\t\trunCtx.isResumeTarget = true\n\t\t\t\t\trunCtx.resumeData = rInfo.id2ResumeData[id_]\n\t\t\t\t\trInfo.mu.Lock()\n\t\t\t\t\trInfo.id2ResumeDataUsed[id_] = true\n\t\t\t\t\trInfo.mu.Unlock()\n\t\t\t\t}\n\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\n\treturn ctx\n}\n\nfunc getResumeInfo(ctx context.Context) (*globalResumeInfo, bool) {\n\tinfo, ok := ctx.Value(globalResumeInfoKey{}).(*globalResumeInfo)\n\treturn info, ok\n}\n\ntype InterruptInfo struct {\n\tInfo        any\n\tIsRootCause bool\n}\n\nfunc (i *InterruptInfo) String() string {\n\tif i == nil {\n\t\treturn \"\"\n\t}\n\treturn fmt.Sprintf(\"interrupt info: Info=%v, IsRootCause=%v\", i.Info, i.IsRootCause)\n}\n"
  },
  {
    "path": "internal/core/interrupt.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage core\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/google/uuid\"\n)\n\ntype CheckPointStore interface {\n\tGet(ctx context.Context, checkPointID string) ([]byte, bool, error)\n\tSet(ctx context.Context, checkPointID string, checkPoint []byte) error\n}\n\ntype InterruptSignal struct {\n\tID string\n\tAddress\n\tInterruptInfo\n\tInterruptState\n\tSubs []*InterruptSignal\n}\n\nfunc (is *InterruptSignal) Error() string {\n\treturn fmt.Sprintf(\"interrupt signal: ID=%s, Addr=%s, Info=%s, State=%s, SubsLen=%d\",\n\t\tis.ID, is.Address.String(), is.InterruptInfo.String(), is.InterruptState.String(), len(is.Subs))\n}\n\ntype InterruptState struct {\n\tState                any\n\tLayerSpecificPayload any\n}\n\nfunc (is *InterruptState) String() string {\n\tif is == nil {\n\t\treturn \"\"\n\t}\n\treturn fmt.Sprintf(\"interrupt state: State=%v, LayerSpecificPayload=%v\", is.State, is.LayerSpecificPayload)\n}\n\n// InterruptConfig holds optional parameters for creating an interrupt.\ntype InterruptConfig struct {\n\tLayerPayload any\n}\n\n// InterruptOption is a function that configures an InterruptConfig.\ntype InterruptOption func(*InterruptConfig)\n\n// WithLayerPayload creates an option to attach layer-specific metadata\n// to the interrupt's state.\nfunc WithLayerPayload(payload any) InterruptOption {\n\treturn func(c *InterruptConfig) {\n\t\tc.LayerPayload = payload\n\t}\n}\n\nfunc Interrupt(ctx context.Context, info any, state any, subContexts []*InterruptSignal, opts ...InterruptOption) (\n\t*InterruptSignal, error) {\n\taddr := GetCurrentAddress(ctx)\n\n\t// Apply options to get config\n\tconfig := &InterruptConfig{}\n\tfor _, opt := range opts {\n\t\topt(config)\n\t}\n\n\tmyPoint := InterruptInfo{\n\t\tInfo: info,\n\t}\n\n\tif len(subContexts) == 0 {\n\t\tmyPoint.IsRootCause = true\n\t\treturn &InterruptSignal{\n\t\t\tID:            uuid.NewString(),\n\t\t\tAddress:       addr,\n\t\t\tInterruptInfo: myPoint,\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState:                state,\n\t\t\t\tLayerSpecificPayload: config.LayerPayload,\n\t\t\t},\n\t\t}, nil\n\t}\n\n\treturn &InterruptSignal{\n\t\tID:            uuid.NewString(),\n\t\tAddress:       addr,\n\t\tInterruptInfo: myPoint,\n\t\tInterruptState: InterruptState{\n\t\t\tState:                state,\n\t\t\tLayerSpecificPayload: config.LayerPayload,\n\t\t},\n\t\tSubs: subContexts,\n\t}, nil\n}\n\n// InterruptCtx provides a complete, user-facing context for a single, resumable interrupt point.\ntype InterruptCtx struct {\n\t// ID is the unique, fully-qualified address of the interrupt point.\n\t// It is constructed by joining the individual Address segments, e.g., \"agent:A;node:graph_a;tool:tool_call_123\".\n\t// This ID should be used when providing resume data via ResumeWithData.\n\tID string\n\t// Address is the structured sequence of AddressSegment segments that leads to the interrupt point.\n\tAddress Address\n\t// Info is the user-facing information associated with the interrupt, provided by the component that triggered it.\n\tInfo any\n\t// IsRootCause indicates whether the interrupt point is the exact root cause for an interruption.\n\tIsRootCause bool\n\t// Parent points to the context of the parent component in the interrupt chain.\n\t// It is nil for the top-level interrupt.\n\tParent *InterruptCtx\n}\n\nfunc (ic *InterruptCtx) EqualsWithoutID(other *InterruptCtx) bool {\n\tif ic == nil && other == nil {\n\t\treturn true\n\t}\n\n\tif ic == nil || other == nil {\n\t\treturn false\n\t}\n\n\tif !ic.Address.Equals(other.Address) {\n\t\treturn false\n\t}\n\n\tif ic.IsRootCause != other.IsRootCause {\n\t\treturn false\n\t}\n\n\tif ic.Info != nil || other.Info != nil {\n\t\tif ic.Info == nil || other.Info == nil {\n\t\t\treturn false\n\t\t}\n\n\t\tif !reflect.DeepEqual(ic.Info, other.Info) {\n\t\t\treturn false\n\t\t}\n\t}\n\n\tif ic.Parent != nil || other.Parent != nil {\n\t\tif ic.Parent == nil || other.Parent == nil {\n\t\t\treturn false\n\t\t}\n\n\t\tif !ic.Parent.EqualsWithoutID(other.Parent) {\n\t\t\treturn false\n\t\t}\n\t}\n\n\treturn true\n}\n\n// InterruptContextsProvider is an interface for errors that contain interrupt contexts.\n// This allows different packages to check for and extract interrupt contexts from errors\n// without needing to know the concrete error type.\ntype InterruptContextsProvider interface {\n\tGetInterruptContexts() []*InterruptCtx\n}\n\n// FromInterruptContexts converts a list of user-facing InterruptCtx objects into an\n// internal InterruptSignal tree. It correctly handles common ancestors and ensures\n// that the resulting tree is consistent with the original interrupt chain.\n//\n// This method is primarily used by components that bridge different execution environments.\n// For example, an `adk.AgentTool` might catch an `adk.InterruptInfo`, extract the\n// `adk.InterruptCtx` objects from it, and then call this method on each one. The resulting\n// error signals are then typically aggregated into a single error using `compose.CompositeInterrupt`\n// to be returned from the tool's `InvokableRun` method.\n// FromInterruptContexts reconstructs a single InterruptSignal tree from a list of\n// user-facing InterruptCtx objects. It correctly merges common ancestors.\nfunc FromInterruptContexts(contexts []*InterruptCtx) *InterruptSignal {\n\tif len(contexts) == 0 {\n\t\treturn nil\n\t}\n\n\tsignalMap := make(map[string]*InterruptSignal)\n\tvar rootSignal *InterruptSignal\n\n\t// getOrCreateSignal is a recursive helper that builds the tree bottom-up.\n\tvar getOrCreateSignal func(*InterruptCtx) *InterruptSignal\n\tgetOrCreateSignal = func(ctx *InterruptCtx) *InterruptSignal {\n\t\tif ctx == nil {\n\t\t\treturn nil\n\t\t}\n\t\t// If we've already created a signal for this context, return it.\n\t\tif signal, exists := signalMap[ctx.ID]; exists {\n\t\t\treturn signal\n\t\t}\n\n\t\t// Create the signal for the current context.\n\t\tnewSignal := &InterruptSignal{\n\t\t\tID:      ctx.ID,\n\t\t\tAddress: ctx.Address,\n\t\t\tInterruptInfo: InterruptInfo{\n\t\t\t\tInfo:        ctx.Info,\n\t\t\t\tIsRootCause: ctx.IsRootCause,\n\t\t\t},\n\t\t}\n\t\tsignalMap[ctx.ID] = newSignal // Cache it immediately.\n\n\t\t// Recursively ensure the parent exists. If it doesn't, this is the root.\n\t\tif parentSignal := getOrCreateSignal(ctx.Parent); parentSignal != nil {\n\t\t\tparentSignal.Subs = append(parentSignal.Subs, newSignal)\n\t\t} else {\n\t\t\trootSignal = newSignal\n\t\t}\n\t\treturn newSignal\n\t}\n\n\t// Process all contexts to ensure all branches of the tree are built.\n\tfor _, ctx := range contexts {\n\t\t_ = getOrCreateSignal(ctx)\n\t}\n\n\treturn rootSignal\n}\n\n// ToInterruptContexts converts the internal InterruptSignal tree into a list of\n// user-facing InterruptCtx objects for the root causes of the interruption.\n// Each returned context has its Parent field populated (if it has a parent),\n// allowing traversal up the interrupt chain.\n//\n// If allowedSegmentTypes is nil, all segment types are kept and addresses are unchanged.\n// If allowedSegmentTypes is provided, it:\n//  1. Filters the parent chain to only keep contexts whose leaf segment type is allowed\n//  2. Strips non-allowed segment types from all addresses\nfunc ToInterruptContexts(is *InterruptSignal, allowedSegmentTypes []AddressSegmentType) []*InterruptCtx {\n\tif is == nil {\n\t\treturn nil\n\t}\n\tvar rootCauseContexts []*InterruptCtx\n\n\tvar buildContexts func(*InterruptSignal, *InterruptCtx)\n\tbuildContexts = func(signal *InterruptSignal, parentCtx *InterruptCtx) {\n\t\tcurrentCtx := &InterruptCtx{\n\t\t\tID:          signal.ID,\n\t\t\tAddress:     signal.Address,\n\t\t\tInfo:        signal.InterruptInfo.Info,\n\t\t\tIsRootCause: signal.InterruptInfo.IsRootCause,\n\t\t\tParent:      parentCtx,\n\t\t}\n\n\t\tif currentCtx.IsRootCause {\n\t\t\trootCauseContexts = append(rootCauseContexts, currentCtx)\n\t\t}\n\n\t\tfor _, subSignal := range signal.Subs {\n\t\t\tbuildContexts(subSignal, currentCtx)\n\t\t}\n\t}\n\n\tbuildContexts(is, nil)\n\n\tif len(allowedSegmentTypes) > 0 {\n\t\tallowedSet := make(map[AddressSegmentType]bool, len(allowedSegmentTypes))\n\t\tfor _, t := range allowedSegmentTypes {\n\t\t\tallowedSet[t] = true\n\t\t}\n\n\t\tfor _, ctx := range rootCauseContexts {\n\t\t\tfilterParentChain(ctx, allowedSet)\n\t\t\tencapsulateContextAddresses(ctx, allowedSet)\n\t\t}\n\t}\n\n\treturn rootCauseContexts\n}\n\nfunc filterParentChain(ctx *InterruptCtx, allowedSet map[AddressSegmentType]bool) {\n\tif ctx == nil {\n\t\treturn\n\t}\n\n\tparent := ctx.Parent\n\tfor parent != nil {\n\t\tif len(parent.Address) > 0 && allowedSet[parent.Address[len(parent.Address)-1].Type] {\n\t\t\tbreak\n\t\t}\n\t\tparent = parent.Parent\n\t}\n\n\tctx.Parent = parent\n\n\tfilterParentChain(parent, allowedSet)\n}\n\nfunc encapsulateContextAddresses(ctx *InterruptCtx, allowedSet map[AddressSegmentType]bool) {\n\tfor c := ctx; c != nil; c = c.Parent {\n\t\tnewAddr := make(Address, 0, len(c.Address))\n\t\tfor _, seg := range c.Address {\n\t\t\tif allowedSet[seg.Type] {\n\t\t\t\tnewAddr = append(newAddr, seg)\n\t\t\t}\n\t\t}\n\t\tc.Address = newAddr\n\t}\n}\n\n// SignalToPersistenceMaps flattens an InterruptSignal tree into two maps suitable for persistence in a checkpoint.\nfunc SignalToPersistenceMaps(is *InterruptSignal) (map[string]Address, map[string]InterruptState) {\n\tid2addr := make(map[string]Address)\n\tid2state := make(map[string]InterruptState)\n\n\tif is == nil {\n\t\treturn id2addr, id2state\n\t}\n\n\tvar traverse func(*InterruptSignal)\n\ttraverse = func(signal *InterruptSignal) {\n\t\t// Add current signal's data to the maps.\n\t\tid2addr[signal.ID] = signal.Address\n\t\tid2state[signal.ID] = signal.InterruptState // The embedded struct\n\n\t\t// Recurse into children.\n\t\tfor _, sub := range signal.Subs {\n\t\t\ttraverse(sub)\n\t\t}\n\t}\n\n\ttraverse(is)\n\treturn id2addr, id2state\n}\n"
  },
  {
    "path": "internal/core/interrupt_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage core\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\n// Define AddressSegmentType constants locally to avoid dependency cycles\nconst (\n\tAddressSegmentAgent AddressSegmentType = \"agent\"\n\tAddressSegmentTool  AddressSegmentType = \"tool\"\n\tAddressSegmentNode  AddressSegmentType = \"node\"\n)\n\nfunc TestInterruptConversion(t *testing.T) {\n\t// Test Case 1: Simple Chain (A -> B -> C)\n\tt.Run(\"SimpleChain\", func(t *testing.T) {\n\t\t// Manually construct the user-facing contexts with parent pointers\n\t\tctxA := &InterruptCtx{ID: \"A\", IsRootCause: false}\n\t\tctxB := &InterruptCtx{ID: \"B\", Parent: ctxA, IsRootCause: false}\n\t\tctxC := &InterruptCtx{ID: \"C\", Parent: ctxB, IsRootCause: true}\n\n\t\t// The input to FromInterruptContexts is just the root cause leaf node\n\t\tcontexts := []*InterruptCtx{ctxC}\n\n\t\t// Convert from user-facing contexts to internal signal tree\n\t\tsignal := FromInterruptContexts(contexts)\n\n\t\t// Assertions for the signal tree structure\n\t\tassert.NotNil(t, signal)\n\t\tassert.Equal(t, \"A\", signal.ID)\n\t\tassert.Len(t, signal.Subs, 1)\n\t\tassert.Equal(t, \"B\", signal.Subs[0].ID)\n\t\tassert.Len(t, signal.Subs[0].Subs, 1)\n\t\tassert.Equal(t, \"C\", signal.Subs[0].Subs[0].ID)\n\t\tassert.True(t, signal.Subs[0].Subs[0].IsRootCause)\n\n\t\t// Convert back from the signal tree to user-facing contexts\n\t\tfinalContexts := ToInterruptContexts(signal, nil)\n\n\t\t// Assertions for the final list of contexts\n\t\tassert.Len(t, finalContexts, 1)\n\t\tfinalC := finalContexts[0]\n\t\tassert.Equal(t, \"C\", finalC.ID)\n\t\tassert.True(t, finalC.IsRootCause)\n\t\tassert.NotNil(t, finalC.Parent)\n\t\tassert.Equal(t, \"B\", finalC.Parent.ID)\n\t\tassert.NotNil(t, finalC.Parent.Parent)\n\t\tassert.Equal(t, \"A\", finalC.Parent.Parent.ID)\n\t\tassert.Nil(t, finalC.Parent.Parent.Parent)\n\t})\n\n\t// Test Case 2: Multiple Root Causes with Shared Parent (B -> D, C -> D)\n\tt.Run(\"MultipleRootsSharedParent\", func(t *testing.T) {\n\t\t// Manually construct the contexts\n\t\tctxD := &InterruptCtx{ID: \"D\", IsRootCause: false}\n\t\tctxB := &InterruptCtx{ID: \"B\", Parent: ctxD, IsRootCause: true}\n\t\tctxC := &InterruptCtx{ID: \"C\", Parent: ctxD, IsRootCause: true}\n\n\t\t// The input contains both root cause leaves\n\t\tcontexts := []*InterruptCtx{ctxB, ctxC}\n\n\t\t// Convert to signal tree\n\t\tsignal := FromInterruptContexts(contexts)\n\n\t\t// Assertions for the signal tree structure (should merge at D)\n\t\tassert.NotNil(t, signal)\n\t\tassert.Equal(t, \"D\", signal.ID)\n\t\tassert.Len(t, signal.Subs, 2)\n\t\t// Order of subs is not guaranteed, so we check for presence\n\t\tsubIDs := []string{signal.Subs[0].ID, signal.Subs[1].ID}\n\t\tassert.Contains(t, subIDs, \"B\")\n\t\tassert.Contains(t, subIDs, \"C\")\n\n\t\t// Convert back to user-facing contexts\n\t\tfinalContexts := ToInterruptContexts(signal, nil)\n\n\t\t// Assertions for the final list of contexts\n\t\tassert.Len(t, finalContexts, 2)\n\t\tfinalIDs := []string{finalContexts[0].ID, finalContexts[1].ID}\n\t\tassert.Contains(t, finalIDs, \"B\")\n\t\tassert.Contains(t, finalIDs, \"C\")\n\n\t\t// Check parent linking for one of the branches\n\t\tvar finalB *InterruptCtx\n\t\tif finalContexts[0].ID == \"B\" {\n\t\t\tfinalB = finalContexts[0]\n\t\t} else {\n\t\t\tfinalB = finalContexts[1]\n\t\t}\n\t\tassert.NotNil(t, finalB.Parent)\n\t\tassert.Equal(t, \"D\", finalB.Parent.ID)\n\t\tassert.Nil(t, finalB.Parent.Parent)\n\t})\n\n\t// Test Case 3: Nil and Empty Inputs\n\tt.Run(\"NilAndEmpty\", func(t *testing.T) {\n\t\tassert.Nil(t, FromInterruptContexts(nil))\n\t\tassert.Nil(t, FromInterruptContexts([]*InterruptCtx{}))\n\t\tassert.Nil(t, ToInterruptContexts(nil, nil))\n\t})\n}\n\nfunc TestSignalToPersistenceMaps(t *testing.T) {\n\t// Test Case 1: Nil Signal\n\tt.Run(\"NilSignal\", func(t *testing.T) {\n\t\tid2addr, id2state := SignalToPersistenceMaps(nil)\n\t\tassert.NotNil(t, id2addr)\n\t\tassert.NotNil(t, id2state)\n\t\tassert.Empty(t, id2addr)\n\t\tassert.Empty(t, id2state)\n\t})\n\n\t// Test Case 2: Single Node Signal\n\tt.Run(\"SingleNode\", func(t *testing.T) {\n\t\tsignal := &InterruptSignal{\n\t\t\tID: \"node1\",\n\t\t\tAddress: Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState:                \"test state\",\n\t\t\t\tLayerSpecificPayload: \"test payload\",\n\t\t\t},\n\t\t}\n\n\t\tid2addr, id2state := SignalToPersistenceMaps(signal)\n\n\t\tassert.Len(t, id2addr, 1)\n\t\tassert.Len(t, id2state, 1)\n\n\t\tassert.Equal(t, signal.Address, id2addr[\"node1\"])\n\t\tassert.Equal(t, signal.InterruptState, id2state[\"node1\"])\n\t})\n\n\t// Test Case 3: Simple Tree Structure\n\tt.Run(\"SimpleTree\", func(t *testing.T) {\n\t\tchild1 := &InterruptSignal{\n\t\t\tID: \"child1\",\n\t\t\tAddress: Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState: \"child1 state\",\n\t\t\t},\n\t\t}\n\n\t\tchild2 := &InterruptSignal{\n\t\t\tID: \"child2\",\n\t\t\tAddress: Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t\t{Type: AddressSegmentTool, ID: \"tool2\"},\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState: \"child2 state\",\n\t\t\t},\n\t\t}\n\n\t\tparent := &InterruptSignal{\n\t\t\tID: \"parent\",\n\t\t\tAddress: Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState: \"parent state\",\n\t\t\t},\n\t\t\tSubs: []*InterruptSignal{child1, child2},\n\t\t}\n\n\t\tid2addr, id2state := SignalToPersistenceMaps(parent)\n\n\t\t// Should contain all 3 nodes\n\t\tassert.Len(t, id2addr, 3)\n\t\tassert.Len(t, id2state, 3)\n\n\t\t// Check parent node\n\t\tassert.Equal(t, parent.Address, id2addr[\"parent\"])\n\t\tassert.Equal(t, parent.InterruptState, id2state[\"parent\"])\n\n\t\t// Check child nodes\n\t\tassert.Equal(t, child1.Address, id2addr[\"child1\"])\n\t\tassert.Equal(t, child1.InterruptState, id2state[\"child1\"])\n\t\tassert.Equal(t, child2.Address, id2addr[\"child2\"])\n\t\tassert.Equal(t, child2.InterruptState, id2state[\"child2\"])\n\t})\n\n\t// Test Case 4: Deeply Nested Tree\n\tt.Run(\"DeeplyNestedTree\", func(t *testing.T) {\n\t\tleaf1 := &InterruptSignal{\n\t\t\tID: \"leaf1\",\n\t\t\tAddress: Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t\t\t{Type: AddressSegmentNode, ID: \"node1\"},\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState: \"leaf1 state\",\n\t\t\t},\n\t\t}\n\n\t\tleaf2 := &InterruptSignal{\n\t\t\tID: \"leaf2\",\n\t\t\tAddress: Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t\t\t{Type: AddressSegmentNode, ID: \"node2\"},\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState: \"leaf2 state\",\n\t\t\t},\n\t\t}\n\n\t\tmiddle := &InterruptSignal{\n\t\t\tID: \"middle\",\n\t\t\tAddress: Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState: \"middle state\",\n\t\t\t},\n\t\t\tSubs: []*InterruptSignal{leaf1, leaf2},\n\t\t}\n\n\t\troot := &InterruptSignal{\n\t\t\tID: \"root\",\n\t\t\tAddress: Address{\n\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState: \"root state\",\n\t\t\t},\n\t\t\tSubs: []*InterruptSignal{middle},\n\t\t}\n\n\t\tid2addr, id2state := SignalToPersistenceMaps(root)\n\n\t\t// Should contain all 4 nodes\n\t\tassert.Len(t, id2addr, 4)\n\t\tassert.Len(t, id2state, 4)\n\n\t\t// Verify all nodes are present\n\t\tassert.Equal(t, root.Address, id2addr[\"root\"])\n\t\tassert.Equal(t, root.InterruptState, id2state[\"root\"])\n\t\tassert.Equal(t, middle.Address, id2addr[\"middle\"])\n\t\tassert.Equal(t, middle.InterruptState, id2state[\"middle\"])\n\t\tassert.Equal(t, leaf1.Address, id2addr[\"leaf1\"])\n\t\tassert.Equal(t, leaf1.InterruptState, id2state[\"leaf1\"])\n\t\tassert.Equal(t, leaf2.Address, id2addr[\"leaf2\"])\n\t\tassert.Equal(t, leaf2.InterruptState, id2state[\"leaf2\"])\n\t})\n\n\t// Test Case 5: Complex Tree with Multiple Branches\n\tt.Run(\"ComplexTree\", func(t *testing.T) {\n\t\t// Create a complex tree structure with multiple branches\n\t\tbranch1Leaf1 := &InterruptSignal{ID: \"b1l1\", Address: Address{{Type: AddressSegmentAgent, ID: \"a1\"}}, InterruptState: InterruptState{State: \"b1l1\"}}\n\t\tbranch1Leaf2 := &InterruptSignal{ID: \"b1l2\", Address: Address{{Type: AddressSegmentAgent, ID: \"a1\"}}, InterruptState: InterruptState{State: \"b1l2\"}}\n\t\tbranch1 := &InterruptSignal{ID: \"b1\", Address: Address{{Type: AddressSegmentAgent, ID: \"a1\"}}, InterruptState: InterruptState{State: \"b1\"}, Subs: []*InterruptSignal{branch1Leaf1, branch1Leaf2}}\n\n\t\tbranch2Leaf1 := &InterruptSignal{ID: \"b2l1\", Address: Address{{Type: AddressSegmentAgent, ID: \"a1\"}}, InterruptState: InterruptState{State: \"b2l1\"}}\n\t\tbranch2 := &InterruptSignal{ID: \"b2\", Address: Address{{Type: AddressSegmentAgent, ID: \"a1\"}}, InterruptState: InterruptState{State: \"b2\"}, Subs: []*InterruptSignal{branch2Leaf1}}\n\n\t\troot := &InterruptSignal{ID: \"root\", Address: Address{{Type: AddressSegmentAgent, ID: \"a1\"}}, InterruptState: InterruptState{State: \"root\"}, Subs: []*InterruptSignal{branch1, branch2}}\n\n\t\tid2addr, id2state := SignalToPersistenceMaps(root)\n\n\t\t// Should contain all 6 nodes\n\t\tassert.Len(t, id2addr, 6)\n\t\tassert.Len(t, id2state, 6)\n\n\t\t// Verify all nodes are present\n\t\texpectedNodes := []string{\"root\", \"b1\", \"b2\", \"b1l1\", \"b1l2\", \"b2l1\"}\n\t\tfor _, nodeID := range expectedNodes {\n\t\t\tassert.Contains(t, id2addr, nodeID)\n\t\t\tassert.Contains(t, id2state, nodeID)\n\t\t}\n\t})\n\n\t// Test Case 6: Empty InterruptState Values\n\tt.Run(\"EmptyInterruptState\", func(t *testing.T) {\n\t\tsignal := &InterruptSignal{\n\t\t\tID:             \"node1\",\n\t\t\tAddress:        Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\t// Empty state values\n\t\t\t},\n\t\t}\n\n\t\tid2addr, id2state := SignalToPersistenceMaps(signal)\n\n\t\tassert.Len(t, id2addr, 1)\n\t\tassert.Len(t, id2state, 1)\n\t\tassert.Equal(t, signal.Address, id2addr[\"node1\"])\n\t\tassert.Equal(t, signal.InterruptState, id2state[\"node1\"])\n\t})\n}\n\nfunc TestGetCurrentAddress(t *testing.T) {\n\t// Test Case 1: No Address in Context\n\tt.Run(\"NoAddressInContext\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\taddr := GetCurrentAddress(ctx)\n\t\tassert.Nil(t, addr)\n\t})\n\n\t// Test Case 2: Address in Context\n\tt.Run(\"AddressInContext\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\texpectedAddr := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t}\n\n\t\t// Create a context with address using internal addrCtx\n\t\trunCtx := &addrCtx{\n\t\t\taddr: expectedAddr,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\taddr := GetCurrentAddress(ctx)\n\t\tassert.Equal(t, expectedAddr, addr)\n\t})\n}\n\nfunc TestGetNextResumptionPoints(t *testing.T) {\n\t// Test Case 1: No Resume Info in Context\n\tt.Run(\"NoResumeInfo\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\t_, err := GetNextResumptionPoints(ctx)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"failed to get resume info\")\n\t})\n\n\t// Test Case 2: Empty Resume Info\n\tt.Run(\"EmptyResumeInfo\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\trInfo := &globalResumeInfo{\n\t\t\tid2Addr: make(map[string]Address),\n\t\t}\n\t\tctx = context.WithValue(ctx, globalResumeInfoKey{}, rInfo)\n\n\t\tpoints, err := GetNextResumptionPoints(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, points)\n\t})\n\n\t// Test Case 3: Valid Resume Points\n\tt.Run(\"ValidResumePoints\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Set up current address\n\t\tcurrentAddr := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t}\n\t\trunCtx := &addrCtx{\n\t\t\taddr: currentAddr,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\t// Set up resume info with child addresses\n\t\trInfo := &globalResumeInfo{\n\t\t\tid2Addr: map[string]Address{\n\t\t\t\t\"child1\": {\n\t\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t\t\t},\n\t\t\t\t\"child2\": {\n\t\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t\t\t{Type: AddressSegmentTool, ID: \"tool2\"},\n\t\t\t\t},\n\t\t\t\t\"unrelated\": {\n\t\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent2\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tctx = context.WithValue(ctx, globalResumeInfoKey{}, rInfo)\n\n\t\tpoints, err := GetNextResumptionPoints(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, points, 2)\n\t\tassert.True(t, points[\"tool1\"])\n\t\tassert.True(t, points[\"tool2\"])\n\t})\n\n\t// Test Case 4: Root Address (Empty Parent)\n\tt.Run(\"RootAddress\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Empty current address (root)\n\t\trunCtx := &addrCtx{\n\t\t\taddr: Address{},\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\t// Set up resume info with various addresses\n\t\trInfo := &globalResumeInfo{\n\t\t\tid2Addr: map[string]Address{\n\t\t\t\t\"agent1\": {\n\t\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t\t},\n\t\t\t\t\"agent2\": {\n\t\t\t\t\t{Type: AddressSegmentAgent, ID: \"agent2\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tctx = context.WithValue(ctx, globalResumeInfoKey{}, rInfo)\n\n\t\tpoints, err := GetNextResumptionPoints(ctx)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, points, 2)\n\t\tassert.True(t, points[\"agent1\"])\n\t\tassert.True(t, points[\"agent2\"])\n\t})\n}\n\nfunc TestBatchResumeWithData(t *testing.T) {\n\t// Test Case 1: New Resume Data\n\tt.Run(\"NewResumeData\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tresumeData := map[string]any{\n\t\t\t\"id1\": \"data1\",\n\t\t\t\"id2\": \"data2\",\n\t\t}\n\n\t\tnewCtx := BatchResumeWithData(ctx, resumeData)\n\n\t\t// Verify the data was set correctly\n\t\trInfo, ok := newCtx.Value(globalResumeInfoKey{}).(*globalResumeInfo)\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, rInfo)\n\t\tassert.Equal(t, \"data1\", rInfo.id2ResumeData[\"id1\"])\n\t\tassert.Equal(t, \"data2\", rInfo.id2ResumeData[\"id2\"])\n\t})\n\n\t// Test Case 2: Merge with Existing Resume Data\n\tt.Run(\"MergeWithExisting\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// First call with initial data\n\t\tinitialData := map[string]any{\n\t\t\t\"id1\": \"initial\",\n\t\t}\n\t\tctx = BatchResumeWithData(ctx, initialData)\n\n\t\t// Second call with additional data\n\t\tadditionalData := map[string]any{\n\t\t\t\"id2\": \"additional\",\n\t\t}\n\t\tnewCtx := BatchResumeWithData(ctx, additionalData)\n\n\t\t// Verify both data sets are present\n\t\trInfo, ok := newCtx.Value(globalResumeInfoKey{}).(*globalResumeInfo)\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, rInfo)\n\t\tassert.Equal(t, \"initial\", rInfo.id2ResumeData[\"id1\"])\n\t\tassert.Equal(t, \"additional\", rInfo.id2ResumeData[\"id2\"])\n\t})\n\n\t// Test Case 3: Empty Resume Data\n\tt.Run(\"EmptyResumeData\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tnewCtx := BatchResumeWithData(ctx, map[string]any{})\n\n\t\trInfo, ok := newCtx.Value(globalResumeInfoKey{}).(*globalResumeInfo)\n\t\tassert.True(t, ok)\n\t\tassert.NotNil(t, rInfo)\n\t\tassert.Empty(t, rInfo.id2ResumeData)\n\t})\n}\n\nfunc TestGetInterruptState(t *testing.T) {\n\t// Test Case 1: No Interrupt State\n\tt.Run(\"NoInterruptState\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\twasInterrupted, hasState, state := GetInterruptState[string](ctx)\n\t\tassert.False(t, wasInterrupted)\n\t\tassert.False(t, hasState)\n\t\tassert.Equal(t, \"\", state)\n\t})\n\n\t// Test Case 2: With Interrupt State\n\tt.Run(\"WithInterruptState\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with interrupt state\n\t\trunCtx := &addrCtx{\n\t\t\tinterruptState: &InterruptState{\n\t\t\t\tState: \"test state\",\n\t\t\t},\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\twasInterrupted, hasState, state := GetInterruptState[string](ctx)\n\t\tassert.True(t, wasInterrupted)\n\t\tassert.True(t, hasState)\n\t\tassert.Equal(t, \"test state\", state)\n\t})\n\n\t// Test Case 3: Wrong Type for Interrupt State\n\tt.Run(\"WrongType\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with interrupt state of wrong type\n\t\trunCtx := &addrCtx{\n\t\t\tinterruptState: &InterruptState{\n\t\t\t\tState: 123, // int instead of string\n\t\t\t},\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\twasInterrupted, hasState, state := GetInterruptState[string](ctx)\n\t\tassert.True(t, wasInterrupted)\n\t\tassert.False(t, hasState) // Should be false due to type mismatch\n\t\tassert.Equal(t, \"\", state)\n\t})\n\n\t// Test Case 4: Nil Interrupt State\n\tt.Run(\"NilInterruptState\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with nil interrupt state\n\t\trunCtx := &addrCtx{\n\t\t\tinterruptState: nil,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\twasInterrupted, hasState, state := GetInterruptState[string](ctx)\n\t\tassert.False(t, wasInterrupted) // Should be false because interruptState is nil\n\t\tassert.False(t, hasState)       // Should be false because state is nil\n\t\tassert.Equal(t, \"\", state)\n\t})\n}\n\nfunc TestGetResumeContext(t *testing.T) {\n\t// Test Case 1: Not Resume Target\n\tt.Run(\"NotResumeTarget\", func(t *testing.T) {\n\t\tctx := context.Background()\n\t\tisResumeTarget, hasData, data := GetResumeContext[string](ctx)\n\t\tassert.False(t, isResumeTarget)\n\t\tassert.False(t, hasData)\n\t\tassert.Equal(t, \"\", data)\n\t})\n\n\t// Test Case 2: Resume Target with Data\n\tt.Run(\"ResumeTargetWithData\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context as resume target with data\n\t\trunCtx := &addrCtx{\n\t\t\tisResumeTarget: true,\n\t\t\tresumeData:     \"resume data\",\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\tisResumeTarget, hasData, data := GetResumeContext[string](ctx)\n\t\tassert.True(t, isResumeTarget)\n\t\tassert.True(t, hasData)\n\t\tassert.Equal(t, \"resume data\", data)\n\t})\n\n\t// Test Case 3: Resume Target without Data\n\tt.Run(\"ResumeTargetWithoutData\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context as resume target without data\n\t\trunCtx := &addrCtx{\n\t\t\tisResumeTarget: true,\n\t\t\tresumeData:     nil,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\tisResumeTarget, hasData, data := GetResumeContext[string](ctx)\n\t\tassert.True(t, isResumeTarget)\n\t\tassert.False(t, hasData)\n\t\tassert.Equal(t, \"\", data)\n\t})\n\n\t// Test Case 4: Wrong Type for Resume Data\n\tt.Run(\"WrongType\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with resume data of wrong type\n\t\trunCtx := &addrCtx{\n\t\t\tisResumeTarget: true,\n\t\t\tresumeData:     123, // int instead of string\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\tisResumeTarget, hasData, data := GetResumeContext[string](ctx)\n\t\tassert.True(t, isResumeTarget)\n\t\tassert.False(t, hasData) // Should be false due to type mismatch\n\t\tassert.Equal(t, \"\", data)\n\t})\n}\n\nfunc TestWithLayerPayload(t *testing.T) {\n\t// Test Case 1: Basic Usage\n\tt.Run(\"BasicUsage\", func(t *testing.T) {\n\t\tconfig := &InterruptConfig{}\n\t\topt := WithLayerPayload(\"test payload\")\n\t\topt(config)\n\t\tassert.Equal(t, \"test payload\", config.LayerPayload)\n\t})\n\n\t// Test Case 2: Nil Payload\n\tt.Run(\"NilPayload\", func(t *testing.T) {\n\t\tconfig := &InterruptConfig{LayerPayload: \"existing\"}\n\t\topt := WithLayerPayload(nil)\n\t\topt(config)\n\t\tassert.Nil(t, config.LayerPayload)\n\t})\n\n\t// Test Case 3: Complex Payload\n\tt.Run(\"ComplexPayload\", func(t *testing.T) {\n\t\tconfig := &InterruptConfig{}\n\t\tpayload := map[string]any{\n\t\t\t\"key1\": \"value1\",\n\t\t\t\"key2\": 123,\n\t\t}\n\t\topt := WithLayerPayload(payload)\n\t\topt(config)\n\t\tassert.Equal(t, payload, config.LayerPayload)\n\t})\n}\n\nfunc TestInterruptFunction(t *testing.T) {\n\t// Test Case 1: Simple Interrupt without SubContexts\n\tt.Run(\"SimpleInterrupt\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a mock address\n\t\texpectedAddr := Address{{Type: AddressSegmentAgent, ID: \"test-agent\"}}\n\t\trunCtx := &addrCtx{\n\t\t\taddr: expectedAddr,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\tinfo := \"test info\"\n\t\tstate := \"test state\"\n\n\t\tsignal, err := Interrupt(ctx, info, state, nil)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, signal)\n\t\tassert.NotEmpty(t, signal.ID)\n\t\tassert.Equal(t, info, signal.Info)\n\t\tassert.Equal(t, state, signal.State)\n\t\tassert.True(t, signal.IsRootCause)\n\t\tassert.Equal(t, expectedAddr, signal.Address)\n\t})\n\n\t// Test Case 2: Interrupt with SubContexts\n\tt.Run(\"InterruptWithSubContexts\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a mock address\n\t\texpectedAddr := Address{{Type: AddressSegmentAgent, ID: \"parent-agent\"}}\n\t\trunCtx := &addrCtx{\n\t\t\taddr: expectedAddr,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\t// Create sub contexts\n\t\tsubContexts := []*InterruptSignal{\n\t\t\t{\n\t\t\t\tID:      \"child1\",\n\t\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"child1\"}},\n\t\t\t},\n\t\t\t{\n\t\t\t\tID:      \"child2\",\n\t\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"child2\"}},\n\t\t\t},\n\t\t}\n\n\t\tinfo := \"parent info\"\n\t\tstate := \"parent state\"\n\n\t\tsignal, err := Interrupt(ctx, info, state, subContexts)\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, signal)\n\t\tassert.NotEmpty(t, signal.ID)\n\t\tassert.Equal(t, info, signal.Info)\n\t\tassert.Equal(t, state, signal.State)\n\t\tassert.False(t, signal.IsRootCause) // Should be false when there are sub contexts\n\t\tassert.Len(t, signal.Subs, 2)\n\t\tassert.Equal(t, \"child1\", signal.Subs[0].ID)\n\t\tassert.Equal(t, \"child2\", signal.Subs[1].ID)\n\t})\n\n\t// Test Case 3: Interrupt with Options\n\tt.Run(\"InterruptWithOptions\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a mock address\n\t\texpectedAddr := Address{{Type: AddressSegmentAgent, ID: \"test-agent\"}}\n\t\trunCtx := &addrCtx{\n\t\t\taddr: expectedAddr,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\tinfo := \"test info\"\n\t\tstate := \"test state\"\n\t\tlayerPayload := \"layer payload\"\n\n\t\tsignal, err := Interrupt(ctx, info, state, nil, WithLayerPayload(layerPayload))\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, signal)\n\t\tassert.Equal(t, layerPayload, signal.LayerSpecificPayload)\n\t})\n\n\t// Test Case 4: Empty SubContexts\n\tt.Run(\"EmptySubContexts\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Create a context with a mock address\n\t\texpectedAddr := Address{{Type: AddressSegmentAgent, ID: \"test-agent\"}}\n\t\trunCtx := &addrCtx{\n\t\t\taddr: expectedAddr,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\tinfo := \"test info\"\n\t\tstate := \"test state\"\n\n\t\tsignal, err := Interrupt(ctx, info, state, []*InterruptSignal{})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, signal)\n\t\tassert.True(t, signal.IsRootCause) // Should be true when sub contexts is empty\n\t\tassert.Empty(t, signal.Subs)\n\t})\n}\n\nfunc TestAddressMethods(t *testing.T) {\n\t// Test Case 1: Address.String()\n\tt.Run(\"AddressString\", func(t *testing.T) {\n\t\taddr := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t\t{Type: AddressSegmentNode, ID: \"node1\", SubID: \"sub1\"},\n\t\t}\n\n\t\tresult := addr.String()\n\t\texpected := \"agent:agent1;tool:tool1;node:node1:sub1\"\n\t\tassert.Equal(t, expected, result)\n\t})\n\n\t// Test Case 2: Address.String() with empty address\n\tt.Run(\"EmptyAddressString\", func(t *testing.T) {\n\t\tvar addr Address\n\t\tresult := addr.String()\n\t\tassert.Equal(t, \"\", result)\n\t})\n\n\t// Test Case 3: Address.Equals() with equal addresses\n\tt.Run(\"AddressEquals\", func(t *testing.T) {\n\t\taddr1 := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t}\n\t\taddr2 := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t}\n\n\t\tassert.True(t, addr1.Equals(addr2))\n\t})\n\n\t// Test Case 4: Address.Equals() with different addresses\n\tt.Run(\"AddressNotEquals\", func(t *testing.T) {\n\t\taddr1 := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t}\n\t\taddr2 := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t{Type: AddressSegmentTool, ID: \"tool2\"},\n\t\t}\n\n\t\tassert.False(t, addr1.Equals(addr2))\n\t})\n\n\t// Test Case 5: Address.Equals() with different lengths\n\tt.Run(\"AddressDifferentLengths\", func(t *testing.T) {\n\t\taddr1 := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t\t{Type: AddressSegmentTool, ID: \"tool1\"},\n\t\t}\n\t\taddr2 := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\"},\n\t\t}\n\n\t\tassert.False(t, addr1.Equals(addr2))\n\t})\n\n\t// Test Case 6: Address.Equals() with SubID differences\n\tt.Run(\"AddressSubIDDifference\", func(t *testing.T) {\n\t\taddr1 := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\", SubID: \"sub1\"},\n\t\t}\n\t\taddr2 := Address{\n\t\t\t{Type: AddressSegmentAgent, ID: \"agent1\", SubID: \"sub2\"},\n\t\t}\n\n\t\tassert.False(t, addr1.Equals(addr2))\n\t})\n}\n\nfunc TestAppendAddressSegment(t *testing.T) {\n\t// Test Case 1: Append to empty address\n\tt.Run(\"AppendToEmpty\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tnewCtx := AppendAddressSegment(ctx, AddressSegmentAgent, \"agent1\", \"\")\n\n\t\taddr := GetCurrentAddress(newCtx)\n\t\tassert.Len(t, addr, 1)\n\t\tassert.Equal(t, AddressSegmentAgent, addr[0].Type)\n\t\tassert.Equal(t, \"agent1\", addr[0].ID)\n\t\tassert.Equal(t, \"\", addr[0].SubID)\n\t})\n\n\t// Test Case 2: Append to existing address\n\tt.Run(\"AppendToExisting\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// First append\n\t\tctx = AppendAddressSegment(ctx, AddressSegmentAgent, \"agent1\", \"\")\n\n\t\t// Second append\n\t\tnewCtx := AppendAddressSegment(ctx, AddressSegmentTool, \"tool1\", \"call1\")\n\n\t\taddr := GetCurrentAddress(newCtx)\n\t\tassert.Len(t, addr, 2)\n\t\tassert.Equal(t, AddressSegmentAgent, addr[0].Type)\n\t\tassert.Equal(t, \"agent1\", addr[0].ID)\n\t\tassert.Equal(t, AddressSegmentTool, addr[1].Type)\n\t\tassert.Equal(t, \"tool1\", addr[1].ID)\n\t\tassert.Equal(t, \"call1\", addr[1].SubID)\n\t})\n\n\t// Test Case 3: Append with SubID\n\tt.Run(\"AppendWithSubID\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tnewCtx := AppendAddressSegment(ctx, AddressSegmentTool, \"tool1\", \"call123\")\n\n\t\taddr := GetCurrentAddress(newCtx)\n\t\tassert.Len(t, addr, 1)\n\t\tassert.Equal(t, AddressSegmentTool, addr[0].Type)\n\t\tassert.Equal(t, \"tool1\", addr[0].ID)\n\t\tassert.Equal(t, \"call123\", addr[0].SubID)\n\t})\n}\n\nfunc TestPopulateInterruptState(t *testing.T) {\n\t// Test Case 1: Populate with matching address\n\tt.Run(\"PopulateMatchingAddress\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Set up current address\n\t\tcurrentAddr := Address{{Type: AddressSegmentAgent, ID: \"agent1\"}}\n\t\trunCtx := &addrCtx{\n\t\t\taddr: currentAddr,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\t// Set up interrupt state data\n\t\tid2Addr := map[string]Address{\n\t\t\t\"interrupt1\": currentAddr,\n\t\t}\n\t\tid2State := map[string]InterruptState{\n\t\t\t\"interrupt1\": {State: \"test state\"},\n\t\t}\n\n\t\tnewCtx := PopulateInterruptState(ctx, id2Addr, id2State)\n\n\t\t// Verify the state was populated\n\t\twasInterrupted, hasState, state := GetInterruptState[string](newCtx)\n\t\tassert.True(t, wasInterrupted)\n\t\tassert.True(t, hasState)\n\t\tassert.Equal(t, \"test state\", state)\n\t})\n\n\t// Test Case 2: Populate with non-matching address\n\tt.Run(\"PopulateNonMatchingAddress\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\t// Set up current address\n\t\tcurrentAddr := Address{{Type: AddressSegmentAgent, ID: \"agent1\"}}\n\t\trunCtx := &addrCtx{\n\t\t\taddr: currentAddr,\n\t\t}\n\t\tctx = context.WithValue(ctx, addrCtxKey{}, runCtx)\n\n\t\t// Set up interrupt state data with different address\n\t\tid2Addr := map[string]Address{\n\t\t\t\"interrupt1\": {{Type: AddressSegmentAgent, ID: \"agent2\"}},\n\t\t}\n\t\tid2State := map[string]InterruptState{\n\t\t\t\"interrupt1\": {State: \"test state\"},\n\t\t}\n\n\t\tnewCtx := PopulateInterruptState(ctx, id2Addr, id2State)\n\n\t\t// Verify the state was NOT populated (no matching address)\n\t\twasInterrupted, hasState, state := GetInterruptState[string](newCtx)\n\t\tassert.False(t, wasInterrupted)\n\t\tassert.False(t, hasState)\n\t\tassert.Equal(t, \"\", state)\n\t})\n\n\t// Test Case 3: Populate with empty data\n\tt.Run(\"PopulateEmptyData\", func(t *testing.T) {\n\t\tctx := context.Background()\n\n\t\tnewCtx := PopulateInterruptState(ctx, map[string]Address{}, map[string]InterruptState{})\n\n\t\t// Verify no state was populated\n\t\twasInterrupted, hasState, state := GetInterruptState[string](newCtx)\n\t\tassert.False(t, wasInterrupted)\n\t\tassert.False(t, hasState)\n\t\tassert.Equal(t, \"\", state)\n\t})\n}\n\nfunc TestStringMethods(t *testing.T) {\n\t// Test Case 1: InterruptSignal.Error()\n\tt.Run(\"InterruptSignalError\", func(t *testing.T) {\n\t\tsignal := &InterruptSignal{\n\t\t\tID:      \"test-id\",\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tInterruptInfo: InterruptInfo{\n\t\t\t\tInfo: \"test info\",\n\t\t\t},\n\t\t\tInterruptState: InterruptState{\n\t\t\t\tState:                \"test state\",\n\t\t\t\tLayerSpecificPayload: \"test payload\",\n\t\t\t},\n\t\t\tSubs: []*InterruptSignal{\n\t\t\t\t{ID: \"sub1\"},\n\t\t\t},\n\t\t}\n\n\t\terrorStr := signal.Error()\n\t\texpectedContains := []string{\n\t\t\t\"interrupt signal:\",\n\t\t\t\"ID=test-id\",\n\t\t\t\"Addr=agent:agent1\",\n\t\t\t\"Info=interrupt info: Info=test info, IsRootCause=false\",\n\t\t\t\"State=interrupt state: State=test state, LayerSpecificPayload=test payload\",\n\t\t\t\"SubsLen=1\",\n\t\t}\n\n\t\tfor _, expected := range expectedContains {\n\t\t\tassert.Contains(t, errorStr, expected)\n\t\t}\n\t})\n\n\t// Test Case 2: InterruptState.String()\n\tt.Run(\"InterruptStateString\", func(t *testing.T) {\n\t\tstate := &InterruptState{\n\t\t\tState:                \"test state\",\n\t\t\tLayerSpecificPayload: \"test payload\",\n\t\t}\n\n\t\tresult := state.String()\n\t\texpected := \"interrupt state: State=test state, LayerSpecificPayload=test payload\"\n\t\tassert.Equal(t, expected, result)\n\t})\n\n\t// Test Case 3: InterruptState.String() with nil\n\tt.Run(\"InterruptStateStringNil\", func(t *testing.T) {\n\t\tvar state *InterruptState\n\t\tresult := state.String()\n\t\tassert.Equal(t, \"\", result)\n\t})\n\n\t// Test Case 4: InterruptInfo.String()\n\tt.Run(\"InterruptInfoString\", func(t *testing.T) {\n\t\tinfo := &InterruptInfo{\n\t\t\tInfo:        \"test info\",\n\t\t\tIsRootCause: true,\n\t\t}\n\n\t\tresult := info.String()\n\t\texpected := \"interrupt info: Info=test info, IsRootCause=true\"\n\t\tassert.Equal(t, expected, result)\n\t})\n\n\t// Test Case 5: InterruptInfo.String() with nil\n\tt.Run(\"InterruptInfoStringNil\", func(t *testing.T) {\n\t\tvar info *InterruptInfo\n\t\tresult := info.String()\n\t\tassert.Equal(t, \"\", result)\n\t})\n}\n\nfunc TestInterruptCtxEqualsWithoutID(t *testing.T) {\n\t// Test Case 1: Equal contexts\n\tt.Run(\"EqualContexts\", func(t *testing.T) {\n\t\tctx1 := &InterruptCtx{\n\t\t\tID:          \"id1\",\n\t\t\tAddress:     Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tInfo:        \"info1\",\n\t\t\tIsRootCause: true,\n\t\t}\n\t\tctx2 := &InterruptCtx{\n\t\t\tID:          \"id2\", // Different ID should be ignored\n\t\t\tAddress:     Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tInfo:        \"info1\",\n\t\t\tIsRootCause: true,\n\t\t}\n\n\t\tassert.True(t, ctx1.EqualsWithoutID(ctx2))\n\t})\n\n\t// Test Case 2: Different addresses\n\tt.Run(\"DifferentAddresses\", func(t *testing.T) {\n\t\tctx1 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t}\n\t\tctx2 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent2\"}},\n\t\t}\n\n\t\tassert.False(t, ctx1.EqualsWithoutID(ctx2))\n\t})\n\n\t// Test Case 3: Different root cause flags\n\tt.Run(\"DifferentRootCause\", func(t *testing.T) {\n\t\tctx1 := &InterruptCtx{\n\t\t\tAddress:     Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tIsRootCause: true,\n\t\t}\n\t\tctx2 := &InterruptCtx{\n\t\t\tAddress:     Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tIsRootCause: false,\n\t\t}\n\n\t\tassert.False(t, ctx1.EqualsWithoutID(ctx2))\n\t})\n\n\t// Test Case 4: Different info\n\tt.Run(\"DifferentInfo\", func(t *testing.T) {\n\t\tctx1 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tInfo:    \"info1\",\n\t\t}\n\t\tctx2 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tInfo:    \"info2\",\n\t\t}\n\n\t\tassert.False(t, ctx1.EqualsWithoutID(ctx2))\n\t})\n\n\t// Test Case 5: Nil contexts\n\tt.Run(\"NilContexts\", func(t *testing.T) {\n\t\tvar ctx1 *InterruptCtx\n\t\tvar ctx2 *InterruptCtx\n\n\t\tassert.True(t, ctx1.EqualsWithoutID(ctx2))\n\n\t\tctx3 := &InterruptCtx{}\n\t\tassert.False(t, ctx1.EqualsWithoutID(ctx3))\n\t\tassert.False(t, ctx3.EqualsWithoutID(ctx1))\n\t})\n\n\t// Test Case 6: With parent contexts\n\tt.Run(\"WithParentContexts\", func(t *testing.T) {\n\t\tparent1 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"parent\"}},\n\t\t}\n\t\tparent2 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"parent\"}},\n\t\t}\n\n\t\tctx1 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tParent:  parent1,\n\t\t}\n\t\tctx2 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tParent:  parent2,\n\t\t}\n\n\t\tassert.True(t, ctx1.EqualsWithoutID(ctx2))\n\t})\n\n\t// Test Case 7: Different parent contexts\n\tt.Run(\"DifferentParentContexts\", func(t *testing.T) {\n\t\tparent1 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"parent1\"}},\n\t\t}\n\t\tparent2 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"parent2\"}},\n\t\t}\n\n\t\tctx1 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tParent:  parent1,\n\t\t}\n\t\tctx2 := &InterruptCtx{\n\t\t\tAddress: Address{{Type: AddressSegmentAgent, ID: \"agent1\"}},\n\t\t\tParent:  parent2,\n\t\t}\n\n\t\tassert.False(t, ctx1.EqualsWithoutID(ctx2))\n\t})\n}\n"
  },
  {
    "path": "internal/core/resume.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage core\n\nimport \"context\"\n\n// GetInterruptState provides a type-safe way to check for and retrieve the persisted state from a previous interruption.\n// It is the primary function a component should use to understand its past state.\n//\n// It returns three values:\n//   - wasInterrupted (bool): True if the node was part of a previous interruption, regardless of whether state was provided.\n//   - state (T): The typed state object, if it was provided and matches type `T`.\n//   - hasState (bool): True if state was provided during the original interrupt and successfully cast to type `T`.\nfunc GetInterruptState[T any](ctx context.Context) (wasInterrupted bool, hasState bool, state T) {\n\trCtx, ok := getRunCtx(ctx)\n\tif !ok || rCtx.interruptState == nil {\n\t\treturn\n\t}\n\n\twasInterrupted = true\n\tif rCtx.interruptState.State == nil {\n\t\treturn\n\t}\n\n\tstate, hasState = rCtx.interruptState.State.(T)\n\treturn\n}\n\n// GetResumeContext checks if the current component is the target of a resume operation\n// and retrieves any data provided by the user for that resumption.\n//\n// This function is typically called *after* a component has already determined it is in a\n// resumed state by calling GetInterruptState.\n//\n// It returns three values:\n//   - isResumeTarget: A boolean that is true if the current component's address OR any of its\n//     descendant addresses was explicitly targeted by a call to Resume() or ResumeWithData().\n//     This allows composite components (like tools containing nested graphs) to know they should\n//     execute their children to reach the actual resume target.\n//   - hasData: A boolean that is true if data was provided for this specific component (i.e., not nil).\n//   - data: The typed data provided by the user.\n//\n// ### How to Use This Function: A Decision Framework\n//\n// The correct usage pattern depends on the application's desired resume strategy.\n//\n// #### Strategy 1: Implicit \"Resume All\"\n// In some use cases, any resume operation implies that *all* interrupted points should proceed.\n// For example, if an application's UI only provides a single \"Continue\" button for a set of\n// interruptions. In this model, a component can often just use `GetInterruptState` to see if\n// `wasInterrupted` is true and then proceed with its logic, as it can assume it is an intended target.\n// It may still call `GetResumeContext` to check for optional data, but the `isResumeFlow` flag is less critical.\n//\n// #### Strategy 2: Explicit \"Targeted Resume\" (Most Common)\n// For applications with multiple, distinct interrupt points that must be resumed independently, it is\n// crucial to differentiate which point is being resumed. This is the primary use case for the `isResumeTarget` flag.\n//   - If `isResumeTarget` is `true`: Your component (or one of its descendants) is the target.\n//     If `hasData` is true, you are the direct target and should consume the data.\n//     If `hasData` is false, a descendant is the target—execute your children to reach it.\n//   - If `isResumeTarget` is `false`: Neither you nor your descendants are the target. You MUST\n//     re-interrupt (e.g., by returning `StatefulInterrupt(...)`) to preserve your state.\n//\n// ### Guidance for Composite Components\n//\n// Composite components (like `Graph` or other `Runnable`s that contain sub-processes) have a dual role:\n//  1. Check for Self-Targeting: A composite component can itself be the target of a resume\n//     operation, for instance, to modify its internal state. It may call `GetResumeContext`\n//     to check for data targeted at its own address.\n//  2. Act as a Conduit: After checking for itself, its primary role is to re-execute its children,\n//     allowing the resume context to flow down to them. It must not consume a resume signal\n//     intended for one of its descendants.\nfunc GetResumeContext[T any](ctx context.Context) (isResumeTarget bool, hasData bool, data T) {\n\trCtx, ok := getRunCtx(ctx)\n\tif !ok {\n\t\treturn\n\t}\n\n\tisResumeTarget = rCtx.isResumeTarget\n\tif !isResumeTarget {\n\t\treturn\n\t}\n\n\t// It is a resume flow, now check for data\n\tif rCtx.resumeData == nil {\n\t\treturn // hasData is false\n\t}\n\n\tdata, hasData = rCtx.resumeData.(T)\n\treturn\n}\n\nfunc getRunCtx(ctx context.Context) (*addrCtx, bool) {\n\trCtx, ok := ctx.Value(addrCtxKey{}).(*addrCtx)\n\treturn rCtx, ok\n}\n"
  },
  {
    "path": "internal/generic/generic.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage generic\n\nimport (\n\t\"reflect\"\n)\n\n// NewInstance create an instance of the given type T.\n// the main purpose of this function is to create an instance of a type, can handle the type of T is a pointer or not.\n// eg. NewInstance[int] returns 0.\n// eg. NewInstance[*int] returns *0 (will be ptr of 0, not nil!).\nfunc NewInstance[T any]() T {\n\ttyp := TypeOf[T]()\n\n\tswitch typ.Kind() {\n\tcase reflect.Map:\n\t\treturn reflect.MakeMap(typ).Interface().(T)\n\tcase reflect.Slice, reflect.Array:\n\t\treturn reflect.MakeSlice(typ, 0, 0).Interface().(T)\n\tcase reflect.Ptr:\n\t\ttyp = typ.Elem()\n\t\torigin := reflect.New(typ)\n\t\tinst := origin\n\n\t\tfor typ.Kind() == reflect.Ptr {\n\t\t\ttyp = typ.Elem()\n\t\t\tinst = inst.Elem()\n\t\t\tinst.Set(reflect.New(typ))\n\t\t}\n\n\t\treturn origin.Interface().(T)\n\tdefault:\n\t\tvar t T\n\t\treturn t\n\t}\n}\n\n// TypeOf returns the type of T.\n// eg. TypeOf[int] returns reflect.TypeOf(int).\n// eg. TypeOf[*int] returns reflect.TypeOf(*int).\nfunc TypeOf[T any]() reflect.Type {\n\treturn reflect.TypeOf((*T)(nil)).Elem()\n}\n\n// PtrOf returns a pointer of T.\n// useful when you want to get a pointer of a value, in some config, for example.\n// eg. PtrOf[int] returns *int.\n// eg. PtrOf[*int] returns **int.\nfunc PtrOf[T any](v T) *T {\n\treturn &v\n}\n\ntype Pair[F, S any] struct {\n\tFirst  F\n\tSecond S\n}\n\n// Reverse returns a new slice with elements in reversed order.\nfunc Reverse[S ~[]E, E any](s S) S {\n\td := make(S, len(s))\n\tfor i := 0; i < len(s); i++ {\n\t\td[i] = s[len(s)-i-1]\n\t}\n\n\treturn d\n}\n\n// CopyMap copies a map to a new map.\nfunc CopyMap[K comparable, V any](src map[K]V) map[K]V {\n\tdst := make(map[K]V, len(src))\n\tfor k, v := range src {\n\t\tdst[k] = v\n\t}\n\treturn dst\n}\n"
  },
  {
    "path": "internal/generic/generic_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage generic\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestNewInstance(t *testing.T) {\n\tt.Run(\"struct\", func(t *testing.T) {\n\t\ttype Test struct{}\n\n\t\tinst := NewInstance[Test]()\n\n\t\tassert.IsType(t, Test{}, inst)\n\t})\n\n\tt.Run(\"pointer\", func(t *testing.T) {\n\t\ttype Test struct{}\n\n\t\tinst := NewInstance[*Test]()\n\n\t\tassert.IsType(t, &Test{}, inst)\n\t})\n\n\tt.Run(\"interface\", func(t *testing.T) {\n\t\ttype Test interface{}\n\n\t\tinst := NewInstance[Test]()\n\t\tassert.IsType(t, Test(nil), inst)\n\t})\n\n\tt.Run(\"pointer of pointer of pointer\", func(t *testing.T) {\n\t\ttype Test struct {\n\t\t\tValue int\n\t\t}\n\n\t\tinst := NewInstance[***Test]()\n\n\t\tptr := &Test{}\n\t\tptrOfPtr := &ptr\n\t\tassert.NotNil(t, inst)\n\t\tassert.NotNil(t, *inst)\n\t\tassert.IsType(t, ptrOfPtr, *inst)\n\t\tassert.NotNil(t, **inst)\n\t\tassert.Equal(t, Test{Value: 0}, ***inst)\n\t})\n\n\tt.Run(\"primitive_map\", func(t *testing.T) {\n\t\tinst := NewInstance[map[string]any]()\n\t\tassert.NotNil(t, inst)\n\t\tinst[\"a\"] = 1\n\t\tassert.Equal(t, map[string]any{\"a\": 1}, inst)\n\t})\n\n\tt.Run(\"primitive_slice\", func(t *testing.T) {\n\t\tinst := NewInstance[[]int]()\n\t\tassert.NotNil(t, inst)\n\t\tinst = append(inst, 1)\n\t\tassert.Equal(t, []int{1}, inst)\n\t})\n\n\tt.Run(\"primitive_string\", func(t *testing.T) {\n\t\tinst := NewInstance[string]()\n\t\tassert.Equal(t, \"\", inst)\n\t})\n\n\tt.Run(\"primitive_int64\", func(t *testing.T) {\n\t\tinst := NewInstance[int64]()\n\t\tassert.Equal(t, int64(0), inst)\n\t})\n}\n\nfunc TestReverse(t *testing.T) {\n\tt.Run(\"reverse int slice\", func(t *testing.T) {\n\t\tinput := []int{1, 2, 3, 4, 5}\n\t\texpected := []int{5, 4, 3, 2, 1}\n\t\tresult := Reverse(input)\n\t\tassert.Equal(t, expected, result)\n\t})\n\n\tt.Run(\"reverse string slice\", func(t *testing.T) {\n\t\tinput := []string{\"a\", \"b\", \"c\"}\n\t\texpected := []string{\"c\", \"b\", \"a\"}\n\t\tresult := Reverse(input)\n\t\tassert.Equal(t, expected, result)\n\t})\n}\n"
  },
  {
    "path": "internal/generic/type_name.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage generic\n\nimport (\n\t\"reflect\"\n\t\"regexp\"\n\t\"runtime\"\n\t\"strings\"\n)\n\nvar (\n\tregOfAnonymousFunc = regexp.MustCompile(`^func[0-9]+`)\n\tregOfNumber        = regexp.MustCompile(`^\\d+$`)\n)\n\n// ParseTypeName returns the name of the type of the given value.\n// It takes a reflect.Value as input and processes it to determine the underlying type. If the type is a pointer, it dereferences it to get the actual type. (the optimization of this function)\n// eg: ParseTypeName(reflect.ValueOf(&&myStruct{})) returns \"myStruct\" (not \"**myStruct\")\n//\n// If the type is a function, it retrieves the function's name, handling both named and anonymous functions.\n// examples of function paths: [package_path].[receiver_type].[func_name]\n// named function: xxx/utils.ParseTypeName\n// method: xxx/utils.(*MyStruct).Method\n// anonymous function: xxx/utils.TestParseTypeName.func6.1\nfunc ParseTypeName(val reflect.Value) string {\n\ttyp := val.Type()\n\n\tfor typ.Kind() == reflect.Pointer {\n\t\ttyp = typ.Elem()\n\t}\n\n\tif typ.Kind() == reflect.Func {\n\t\tfuncName := runtime.FuncForPC(val.Pointer()).Name()\n\t\tidx := strings.LastIndex(funcName, \".\")\n\t\tif idx < 0 {\n\t\t\tif funcName != \"\" {\n\t\t\t\treturn funcName\n\t\t\t}\n\t\t\treturn \"\"\n\t\t}\n\n\t\tname := funcName[idx+1:]\n\n\t\tif regOfAnonymousFunc.MatchString(name) {\n\t\t\treturn \"\"\n\t\t}\n\n\t\tif regOfNumber.MatchString(name) {\n\t\t\treturn \"\"\n\t\t}\n\n\t\treturn name\n\t}\n\n\treturn typ.Name()\n}\n"
  },
  {
    "path": "internal/generic/type_name_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage generic\n\nimport (\n\t\"reflect\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestParseTypeName(t *testing.T) {\n\tt.Run(\"named_struct\", func(t *testing.T) {\n\t\ttype OpenAI struct{}\n\t\tmodel := &OpenAI{}\n\t\tname := ParseTypeName(reflect.Indirect(reflect.ValueOf(model)))\n\t\tassert.Equal(t, \"OpenAI\", name)\n\t})\n\n\tt.Run(\"anonymous_struct\", func(t *testing.T) {\n\t\tmodel := &struct{}{}\n\t\tname := ParseTypeName(reflect.ValueOf(model))\n\t\tassert.Equal(t, \"\", name)\n\t})\n\n\tt.Run(\"anonymous_struct_from_func\", func(t *testing.T) {\n\t\tmodel := genStruct()\n\t\tname := ParseTypeName(reflect.ValueOf(model))\n\t\tassert.Equal(t, \"\", name)\n\t})\n\n\tt.Run(\"named_interface\", func(t *testing.T) {\n\t\ttype OpenAI interface{}\n\t\tmodel := OpenAI(&struct{}{})\n\t\tname := ParseTypeName(reflect.ValueOf(model))\n\t\tassert.Equal(t, \"\", name)\n\n\t\tname = ParseTypeName(reflect.ValueOf((*OpenAI)(nil)))\n\t\tassert.Equal(t, \"OpenAI\", name)\n\t})\n\n\tt.Run(\"named_function\", func(t *testing.T) {\n\t\tf := genOpenAI\n\t\tname := ParseTypeName(reflect.ValueOf(f))\n\t\tassert.Equal(t, \"genOpenAI\", name)\n\t})\n\n\tt.Run(\"anonymous_function\", func(t *testing.T) {\n\t\tf := genAnonymousFunc()\n\t\tname := ParseTypeName(reflect.ValueOf(f))\n\t\tassert.Equal(t, \"\", name)\n\n\t\tff := func(n string) {\n\t\t\t_ = n\n\t\t}\n\n\t\tname = ParseTypeName(reflect.ValueOf(ff))\n\t\tassert.Equal(t, \"\", name)\n\t})\n}\n\nfunc genStruct() *struct{} {\n\treturn &struct{}{}\n}\n\nfunc genOpenAI() {}\n\nfunc genAnonymousFunc() func(n string) {\n\treturn func(n string) {\n\t\t_ = n\n\t}\n}\n"
  },
  {
    "path": "internal/gmap/gmap.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage gmap\n\n// Concat returns the unions of maps as a new map.\n//\n// 💡 NOTE:\n//\n//   - Once the key conflicts, the newer value always replace the older one ([DiscardOld]),\n//   - If the result is an empty set, always return an empty map instead of nil\n//\n// 🚀 EXAMPLE:\n//\n//\tm := map[int]int{1: 1, 2: 2}\n//\tConcat(m, nil)             ⏩ map[int]int{1: 1, 2: 2}\n//\tConcat(m, map[int]{3: 3})  ⏩ map[int]int{1: 1, 2: 2, 3: 3}\n//\tConcat(m, map[int]{2: -1}) ⏩ map[int]int{1: 1, 2: -1} // \"2:2\" is replaced by the newer \"2:-1\"\n//\n// 💡 AKA: Merge, Union, Combine\nfunc Concat[K comparable, V any](ms ...map[K]V) map[K]V {\n\t// FastPath: no map or only one map given.\n\tif len(ms) == 0 {\n\t\treturn make(map[K]V)\n\t}\n\tif len(ms) == 1 {\n\t\treturn cloneWithoutNilCheck(ms[0])\n\t}\n\n\tvar maxLen int\n\tfor _, m := range ms {\n\t\tif len(m) > maxLen {\n\t\t\tmaxLen = len(m)\n\t\t}\n\t}\n\tret := make(map[K]V, maxLen)\n\t// FastPath: all maps are empty.\n\tif maxLen == 0 {\n\t\treturn ret\n\t}\n\n\t// Concat all maps.\n\tfor _, m := range ms {\n\t\tfor k, v := range m {\n\t\t\tret[k] = v\n\t\t}\n\t}\n\treturn ret\n}\n\n// Map applies function f to each key and value of map m.\n// Results of f are returned as a new map.\n//\n// 🚀 EXAMPLE:\n//\n//\tf := func(k, v int) (string, string) { return strconv.Itoa(k), strconv.Itoa(v) }\n//\tMap(map[int]int{1: 1}, f) ⏩ map[string]string{\"1\": \"1\"}\n//\tMap(map[int]int{}, f)     ⏩ map[string]string{}\nfunc Map[K1, K2 comparable, V1, V2 any](m map[K1]V1, f func(K1, V1) (K2, V2)) map[K2]V2 {\n\tr := make(map[K2]V2, len(m))\n\tfor k, v := range m {\n\t\tk2, v2 := f(k, v)\n\t\tr[k2] = v2\n\t}\n\treturn r\n}\n\n// Values returns the values of the map m.\n//\n// 🚀 EXAMPLE:\n//\n//\tm := map[int]string{1: \"1\", 2: \"2\", 3: \"3\", 4: \"4\"}\n//\tValues(m) ⏩ []string{\"1\", \"4\", \"2\", \"3\"} //⚠️INDETERMINATE ORDER⚠️\n//\n// ⚠️  WARNING: The keys values be in an indeterminate order,\nfunc Values[K comparable, V any](m map[K]V) []V {\n\tr := make([]V, 0, len(m))\n\tfor _, v := range m {\n\t\tr = append(r, v)\n\t}\n\treturn r\n}\n\n// Clone returns a shallow copy of map.\n// If the given map is nil, nil is returned.\n//\n// 🚀 EXAMPLE:\n//\n//\tClone(map[int]int{1: 1, 2: 2}) ⏩ map[int]int{1: 1, 2: 2}\n//\tClone(map[int]int{})           ⏩ map[int]int{}\n//\tClone[int, int](nil)           ⏩ nil\n//\n// 💡 HINT: Both keys and values are copied using assignment (=), so this is a shallow clone.\n// 💡 AKA: Copy\nfunc Clone[K comparable, V any, M ~map[K]V](m M) M {\n\tif m == nil {\n\t\treturn nil\n\t}\n\treturn cloneWithoutNilCheck(m)\n}\n\nfunc cloneWithoutNilCheck[K comparable, V any, M ~map[K]V](m M) M {\n\tr := make(M, len(m))\n\tfor k, v := range m {\n\t\tr[k] = v\n\t}\n\treturn r\n}\n"
  },
  {
    "path": "internal/gmap/gmap_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage gmap\n\nimport (\n\t\"fmt\"\n\t\"sort\"\n\t\"strconv\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestMerge(t *testing.T) {\n\tassert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4},\n\t\tConcat(map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, nil))\n\tassert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4},\n\t\tConcat[int, int](nil, map[int]int{1: 1, 2: 2, 3: 3, 4: 4}))\n\tassert.Equal(t, map[int]int{}, Concat[int, int](nil, nil))\n\tassert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4},\n\t\tConcat(map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, map[int]int{1: 1, 2: 2, 3: 3, 4: 4}))\n\tassert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4},\n\t\tConcat(map[int]int{1: 0, 2: 0}, map[int]int{1: 1, 2: 2, 3: 3, 4: 4}))\n\tassert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4},\n\t\tConcat(map[int]int{1: 1, 2: 1}, map[int]int{2: 2, 3: 3, 4: 4}))\n}\n\nfunc TestMap(t *testing.T) {\n\tassert.Equal(t,\n\t\tmap[string]string{\"1\": \"1\", \"2\": \"2\"},\n\t\tMap(map[int]int{1: 1, 2: 2}, func(k, v int) (string, string) {\n\t\t\treturn strconv.Itoa(k), strconv.Itoa(v)\n\t\t}))\n\tassert.Equal(t,\n\t\tmap[string]string{},\n\t\tMap(map[int]int{}, func(k, v int) (string, string) {\n\t\t\treturn strconv.Itoa(k), strconv.Itoa(v)\n\t\t}))\n}\n\nfunc TestValues(t *testing.T) {\n\t{\n\t\tkeys := Values(map[int]string{1: \"1\", 2: \"2\", 3: \"3\", 4: \"4\"})\n\t\tsort.Strings(keys)\n\t\tassert.Equal(t, []string{\"1\", \"2\", \"3\", \"4\"}, keys)\n\t}\n\tassert.Equal(t, []string{}, Values(map[int]string{}))\n\tassert.Equal(t, []string{}, Values[int, string](nil))\n}\n\nfunc TestClone(t *testing.T) {\n\tassert.Equal(t, map[int]int{1: 1, 2: 2}, Clone(map[int]int{1: 1, 2: 2}))\n\tvar nilMap map[int]int\n\tassert.Equal(t, map[int]int{}, Clone(map[int]int{}))\n\tassert.NotEqual(t, (map[int]int)(nil), Clone(map[int]int{}))\n\tassert.Equal(t, (map[int]int)(nil), Clone(nilMap))\n\tassert.NotEqual(t, map[int]int{}, Clone(nilMap))\n\n\t// Test new type.\n\ttype I2I map[int]int\n\tassert.Equal(t, I2I{1: 1, 2: 2}, Clone(I2I{1: 1, 2: 2}))\n\tassert.Equal(t, \"gmap.I2I\", fmt.Sprintf(\"%T\", Clone(I2I{})))\n\n\t// Test shallow clone.\n\tsrc := map[int]*int{1: ptr(1), 2: ptr(2)}\n\tdst := Clone(src)\n\tassert.Equal(t, src, dst)\n\tassert.True(t, src[1] == dst[1])\n\tassert.True(t, src[2] == dst[2])\n}\n\n// Ptr returns a pointer to the given value.\nfunc ptr[T any](v T) *T {\n\treturn &v\n}\n"
  },
  {
    "path": "internal/gslice/gslice.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage gslice\n\n// ToMap collects elements of slice to map, both map keys and values are produced\n// by mapping function f.\n//\n// 🚀 EXAMPLE:\n//\n//\ttype Foo struct {\n//\t\tID   int\n//\t\tName string\n//\t}\n//\tmapper := func(f Foo) (int, string) { return f.ID, f.Name }\n//\tToMap([]Foo{}, mapper) ⏩ map[int]string{}\n//\ts := []Foo{{1, \"one\"}, {2, \"two\"}, {3, \"three\"}}\n//\tToMap(s, mapper)       ⏩ map[int]string{1: \"one\", 2: \"two\", 3: \"three\"}\nfunc ToMap[T, V any, K comparable](s []T, f func(T) (K, V)) map[K]V {\n\tm := make(map[K]V, len(s))\n\tfor _, e := range s {\n\t\tk, v := f(e)\n\t\tm[k] = v\n\t}\n\treturn m\n}\n"
  },
  {
    "path": "internal/gslice/gslice_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage gslice\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestToMap(t *testing.T) {\n\ttype Foo struct {\n\t\tID   int\n\t\tName string\n\t}\n\tmapper := func(f Foo) (int, string) { return f.ID, f.Name }\n\tassert.Equal(t, map[int]string{}, ToMap([]Foo{}, mapper))\n\tassert.Equal(t, map[int]string{}, ToMap(nil, mapper))\n\tassert.Equal(t,\n\t\tmap[int]string{1: \"one\", 2: \"two\", 3: \"three\"},\n\t\tToMap([]Foo{{1, \"one\"}, {2, \"two\"}, {3, \"three\"}}, mapper))\n}\n"
  },
  {
    "path": "internal/merge.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage internal\n\nimport (\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\nvar mergeFuncs = map[reflect.Type]any{}\n\nfunc RegisterValuesMergeFunc[T any](fn func([]T) (T, error)) {\n\tmergeFuncs[generic.TypeOf[T]()] = fn\n}\n\nfunc GetMergeFunc(typ reflect.Type) func([]any) (any, error) {\n\tif fn, ok := mergeFuncs[typ]; ok {\n\t\treturn func(vs []any) (any, error) {\n\t\t\trvs := reflect.MakeSlice(reflect.SliceOf(typ), 0, len(vs))\n\t\t\tfor _, v := range vs {\n\t\t\t\tif t := reflect.TypeOf(v); t != typ {\n\t\t\t\t\treturn nil, fmt.Errorf(\n\t\t\t\t\t\t\"(values merge) field type mismatch. expected: '%v', got: '%v'\", typ, t)\n\t\t\t\t}\n\t\t\t\trvs = reflect.Append(rvs, reflect.ValueOf(v))\n\t\t\t}\n\n\t\t\trets := reflect.ValueOf(fn).Call([]reflect.Value{rvs})\n\t\t\tvar err error\n\t\t\tif !rets[1].IsNil() {\n\t\t\t\terr = rets[1].Interface().(error)\n\t\t\t}\n\t\t\treturn rets[0].Interface(), err\n\t\t}\n\t}\n\n\tif typ.Kind() == reflect.Map {\n\t\treturn func(vs []any) (any, error) {\n\t\t\treturn mergeMap(typ, vs)\n\t\t}\n\t}\n\n\treturn nil\n}\n\nfunc mergeMap(typ reflect.Type, vs []any) (any, error) {\n\tmerged := reflect.MakeMap(typ)\n\tfor _, v := range vs {\n\t\tif t := reflect.TypeOf(v); t != typ {\n\t\t\treturn nil, fmt.Errorf(\n\t\t\t\t\"(values merge map) field type mismatch. expected: '%v', got: '%v'\", typ, t)\n\t\t}\n\n\t\titer := reflect.ValueOf(v).MapRange()\n\t\tfor iter.Next() {\n\t\t\tkey, val := iter.Key(), iter.Value()\n\t\t\tif merged.MapIndex(key).IsValid() {\n\t\t\t\treturn nil, fmt.Errorf(\"(values merge map) duplicated key ('%v') found\", key.Interface())\n\t\t\t}\n\t\t\tmerged.SetMapIndex(key, val)\n\t\t}\n\t}\n\n\treturn merged.Interface(), nil\n}\n"
  },
  {
    "path": "internal/mock/adk/Agent_mock.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Code generated by MockGen. DO NOT EDIT.\n// Source: interface.go\n//\n// Generated by this command:\n//\n//\tmockgen -destination ../internal/mock/adk/Agent_mock.go --package adk -source interface.go\n//\n\n// Package adk is a generated GoMock package.\npackage adk\n\nimport (\n\tcontext \"context\"\n\treflect \"reflect\"\n\n\tadk \"github.com/cloudwego/eino/adk\"\n\tgomock \"go.uber.org/mock/gomock\"\n)\n\n// MockAgent is a mock of Agent interface.\ntype MockAgent struct {\n\tctrl     *gomock.Controller\n\trecorder *MockAgentMockRecorder\n\tisgomock struct{}\n}\n\n// MockAgentMockRecorder is the mock recorder for MockAgent.\ntype MockAgentMockRecorder struct {\n\tmock *MockAgent\n}\n\n// NewMockAgent creates a new mock instance.\nfunc NewMockAgent(ctrl *gomock.Controller) *MockAgent {\n\tmock := &MockAgent{ctrl: ctrl}\n\tmock.recorder = &MockAgentMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockAgent) EXPECT() *MockAgentMockRecorder {\n\treturn m.recorder\n}\n\n// Description mocks base method.\nfunc (m *MockAgent) Description(ctx context.Context) string {\n\tm.ctrl.T.Helper()\n\tret := m.ctrl.Call(m, \"Description\", ctx)\n\tret0, _ := ret[0].(string)\n\treturn ret0\n}\n\n// Description indicates an expected call of Description.\nfunc (mr *MockAgentMockRecorder) Description(ctx any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Description\", reflect.TypeOf((*MockAgent)(nil).Description), ctx)\n}\n\n// Name mocks base method.\nfunc (m *MockAgent) Name(ctx context.Context) string {\n\tm.ctrl.T.Helper()\n\tret := m.ctrl.Call(m, \"Name\", ctx)\n\tret0, _ := ret[0].(string)\n\treturn ret0\n}\n\n// Name indicates an expected call of Name.\nfunc (mr *MockAgentMockRecorder) Name(ctx any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Name\", reflect.TypeOf((*MockAgent)(nil).Name), ctx)\n}\n\n// Run mocks base method.\nfunc (m *MockAgent) Run(ctx context.Context, input *adk.AgentInput, options ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, input}\n\tfor _, a := range options {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Run\", varargs...)\n\tret0, _ := ret[0].(*adk.AsyncIterator[*adk.AgentEvent])\n\treturn ret0\n}\n\n// Run indicates an expected call of Run.\nfunc (mr *MockAgentMockRecorder) Run(ctx, input any, options ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, input}, options...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Run\", reflect.TypeOf((*MockAgent)(nil).Run), varargs...)\n}\n\n// MockOnSubAgents is a mock of OnSubAgents interface.\ntype MockOnSubAgents struct {\n\tctrl     *gomock.Controller\n\trecorder *MockOnSubAgentsMockRecorder\n\tisgomock struct{}\n}\n\n// MockOnSubAgentsMockRecorder is the mock recorder for MockOnSubAgents.\ntype MockOnSubAgentsMockRecorder struct {\n\tmock *MockOnSubAgents\n}\n\n// NewMockOnSubAgents creates a new mock instance.\nfunc NewMockOnSubAgents(ctrl *gomock.Controller) *MockOnSubAgents {\n\tmock := &MockOnSubAgents{ctrl: ctrl}\n\tmock.recorder = &MockOnSubAgentsMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockOnSubAgents) EXPECT() *MockOnSubAgentsMockRecorder {\n\treturn m.recorder\n}\n\n// OnDisallowTransferToParent mocks base method.\nfunc (m *MockOnSubAgents) OnDisallowTransferToParent(ctx context.Context) error {\n\tm.ctrl.T.Helper()\n\tret := m.ctrl.Call(m, \"OnDisallowTransferToParent\", ctx)\n\tret0, _ := ret[0].(error)\n\treturn ret0\n}\n\n// OnDisallowTransferToParent indicates an expected call of OnDisallowTransferToParent.\nfunc (mr *MockOnSubAgentsMockRecorder) OnDisallowTransferToParent(ctx any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"OnDisallowTransferToParent\", reflect.TypeOf((*MockOnSubAgents)(nil).OnDisallowTransferToParent), ctx)\n}\n\n// OnSetAsSubAgent mocks base method.\nfunc (m *MockOnSubAgents) OnSetAsSubAgent(ctx context.Context, parent adk.Agent) error {\n\tm.ctrl.T.Helper()\n\tret := m.ctrl.Call(m, \"OnSetAsSubAgent\", ctx, parent)\n\tret0, _ := ret[0].(error)\n\treturn ret0\n}\n\n// OnSetAsSubAgent indicates an expected call of OnSetAsSubAgent.\nfunc (mr *MockOnSubAgentsMockRecorder) OnSetAsSubAgent(ctx, parent any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"OnSetAsSubAgent\", reflect.TypeOf((*MockOnSubAgents)(nil).OnSetAsSubAgent), ctx, parent)\n}\n\n// OnSetSubAgents mocks base method.\nfunc (m *MockOnSubAgents) OnSetSubAgents(ctx context.Context, subAgents []adk.Agent) error {\n\tm.ctrl.T.Helper()\n\tret := m.ctrl.Call(m, \"OnSetSubAgents\", ctx, subAgents)\n\tret0, _ := ret[0].(error)\n\treturn ret0\n}\n\n// OnSetSubAgents indicates an expected call of OnSetSubAgents.\nfunc (mr *MockOnSubAgentsMockRecorder) OnSetSubAgents(ctx, subAgents any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"OnSetSubAgents\", reflect.TypeOf((*MockOnSubAgents)(nil).OnSetSubAgents), ctx, subAgents)\n}\n"
  },
  {
    "path": "internal/mock/components/document/document_mock.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Code generated by MockGen. DO NOT EDIT.\n// Source: interface.go\n\n// Package document is a generated GoMock package.\npackage document\n\nimport (\n\tcontext \"context\"\n\treflect \"reflect\"\n\n\tdocument \"github.com/cloudwego/eino/components/document\"\n\tschema \"github.com/cloudwego/eino/schema\"\n\tgomock \"go.uber.org/mock/gomock\"\n)\n\n// MockLoader is a mock of Loader interface.\ntype MockLoader struct {\n\tctrl     *gomock.Controller\n\trecorder *MockLoaderMockRecorder\n}\n\n// MockLoaderMockRecorder is the mock recorder for MockLoader.\ntype MockLoaderMockRecorder struct {\n\tmock *MockLoader\n}\n\n// NewMockLoader creates a new mock instance.\nfunc NewMockLoader(ctrl *gomock.Controller) *MockLoader {\n\tmock := &MockLoader{ctrl: ctrl}\n\tmock.recorder = &MockLoaderMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockLoader) EXPECT() *MockLoaderMockRecorder {\n\treturn m.recorder\n}\n\n// Load mocks base method.\nfunc (m *MockLoader) Load(ctx context.Context, src document.Source, opts ...document.LoaderOption) ([]*schema.Document, error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []interface{}{ctx, src}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Load\", varargs...)\n\tret0, _ := ret[0].([]*schema.Document)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Load indicates an expected call of Load.\nfunc (mr *MockLoaderMockRecorder) Load(ctx, src interface{}, opts ...interface{}) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]interface{}{ctx, src}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Load\", reflect.TypeOf((*MockLoader)(nil).Load), varargs...)\n}\n\n// MockTransformer is a mock of Transformer interface.\ntype MockTransformer struct {\n\tctrl     *gomock.Controller\n\trecorder *MockTransformerMockRecorder\n}\n\n// MockTransformerMockRecorder is the mock recorder for MockTransformer.\ntype MockTransformerMockRecorder struct {\n\tmock *MockTransformer\n}\n\n// NewMockTransformer creates a new mock instance.\nfunc NewMockTransformer(ctrl *gomock.Controller) *MockTransformer {\n\tmock := &MockTransformer{ctrl: ctrl}\n\tmock.recorder = &MockTransformerMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockTransformer) EXPECT() *MockTransformerMockRecorder {\n\treturn m.recorder\n}\n\n// Transform mocks base method.\nfunc (m *MockTransformer) Transform(ctx context.Context, src []*schema.Document, opts ...document.TransformerOption) ([]*schema.Document, error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []interface{}{ctx, src}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Transform\", varargs...)\n\tret0, _ := ret[0].([]*schema.Document)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Transform indicates an expected call of Transform.\nfunc (mr *MockTransformerMockRecorder) Transform(ctx, src interface{}, opts ...interface{}) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]interface{}{ctx, src}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Transform\", reflect.TypeOf((*MockTransformer)(nil).Transform), varargs...)\n}\n"
  },
  {
    "path": "internal/mock/components/embedding/Embedding_mock.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Code generated by MockGen. DO NOT EDIT.\n// Source: interface.go\n//\n// Generated by this command:\n//\n//\tmockgen -destination ../../internal/mock/components/embedding/Embedding_mock.go --package embedding -source interface.go\n//\n\n// Package embedding is a generated GoMock package.\npackage embedding\n\nimport (\n\tcontext \"context\"\n\treflect \"reflect\"\n\n\tembedding \"github.com/cloudwego/eino/components/embedding\"\n\tgomock \"go.uber.org/mock/gomock\"\n)\n\n// MockEmbedder is a mock of Embedder interface.\ntype MockEmbedder struct {\n\tctrl     *gomock.Controller\n\trecorder *MockEmbedderMockRecorder\n}\n\n// MockEmbedderMockRecorder is the mock recorder for MockEmbedder.\ntype MockEmbedderMockRecorder struct {\n\tmock *MockEmbedder\n}\n\n// NewMockEmbedder creates a new mock instance.\nfunc NewMockEmbedder(ctrl *gomock.Controller) *MockEmbedder {\n\tmock := &MockEmbedder{ctrl: ctrl}\n\tmock.recorder = &MockEmbedderMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockEmbedder) EXPECT() *MockEmbedderMockRecorder {\n\treturn m.recorder\n}\n\n// EmbedStrings mocks base method.\nfunc (m *MockEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, texts}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"EmbedStrings\", varargs...)\n\tret0, _ := ret[0].([][]float64)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// EmbedStrings indicates an expected call of EmbedStrings.\nfunc (mr *MockEmbedderMockRecorder) EmbedStrings(ctx, texts any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, texts}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"EmbedStrings\", reflect.TypeOf((*MockEmbedder)(nil).EmbedStrings), varargs...)\n}\n"
  },
  {
    "path": "internal/mock/components/indexer/indexer_mock.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Code generated by MockGen. DO NOT EDIT.\n// Source: interface.go\n//\n// Generated by this command:\n//\n//\tmockgen -destination ../../internal/mock/components/indexer/indexer_mock.go --package indexer -source interface.go\n//\n\n// Package indexer is a generated GoMock package.\npackage indexer\n\nimport (\n\tcontext \"context\"\n\treflect \"reflect\"\n\n\tindexer \"github.com/cloudwego/eino/components/indexer\"\n\tschema \"github.com/cloudwego/eino/schema\"\n\tgomock \"go.uber.org/mock/gomock\"\n)\n\n// MockIndexer is a mock of Indexer interface.\ntype MockIndexer struct {\n\tctrl     *gomock.Controller\n\trecorder *MockIndexerMockRecorder\n}\n\n// MockIndexerMockRecorder is the mock recorder for MockIndexer.\ntype MockIndexerMockRecorder struct {\n\tmock *MockIndexer\n}\n\n// NewMockIndexer creates a new mock instance.\nfunc NewMockIndexer(ctrl *gomock.Controller) *MockIndexer {\n\tmock := &MockIndexer{ctrl: ctrl}\n\tmock.recorder = &MockIndexerMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockIndexer) EXPECT() *MockIndexerMockRecorder {\n\treturn m.recorder\n}\n\n// Store mocks base method.\nfunc (m *MockIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) ([]string, error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, docs}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Store\", varargs...)\n\tret0, _ := ret[0].([]string)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Store indicates an expected call of Store.\nfunc (mr *MockIndexerMockRecorder) Store(ctx, docs any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, docs}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Store\", reflect.TypeOf((*MockIndexer)(nil).Store), varargs...)\n}\n"
  },
  {
    "path": "internal/mock/components/model/ChatModel_mock.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Code generated by MockGen. DO NOT EDIT.\n// Source: interface.go\n//\n// Generated by this command:\n//\n//\tmockgen -destination ../../internal/mock/components/model/ChatModel_mock.go --package model -source interface.go\n//\n\n// Package model is a generated GoMock package.\npackage model\n\nimport (\n\tcontext \"context\"\n\treflect \"reflect\"\n\n\tmodel \"github.com/cloudwego/eino/components/model\"\n\tschema \"github.com/cloudwego/eino/schema\"\n\tgomock \"go.uber.org/mock/gomock\"\n)\n\n// MockBaseChatModel is a mock of BaseChatModel interface.\ntype MockBaseChatModel struct {\n\tctrl     *gomock.Controller\n\trecorder *MockBaseChatModelMockRecorder\n\tisgomock struct{}\n}\n\n// MockBaseChatModelMockRecorder is the mock recorder for MockBaseChatModel.\ntype MockBaseChatModelMockRecorder struct {\n\tmock *MockBaseChatModel\n}\n\n// NewMockBaseChatModel creates a new mock instance.\nfunc NewMockBaseChatModel(ctrl *gomock.Controller) *MockBaseChatModel {\n\tmock := &MockBaseChatModel{ctrl: ctrl}\n\tmock.recorder = &MockBaseChatModelMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockBaseChatModel) EXPECT() *MockBaseChatModelMockRecorder {\n\treturn m.recorder\n}\n\n// Generate mocks base method.\nfunc (m *MockBaseChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, input}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Generate\", varargs...)\n\tret0, _ := ret[0].(*schema.Message)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Generate indicates an expected call of Generate.\nfunc (mr *MockBaseChatModelMockRecorder) Generate(ctx, input any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, input}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Generate\", reflect.TypeOf((*MockBaseChatModel)(nil).Generate), varargs...)\n}\n\n// Stream mocks base method.\nfunc (m *MockBaseChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, input}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Stream\", varargs...)\n\tret0, _ := ret[0].(*schema.StreamReader[*schema.Message])\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Stream indicates an expected call of Stream.\nfunc (mr *MockBaseChatModelMockRecorder) Stream(ctx, input any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, input}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Stream\", reflect.TypeOf((*MockBaseChatModel)(nil).Stream), varargs...)\n}\n\n// MockChatModel is a mock of ChatModel interface.\ntype MockChatModel struct {\n\tctrl     *gomock.Controller\n\trecorder *MockChatModelMockRecorder\n\tisgomock struct{}\n}\n\n// MockChatModelMockRecorder is the mock recorder for MockChatModel.\ntype MockChatModelMockRecorder struct {\n\tmock *MockChatModel\n}\n\n// NewMockChatModel creates a new mock instance.\nfunc NewMockChatModel(ctrl *gomock.Controller) *MockChatModel {\n\tmock := &MockChatModel{ctrl: ctrl}\n\tmock.recorder = &MockChatModelMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockChatModel) EXPECT() *MockChatModelMockRecorder {\n\treturn m.recorder\n}\n\n// BindTools mocks base method.\nfunc (m *MockChatModel) BindTools(tools []*schema.ToolInfo) error {\n\tm.ctrl.T.Helper()\n\tret := m.ctrl.Call(m, \"BindTools\", tools)\n\tret0, _ := ret[0].(error)\n\treturn ret0\n}\n\n// BindTools indicates an expected call of BindTools.\nfunc (mr *MockChatModelMockRecorder) BindTools(tools any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"BindTools\", reflect.TypeOf((*MockChatModel)(nil).BindTools), tools)\n}\n\n// Generate mocks base method.\nfunc (m *MockChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, input}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Generate\", varargs...)\n\tret0, _ := ret[0].(*schema.Message)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Generate indicates an expected call of Generate.\nfunc (mr *MockChatModelMockRecorder) Generate(ctx, input any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, input}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Generate\", reflect.TypeOf((*MockChatModel)(nil).Generate), varargs...)\n}\n\n// Stream mocks base method.\nfunc (m *MockChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, input}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Stream\", varargs...)\n\tret0, _ := ret[0].(*schema.StreamReader[*schema.Message])\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Stream indicates an expected call of Stream.\nfunc (mr *MockChatModelMockRecorder) Stream(ctx, input any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, input}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Stream\", reflect.TypeOf((*MockChatModel)(nil).Stream), varargs...)\n}\n\n// MockToolCallingChatModel is a mock of ToolCallingChatModel interface.\ntype MockToolCallingChatModel struct {\n\tctrl     *gomock.Controller\n\trecorder *MockToolCallingChatModelMockRecorder\n\tisgomock struct{}\n}\n\n// MockToolCallingChatModelMockRecorder is the mock recorder for MockToolCallingChatModel.\ntype MockToolCallingChatModelMockRecorder struct {\n\tmock *MockToolCallingChatModel\n}\n\n// NewMockToolCallingChatModel creates a new mock instance.\nfunc NewMockToolCallingChatModel(ctrl *gomock.Controller) *MockToolCallingChatModel {\n\tmock := &MockToolCallingChatModel{ctrl: ctrl}\n\tmock.recorder = &MockToolCallingChatModelMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockToolCallingChatModel) EXPECT() *MockToolCallingChatModelMockRecorder {\n\treturn m.recorder\n}\n\n// Generate mocks base method.\nfunc (m *MockToolCallingChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, input}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Generate\", varargs...)\n\tret0, _ := ret[0].(*schema.Message)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Generate indicates an expected call of Generate.\nfunc (mr *MockToolCallingChatModelMockRecorder) Generate(ctx, input any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, input}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Generate\", reflect.TypeOf((*MockToolCallingChatModel)(nil).Generate), varargs...)\n}\n\n// Stream mocks base method.\nfunc (m *MockToolCallingChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, input}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Stream\", varargs...)\n\tret0, _ := ret[0].(*schema.StreamReader[*schema.Message])\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Stream indicates an expected call of Stream.\nfunc (mr *MockToolCallingChatModelMockRecorder) Stream(ctx, input any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, input}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Stream\", reflect.TypeOf((*MockToolCallingChatModel)(nil).Stream), varargs...)\n}\n\n// WithTools mocks base method.\nfunc (m *MockToolCallingChatModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {\n\tm.ctrl.T.Helper()\n\tret := m.ctrl.Call(m, \"WithTools\", tools)\n\tret0, _ := ret[0].(model.ToolCallingChatModel)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// WithTools indicates an expected call of WithTools.\nfunc (mr *MockToolCallingChatModelMockRecorder) WithTools(tools any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"WithTools\", reflect.TypeOf((*MockToolCallingChatModel)(nil).WithTools), tools)\n}\n"
  },
  {
    "path": "internal/mock/components/retriever/retriever_mock.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Code generated by MockGen. DO NOT EDIT.\n// Source: interface.go\n//\n// Generated by this command:\n//\n//\tmockgen -destination ../../internal/mock/components/retriever/retriever_mock.go --package retriever -source interface.go\n//\n\n// Package retriever is a generated GoMock package.\npackage retriever\n\nimport (\n\tcontext \"context\"\n\treflect \"reflect\"\n\n\tretriever \"github.com/cloudwego/eino/components/retriever\"\n\tschema \"github.com/cloudwego/eino/schema\"\n\tgomock \"go.uber.org/mock/gomock\"\n)\n\n// MockRetriever is a mock of Retriever interface.\ntype MockRetriever struct {\n\tctrl     *gomock.Controller\n\trecorder *MockRetrieverMockRecorder\n}\n\n// MockRetrieverMockRecorder is the mock recorder for MockRetriever.\ntype MockRetrieverMockRecorder struct {\n\tmock *MockRetriever\n}\n\n// NewMockRetriever creates a new mock instance.\nfunc NewMockRetriever(ctrl *gomock.Controller) *MockRetriever {\n\tmock := &MockRetriever{ctrl: ctrl}\n\tmock.recorder = &MockRetrieverMockRecorder{mock}\n\treturn mock\n}\n\n// EXPECT returns an object that allows the caller to indicate expected use.\nfunc (m *MockRetriever) EXPECT() *MockRetrieverMockRecorder {\n\treturn m.recorder\n}\n\n// Retrieve mocks base method.\nfunc (m *MockRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {\n\tm.ctrl.T.Helper()\n\tvarargs := []any{ctx, query}\n\tfor _, a := range opts {\n\t\tvarargs = append(varargs, a)\n\t}\n\tret := m.ctrl.Call(m, \"Retrieve\", varargs...)\n\tret0, _ := ret[0].([]*schema.Document)\n\tret1, _ := ret[1].(error)\n\treturn ret0, ret1\n}\n\n// Retrieve indicates an expected call of Retrieve.\nfunc (mr *MockRetrieverMockRecorder) Retrieve(ctx, query any, opts ...any) *gomock.Call {\n\tmr.mock.ctrl.T.Helper()\n\tvarargs := append([]any{ctx, query}, opts...)\n\treturn mr.mock.ctrl.RecordCallWithMethodType(mr.mock, \"Retrieve\", reflect.TypeOf((*MockRetriever)(nil).Retrieve), varargs...)\n}\n"
  },
  {
    "path": "internal/mock/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package mock provides mock implementations for testing purposes.\n//\n// This package aims to provide mock implementations for interfaces in the components package,\n// making it easier to use in testing environments. It includes mock implementations for\n// various core components such as retrievers, tools, message handlers, and graph runners.\n//\n// Directory Structure:\n//   - components/: Contains mock implementations for various components\n//   - retriever/: Provides mock implementation for the Retriever interface\n//   - retriever_mock.go: Mock implementation for document retrieval\n//   - tool/: Mock implementations for tool-related interfaces\n//   - message/: Mock implementations for message handling components\n//   - graph/: Mock implementations for graph execution components\n//   - stream/: Mock implementations for streaming components\n//\n// Usage:\n// These mock implementations are primarily used in unit tests and integration tests,\n// allowing developers to conduct tests without depending on actual external services.\n// Each mock component strictly follows the contract of its corresponding interface\n// while providing controllable behaviors and results.\n//\n// Examples:\n//\n//   - Using mock retriever:\n//     retriever := mock.NewMockRetriever()\n//     // Configure retriever behavior\n//\n//   - Using mock tool:\n//     tool := mock.NewMockTool()\n//     // Configure tool behavior\n//\n//   - Using mock graph runner:\n//     runner := mock.NewMockGraphRunner()\n//     // Configure runner behavior\npackage mock\n"
  },
  {
    "path": "internal/safe/panic.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage safe\n\nimport (\n\t\"fmt\"\n)\n\ntype panicErr struct {\n\tinfo  any\n\tstack []byte\n}\n\nfunc (p *panicErr) Error() string {\n\treturn fmt.Sprintf(\"panic error: %v, \\nstack: %s\", p.info, string(p.stack))\n}\n\n// NewPanicErr creates a new panic error.\n// panicErr is a wrapper of panic info and stack trace.\n// it implements the error interface, can print error message of info and stack trace.\nfunc NewPanicErr(info any, stack []byte) error {\n\treturn &panicErr{\n\t\tinfo:  info,\n\t\tstack: stack,\n\t}\n}\n"
  },
  {
    "path": "internal/safe/panic_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage safe\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestPanicErr(t *testing.T) {\n\terr := NewPanicErr(\"info\", []byte(\"stack\"))\n\tassert.Equal(t, \"panic error: info, \\nstack: stack\", err.Error())\n}\n"
  },
  {
    "path": "internal/serialization/serialization.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage serialization\n\nimport (\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"reflect\"\n\n\t\"github.com/bytedance/sonic\"\n)\n\nvar m = map[string]reflect.Type{}\nvar rm = map[reflect.Type]string{}\n\nfunc init() {\n\t_ = GenericRegister[int](\"_eino_int\")\n\t_ = GenericRegister[int8](\"_eino_int8\")\n\t_ = GenericRegister[int16](\"_eino_int16\")\n\t_ = GenericRegister[int32](\"_eino_int32\")\n\t_ = GenericRegister[int64](\"_eino_int64\")\n\t_ = GenericRegister[uint](\"_eino_uint\")\n\t_ = GenericRegister[uint8](\"_eino_uint8\")\n\t_ = GenericRegister[uint16](\"_eino_uint16\")\n\t_ = GenericRegister[uint32](\"_eino_uint32\")\n\t_ = GenericRegister[uint64](\"_eino_uint64\")\n\t_ = GenericRegister[float32](\"_eino_float32\")\n\t_ = GenericRegister[float64](\"_eino_float64\")\n\t_ = GenericRegister[complex64](\"_eino_complex64\")\n\t_ = GenericRegister[complex128](\"_eino_complex128\")\n\t_ = GenericRegister[uintptr](\"_eino_uintptr\")\n\t_ = GenericRegister[bool](\"_eino_bool\")\n\t_ = GenericRegister[string](\"_eino_string\")\n\t_ = GenericRegister[any](\"_eino_any\")\n}\n\nfunc GenericRegister[T any](key string) error {\n\tt := reflect.TypeOf((*T)(nil)).Elem()\n\tfor t.Kind() == reflect.Ptr {\n\t\tt = t.Elem()\n\t}\n\tif nt, ok := m[key]; ok {\n\t\treturn fmt.Errorf(\"key[%s] already registered to %s\", key, nt.String())\n\t}\n\tif nk, ok := rm[t]; ok {\n\t\treturn fmt.Errorf(\"type[%s] already registered to %s\", t.String(), nk)\n\t}\n\tm[key] = t\n\trm[t] = key\n\treturn nil\n}\n\ntype InternalSerializer struct{}\n\nfunc (i *InternalSerializer) Marshal(v any) ([]byte, error) {\n\tis, err := internalMarshal(v, nil)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn sonic.Marshal(is)\n}\n\nfunc (i *InternalSerializer) Unmarshal(data []byte, v any) error {\n\tval, err := unmarshal(data, reflect.TypeOf(v))\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to unmarshal: %w\", err)\n\t}\n\n\trv := reflect.ValueOf(v)\n\tif rv.Kind() != reflect.Ptr || rv.IsNil() {\n\t\treturn fmt.Errorf(\"failed to unmarshal: value must be a non-nil pointer\")\n\t}\n\n\ttarget := rv.Elem()\n\tif !target.CanSet() {\n\t\treturn fmt.Errorf(\"failed to unmarshal: output value must be settable\")\n\t}\n\n\tif val == nil {\n\t\ttarget.Set(reflect.Zero(target.Type()))\n\t\treturn nil\n\t}\n\n\tsource := reflect.ValueOf(val)\n\n\tvar set func(target, source reflect.Value) bool\n\tset = func(target, source reflect.Value) bool {\n\t\tif !source.IsValid() {\n\t\t\ttarget.Set(reflect.Zero(target.Type()))\n\t\t\treturn true\n\t\t}\n\t\tif source.Type().AssignableTo(target.Type()) {\n\t\t\ttarget.Set(source)\n\t\t\treturn true\n\t\t}\n\n\t\tif target.Kind() == reflect.Ptr {\n\t\t\tif target.IsNil() {\n\t\t\t\tif !target.CanSet() {\n\t\t\t\t\treturn false\n\t\t\t\t}\n\t\t\t\ttarget.Set(reflect.New(target.Type().Elem()))\n\t\t\t}\n\t\t\treturn set(target.Elem(), source)\n\t\t}\n\n\t\tif source.Kind() == reflect.Ptr {\n\t\t\tif source.IsNil() {\n\t\t\t\ttarget.Set(reflect.Zero(target.Type()))\n\t\t\t\treturn true\n\t\t\t}\n\t\t\treturn set(target, source.Elem())\n\t\t}\n\n\t\tif source.Type().ConvertibleTo(target.Type()) {\n\t\t\ttarget.Set(source.Convert(target.Type()))\n\t\t\treturn true\n\t\t}\n\n\t\treturn false\n\t}\n\n\tif set(target, source) {\n\t\treturn nil\n\t}\n\n\treturn fmt.Errorf(\"failed to unmarshal: cannot assign %s to %s\", reflect.TypeOf(val), target.Type())\n}\n\nfunc unmarshal(data []byte, t reflect.Type) (any, error) {\n\tis := &internalStruct{}\n\terr := sonic.Unmarshal(data, is)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn internalUnmarshal(is, t)\n}\n\ntype internalStruct struct {\n\tType *valueType `json:\",omitempty\"`\n\n\tJSONValue json.RawMessage `json:\",omitempty\"`\n\n\t// map or struct\n\t// in map, the key is the serialized map key anyway todo: if key is string, don't serialize\n\t// in struct, the key is the original field name\n\tMapValues map[string]*internalStruct `json:\",omitempty\"`\n\n\t// slice\n\tSliceValues []*internalStruct `json:\",omitempty\"`\n}\n\ntype valueType struct {\n\tPointerNum uint32 `json:\",omitempty\"`\n\n\tSimpleType string `json:\",omitempty\"`\n\n\tStructType string `json:\",omitempty\"`\n\n\tMapKeyType   *valueType `json:\",omitempty\"`\n\tMapValueType *valueType `json:\",omitempty\"`\n\n\tSliceValueType *valueType `json:\",omitempty\"`\n}\n\nfunc extractType(t reflect.Type) (*valueType, error) {\n\tret := &valueType{}\n\tfor t.Kind() == reflect.Ptr {\n\t\tret.PointerNum += 1\n\t\tt = t.Elem()\n\t}\n\tvar err error\n\tif t.Kind() == reflect.Map {\n\t\tret.MapKeyType, err = extractType(t.Key())\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tret.MapValueType, err = extractType(t.Elem())\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t} else if t.Kind() == reflect.Slice || t.Kind() == reflect.Array {\n\t\tret.SliceValueType, err = extractType(t.Elem())\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t} else {\n\t\tkey, ok := rm[t]\n\t\tif !ok {\n\t\t\treturn ret, fmt.Errorf(\"unknown type: %s\", t.String())\n\t\t}\n\t\tret.SimpleType = key\n\t}\n\treturn ret, nil\n}\n\nfunc restoreType(vt *valueType) (reflect.Type, error) {\n\tif vt.SimpleType != \"\" {\n\t\trt, ok := m[vt.SimpleType]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"unknown type: %s\", vt.SimpleType)\n\t\t}\n\t\treturn resolvePointerNum(vt.PointerNum, rt), nil\n\t}\n\tif vt.StructType != \"\" {\n\t\trt, ok := m[vt.StructType]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"unknown type: %s\", vt.StructType)\n\t\t}\n\t\treturn resolvePointerNum(vt.PointerNum, rt), nil\n\t}\n\tif vt.MapKeyType != nil {\n\t\trkt, err := restoreType(vt.MapKeyType)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\trvt, err := restoreType(vt.MapValueType)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn resolvePointerNum(vt.PointerNum, reflect.MapOf(rkt, rvt)), nil\n\t}\n\tif vt.SliceValueType != nil {\n\t\trt, err := restoreType(vt.SliceValueType)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn resolvePointerNum(vt.PointerNum, reflect.SliceOf(rt)), nil\n\t}\n\treturn nil, fmt.Errorf(\"empty value\")\n}\n\nfunc internalMarshal(v any, fieldType reflect.Type) (*internalStruct, error) {\n\tif v == nil ||\n\t\t(reflect.ValueOf(v).IsZero() && fieldType != nil && fieldType.Kind() != reflect.Interface) {\n\t\treturn nil, nil\n\t}\n\n\tret := &internalStruct{}\n\trv := reflect.ValueOf(v)\n\trt := rv.Type()\n\ttypeUnspecific := fieldType == nil || fieldType.Kind() == reflect.Interface\n\n\tvar pointerNum uint32\n\tfor rt.Kind() == reflect.Ptr {\n\t\tpointerNum++\n\t\tif !rv.IsNil() {\n\t\t\trv = rv.Elem()\n\t\t\trt = rt.Elem()\n\t\t\tcontinue\n\t\t}\n\t\tfor rt.Kind() == reflect.Ptr {\n\t\t\trt = rt.Elem()\n\t\t}\n\t\tif typeUnspecific {\n\t\t\t// need type registered\n\t\t\tkey, ok := rm[rt]\n\t\t\tif !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"unknown type: %v\", rt)\n\t\t\t}\n\t\t\tret.Type = &valueType{\n\t\t\t\tPointerNum: pointerNum,\n\t\t\t\tSimpleType: key,\n\t\t\t}\n\t\t}\n\t\tret.JSONValue = json.RawMessage(\"null\")\n\t\treturn ret, nil\n\t}\n\n\tswitch rt.Kind() {\n\tcase reflect.Struct:\n\t\tif typeUnspecific {\n\t\t\t// need type registered\n\t\t\tkey, ok := rm[rt]\n\t\t\tif !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"unknown type: %v\", rt)\n\t\t\t}\n\n\t\t\tif checkMarshaler(rt) {\n\t\t\t\tret.Type = &valueType{\n\t\t\t\t\tPointerNum: pointerNum,\n\t\t\t\t\tSimpleType: key,\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tret.Type = &valueType{\n\t\t\t\t\tPointerNum: pointerNum,\n\t\t\t\t\tStructType: key,\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tif checkMarshaler(rt) {\n\t\t\tjsonBytes, err := json.Marshal(rv.Interface())\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tret.JSONValue = jsonBytes\n\t\t\treturn ret, nil\n\t\t}\n\n\t\tret.MapValues = make(map[string]*internalStruct)\n\n\t\tfor i := 0; i < rt.NumField(); i++ {\n\t\t\tfield := rt.Field(i)\n\t\t\t// only handle exported fields\n\t\t\tif field.PkgPath == \"\" {\n\t\t\t\tk := field.Name\n\t\t\t\tv := rv.Field(i)\n\n\t\t\t\tinternalValue, err := internalMarshal(v.Interface(), field.Type)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\n\t\t\t\tret.MapValues[k] = internalValue\n\t\t\t}\n\t\t}\n\n\t\treturn ret, nil\n\tcase reflect.Map:\n\t\tif typeUnspecific {\n\t\t\tvar err error\n\t\t\tret.Type = &valueType{\n\t\t\t\tPointerNum: pointerNum,\n\t\t\t}\n\t\t\t// map key type\n\t\t\tret.Type.MapKeyType, err = extractType(rt.Key())\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\t// map value type\n\t\t\tret.Type.MapValueType, err = extractType(rt.Elem())\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\n\t\tret.MapValues = make(map[string]*internalStruct)\n\n\t\titer := rv.MapRange()\n\t\tfor iter.Next() {\n\t\t\tk := iter.Key()\n\t\t\tv := iter.Value()\n\n\t\t\tinternalValue, err := internalMarshal(v.Interface(), rt.Elem())\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\tkeyStr, err := sonic.MarshalString(k.Interface())\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"marshaling map key[%v] fail: %v\", k.Interface(), err)\n\t\t\t}\n\t\t\tret.MapValues[keyStr] = internalValue\n\t\t}\n\n\t\treturn ret, nil\n\tcase reflect.Slice, reflect.Array:\n\t\tif typeUnspecific {\n\t\t\tvar err error\n\t\t\tret.Type = &valueType{PointerNum: pointerNum}\n\t\t\tret.Type.SliceValueType, err = extractType(rt.Elem())\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\n\t\tlength := rv.Len()\n\t\tret.SliceValues = make([]*internalStruct, length)\n\n\t\tfor i := 0; i < length; i++ {\n\t\t\tinternalValue, err := internalMarshal(rv.Index(i).Interface(), rt.Elem())\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tret.SliceValues[i] = internalValue\n\t\t}\n\n\t\treturn ret, nil\n\n\tdefault:\n\t\tif typeUnspecific {\n\t\t\tkey, ok := rm[rv.Type()]\n\t\t\tif !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"unknown type: %v\", rt)\n\t\t\t}\n\t\t\tret.Type = &valueType{\n\t\t\t\tPointerNum: pointerNum,\n\t\t\t\tSimpleType: key,\n\t\t\t}\n\t\t}\n\n\t\tjsonBytes, err := json.Marshal(rv.Interface())\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tret.JSONValue = jsonBytes\n\t\treturn ret, nil\n\t}\n}\n\nfunc internalUnmarshal(v *internalStruct, typ reflect.Type) (any, error) {\n\tif v == nil {\n\t\treturn nil, nil\n\t}\n\n\tif v.Type == nil {\n\t\t// specific type\n\t\tif checkMarshaler(typ) {\n\t\t\tpv := reflect.New(typ)\n\t\t\terr := json.Unmarshal(v.JSONValue, pv.Interface())\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\treturn pv.Elem().Interface(), nil\n\t\t}\n\t\treturn internalSpecificTypeUnmarshal(v, typ)\n\t}\n\n\tif len(v.Type.SimpleType) != 0 {\n\t\t// based type\n\t\tt, ok := m[v.Type.SimpleType]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"unknown type key: %v\", v.Type)\n\t\t}\n\t\tpResult := reflect.New(resolvePointerNum(v.Type.PointerNum, t))\n\t\terr := sonic.Unmarshal(v.JSONValue, pResult.Interface())\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"unmarshal type[%s] fail: %v, data: %s\", t.String(), err, string(v.JSONValue))\n\t\t}\n\t\treturn pResult.Elem().Interface(), nil\n\t}\n\n\tif len(v.Type.StructType) > 0 {\n\t\t// struct\n\t\trt, ok := m[v.Type.StructType]\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"unknown type key: %v\", v.Type.StructType)\n\t\t}\n\t\tresult, dResult := createValueFromType(resolvePointerNum(v.Type.PointerNum, rt))\n\n\t\terr := setStructFields(dResult, v.MapValues)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn result.Interface(), nil\n\t}\n\n\tif v.Type.MapKeyType != nil {\n\t\t// map\n\t\trkt, err := restoreType(v.Type.MapKeyType)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\trvt, err := restoreType(v.Type.MapValueType)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tresult, dResult := createValueFromType(reflect.MapOf(rkt, rvt))\n\t\terr = setMapKVs(dResult, v.MapValues)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn result.Interface(), nil\n\t}\n\n\t// slice\n\trvt, err := restoreType(v.Type.SliceValueType)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tresult, dResult := createValueFromType(reflect.SliceOf(rvt))\n\terr = setSliceElems(dResult, v.SliceValues)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn result.Interface(), nil\n}\n\nfunc internalSpecificTypeUnmarshal(is *internalStruct, typ reflect.Type) (any, error) {\n\t_, dtyp := derefPointerNum(typ)\n\tresult, dResult := createValueFromType(typ)\n\n\tif dtyp.Kind() == reflect.Struct {\n\t\terr := setStructFields(dResult, is.MapValues)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn result.Interface(), nil\n\t} else if dtyp.Kind() == reflect.Map {\n\t\terr := setMapKVs(dResult, is.MapValues)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn result.Interface(), nil\n\t} else if dtyp.Kind() == reflect.Array || dtyp.Kind() == reflect.Slice {\n\t\terr := setSliceElems(dResult, is.SliceValues)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn result.Interface(), nil\n\t}\n\t// simple type\n\tv := reflect.New(typ)\n\terr := sonic.Unmarshal(is.JSONValue, v.Interface())\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"unmarshal type[%s] fail: %v\", typ.String(), err)\n\t}\n\treturn v.Elem().Interface(), nil\n}\n\nfunc setSliceElems(dResult reflect.Value, values []*internalStruct) error {\n\tt := dResult.Type()\n\n\t// Handle arrays differently from slices\n\t// Arrays have fixed size and cannot use reflect.Append\n\tif dResult.Kind() == reflect.Array {\n\t\tfor i, internalValue := range values {\n\t\t\tif i >= dResult.Len() {\n\t\t\t\treturn fmt.Errorf(\"array index out of bounds: trying to set index %d in array of length %d\", i, dResult.Len())\n\t\t\t}\n\t\t\tvalue, err := internalUnmarshal(internalValue, t.Elem())\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"unmarshal array[%s] element %d fail: %v\", t.Elem(), i, err)\n\t\t\t}\n\t\t\tif value == nil {\n\t\t\t\tdResult.Index(i).Set(reflect.Zero(t.Elem()))\n\t\t\t} else {\n\t\t\t\tdResult.Index(i).Set(reflect.ValueOf(value))\n\t\t\t}\n\t\t}\n\t\treturn nil\n\t}\n\n\t// For slices, use Append as before\n\tfor _, internalValue := range values {\n\t\tvalue, err := internalUnmarshal(internalValue, t.Elem())\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"unmarshal slice[%s] fail: %v\", t.Elem(), err)\n\t\t}\n\t\tif value == nil {\n\t\t\t// empty value\n\t\t\tdResult.Set(reflect.Append(dResult, reflect.New(t.Elem()).Elem()))\n\t\t} else {\n\t\t\tdResult.Set(reflect.Append(dResult, reflect.ValueOf(value)))\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc setMapKVs(dResult reflect.Value, values map[string]*internalStruct) error {\n\tt := dResult.Type()\n\tfor marshaledMapKey, internalValue := range values {\n\t\tprkv := reflect.New(t.Key())\n\t\terr := sonic.UnmarshalString(marshaledMapKey, prkv.Interface())\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"unmarshal map key[%v] to type[%s] fail: %v\", marshaledMapKey, t.Key(), err)\n\t\t}\n\n\t\tvalue, err := internalUnmarshal(internalValue, t.Elem())\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"unmarshal map value fail: %v\", err)\n\t\t}\n\t\tif value == nil {\n\t\t\tdResult.SetMapIndex(prkv.Elem(), reflect.New(t.Elem()).Elem())\n\t\t} else {\n\t\t\tdResult.SetMapIndex(prkv.Elem(), reflect.ValueOf(value))\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc setStructFields(dResult reflect.Value, values map[string]*internalStruct) error {\n\tt := dResult.Type()\n\tfor k, internalValue := range values {\n\t\tsf, ok := t.FieldByName(k)\n\t\tif !ok {\n\t\t\tcontinue\n\t\t}\n\t\tvalue, err := internalUnmarshal(internalValue, sf.Type)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"unmarshal map field[%v] fail: %v\", k, err)\n\t\t}\n\t\terr = setStructField(t, dResult, k, value)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc setStructField(t reflect.Type, s reflect.Value, fieldName string, val any) error {\n\tfield := s.FieldByName(fieldName)\n\tif !field.CanSet() {\n\t\treturn fmt.Errorf(\"unmarshal map fail, can not set field %v\", fieldName)\n\t}\n\tif val == nil {\n\t\trft, ok := t.FieldByName(fieldName)\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"unmarshal map fail, cannot find field: %v\", fieldName)\n\t\t}\n\t\tfield.Set(reflect.New(rft.Type).Elem())\n\t} else {\n\t\tfield.Set(reflect.ValueOf(val))\n\t}\n\treturn nil\n}\n\nfunc resolvePointerNum(pointerNum uint32, t reflect.Type) reflect.Type {\n\tfor i := uint32(0); i < pointerNum; i++ {\n\t\tt = reflect.PointerTo(t)\n\t}\n\treturn t\n}\n\nfunc derefPointerNum(t reflect.Type) (uint32, reflect.Type) {\n\tvar ptrCount uint32 = 0\n\n\tfor t != nil && t.Kind() == reflect.Ptr {\n\t\tt = t.Elem()\n\t\tptrCount++\n\t}\n\n\treturn ptrCount, t\n}\n\nfunc createValueFromType(t reflect.Type) (value reflect.Value, derefValue reflect.Value) {\n\tvalue = reflect.New(t).Elem()\n\n\tderefValue = value\n\tfor derefValue.Kind() == reflect.Ptr {\n\t\tif derefValue.IsNil() {\n\t\t\tderefValue.Set(reflect.New(derefValue.Type().Elem()))\n\t\t}\n\t\tderefValue = derefValue.Elem()\n\t}\n\n\tif derefValue.Kind() == reflect.Map && derefValue.IsNil() {\n\t\tderefValue.Set(reflect.MakeMap(derefValue.Type()))\n\t}\n\n\t// Use Len() == 0 instead of IsNil() for slices to avoid panic\n\t// IsNil() can panic on uninitialized slice values created via reflect.New().Elem()\n\tif derefValue.Kind() == reflect.Slice {\n\t\tif derefValue.Len() == 0 && derefValue.Cap() == 0 {\n\t\t\tderefValue.Set(reflect.MakeSlice(derefValue.Type(), 0, 0))\n\t\t}\n\t}\n\t// Arrays cannot be nil and don't need initialization\n\n\treturn value, derefValue\n}\n\nvar marshalerType = reflect.TypeOf((*json.Marshaler)(nil)).Elem()\nvar unmarshalerType = reflect.TypeOf((*json.Unmarshaler)(nil)).Elem()\n\nfunc checkMarshaler(t reflect.Type) bool {\n\tfor t.Kind() == reflect.Ptr {\n\t\tt = t.Elem()\n\t}\n\n\tif (t.Implements(marshalerType) || reflect.PointerTo(t).Implements(marshalerType)) &&\n\t\t(t.Implements(unmarshalerType) || reflect.PointerTo(t).Implements(unmarshalerType)) {\n\t\treturn true\n\t}\n\treturn false\n}\n"
  },
  {
    "path": "internal/serialization/serialization_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage serialization\n\nimport (\n\t\"reflect\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n)\n\ntype myInterface interface {\n\tMethod()\n}\ntype myStruct struct {\n\tA string\n}\n\nfunc (m *myStruct) Method() {}\n\ntype myStruct2 struct {\n\tA any\n\tB myInterface\n\tC map[string]**myStruct\n\tD map[myStruct]any\n\tE []any\n\tf string\n\tG myStruct3\n\tH *myStruct4\n\tI []*myStruct3\n\tJ map[string]myStruct3\n\tK myStruct4\n\tL []*myStruct4\n\tM map[string]myStruct4\n}\n\ntype myStruct3 struct {\n\tFieldA string\n}\n\ntype myStruct4 struct {\n\tFieldA string\n}\n\nfunc (m *myStruct4) UnmarshalJSON(bytes []byte) error {\n\tm.FieldA = string(bytes)\n\treturn nil\n}\n\nfunc (m myStruct4) MarshalJSON() ([]byte, error) {\n\treturn []byte(m.FieldA), nil\n}\n\nfunc TestSerialization(t *testing.T) {\n\t_ = GenericRegister[myStruct](\"myStruct\")\n\t_ = GenericRegister[myStruct2](\"myStruct2\")\n\t_ = GenericRegister[myInterface](\"myInterface\")\n\tms := myStruct{A: \"test\"}\n\tpms := &ms\n\tpointerOfPointerOfMyStruct := &pms\n\n\tms1 := myStruct{A: \"1\"}\n\tms2 := myStruct{A: \"2\"}\n\tms3 := myStruct{A: \"3\"}\n\tms4 := myStruct{A: \"4\"}\n\tvalues := []any{\n\t\t10,\n\t\t\"test\",\n\t\tms,\n\t\tpms,\n\t\tpointerOfPointerOfMyStruct,\n\t\tmyInterface(pms),\n\t\t[]int{1, 2, 3},\n\t\t[]any{1, \"test\"},\n\t\t[]myInterface{nil, &myStruct{A: \"1\"}, &myStruct{A: \"2\"}},\n\t\tmap[string]string{\"123\": \"123\", \"abc\": \"abc\"},\n\t\tmap[string]myInterface{\"1\": nil, \"2\": pms},\n\t\tmap[string]any{\"123\": 1, \"abc\": &myStruct{A: \"1\"}, \"bcd\": nil},\n\t\tmap[myStruct]any{\n\t\t\tms1: 1,\n\t\t\tms2: &myStruct{\n\t\t\t\tA: \"2\",\n\t\t\t},\n\t\t\tms3: nil,\n\t\t\tms4: []any{\n\t\t\t\t1,\n\t\t\t\tpointerOfPointerOfMyStruct,\n\t\t\t\t\"123\", &myStruct{\n\t\t\t\t\tA: \"1\",\n\t\t\t\t},\n\t\t\t\tnil,\n\t\t\t\tmap[myStruct]any{\n\t\t\t\t\tms1: 1,\n\t\t\t\t\tms2: nil,\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\tmyStruct2{\n\t\t\tA: \"123\",\n\t\t\tB: &myStruct{\n\t\t\t\tA: \"test\",\n\t\t\t},\n\t\t\tC: map[string]**myStruct{\n\t\t\t\t\"a\": pointerOfPointerOfMyStruct,\n\t\t\t},\n\t\t\tD: map[myStruct]any{{\"a\"}: 1},\n\t\t\tE: []any{1, \"2\", 3},\n\t\t\tf: \"\",\n\t\t\tG: myStruct3{\n\t\t\t\tFieldA: \"1\",\n\t\t\t},\n\t\t\tH: nil,\n\t\t\tI: []*myStruct3{\n\t\t\t\t{FieldA: \"2\"}, {FieldA: \"3\"},\n\t\t\t},\n\t\t\tJ: map[string]myStruct3{\n\t\t\t\t\"1\": {FieldA: \"4\"},\n\t\t\t\t\"2\": {FieldA: \"5\"},\n\t\t\t},\n\t\t\tK: myStruct4{\n\t\t\t\tFieldA: \"1\",\n\t\t\t},\n\t\t\tL: []*myStruct4{\n\t\t\t\t{FieldA: \"2\"}, {FieldA: \"3\"},\n\t\t\t},\n\t\t\tM: map[string]myStruct4{\n\t\t\t\t\"1\": {FieldA: \"4\"},\n\t\t\t\t\"2\": {FieldA: \"5\"},\n\t\t\t},\n\t\t},\n\t\tmap[string]map[string][]map[string][][]string{\n\t\t\t\"1\": {\n\t\t\t\t\"a\": []map[string][][]string{\n\t\t\t\t\t{\"b\": {\n\t\t\t\t\t\t{\"c\"},\n\t\t\t\t\t\t{\"d\"},\n\t\t\t\t\t}},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t[]*myStruct{},\n\t\t&myStruct{},\n\t}\n\n\tfor _, value := range values {\n\t\tdata, err := (&InternalSerializer{}).Marshal(value)\n\t\tassert.NoError(t, err)\n\t\tv := reflect.New(reflect.TypeOf(value)).Interface()\n\t\terr = (&InternalSerializer{}).Unmarshal(data, v)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, value, reflect.ValueOf(v).Elem().Interface())\n\t}\n}\n\ntype myStruct5 struct {\n\tFieldA string\n}\n\nfunc (m *myStruct5) UnmarshalJSON(bytes []byte) error {\n\tm.FieldA = \"FieldA\"\n\treturn nil\n}\n\nfunc (m myStruct5) MarshalJSON() ([]byte, error) {\n\treturn []byte(\"1\"), nil\n}\n\nfunc TestMarshalStruct(t *testing.T) {\n\tassert.NoError(t, GenericRegister[myStruct5](\"myStruct5\"))\n\ts := myStruct5{FieldA: \"1\"}\n\tdata, err := (&InternalSerializer{}).Marshal(s)\n\tassert.NoError(t, err)\n\tresult := &myStruct5{}\n\terr = (&InternalSerializer{}).Unmarshal(data, result)\n\tassert.NoError(t, err)\n\tassert.Equal(t, myStruct5{FieldA: \"FieldA\"}, *result)\n\n\tma := map[string]any{\n\t\t\"1\": s,\n\t}\n\tdata, err = (&InternalSerializer{}).Marshal(ma)\n\tassert.NoError(t, err)\n\tresult2 := map[string]any{}\n\terr = (&InternalSerializer{}).Unmarshal(data, &result2)\n\tassert.NoError(t, err)\n\tassert.Equal(t, map[string]any{\n\t\t\"1\": myStruct5{FieldA: \"FieldA\"},\n\t}, result2)\n}\n\ntype unmarshalTestStruct struct {\n\tFoo string\n\tBar int\n}\n\nfunc init() {\n\t// Register types for the serializer to work.\n\t// This is necessary for the serializer to know how to handle custom struct types.\n\terr := GenericRegister[unmarshalTestStruct](\"unmarshalTestStruct\")\n\tif err != nil {\n\t\tpanic(err)\n\t}\n}\n\nfunc TestInternalSerializer_Unmarshal(t *testing.T) {\n\ts := InternalSerializer{}\n\n\tt.Run(\"success cases\", func(t *testing.T) {\n\t\t// Helper to create a pointer to a value, needed for the expected value in one test case.\n\t\tptr := func(i int) *int { return &i }\n\n\t\ttestCases := []struct {\n\t\t\tname        string\n\t\t\tinputValue  any\n\t\t\toutputPtr   any\n\t\t\texpectedVal any\n\t\t}{\n\t\t\t{\n\t\t\t\tname:        \"simple type\",\n\t\t\t\tinputValue:  123,\n\t\t\t\toutputPtr:   new(int),\n\t\t\t\texpectedVal: 123,\n\t\t\t},\n\t\t\t{\n\t\t\t\tname:        \"struct type\",\n\t\t\t\tinputValue:  unmarshalTestStruct{Foo: \"hello\", Bar: 42},\n\t\t\t\toutputPtr:   new(unmarshalTestStruct),\n\t\t\t\texpectedVal: unmarshalTestStruct{Foo: \"hello\", Bar: 42},\n\t\t\t},\n\t\t\t{\n\t\t\t\tname:        \"pointer to struct\",\n\t\t\t\tinputValue:  &unmarshalTestStruct{Foo: \"world\", Bar: 99},\n\t\t\t\toutputPtr:   new(*unmarshalTestStruct),\n\t\t\t\texpectedVal: &unmarshalTestStruct{Foo: \"world\", Bar: 99},\n\t\t\t},\n\t\t\t{\n\t\t\t\tname:        \"unmarshal pointer to value\",\n\t\t\t\tinputValue:  &unmarshalTestStruct{Foo: \"p2v\", Bar: 1},\n\t\t\t\toutputPtr:   new(unmarshalTestStruct),\n\t\t\t\texpectedVal: unmarshalTestStruct{Foo: \"p2v\", Bar: 1},\n\t\t\t},\n\t\t\t{\n\t\t\t\tname:        \"unmarshal value to pointer\",\n\t\t\t\tinputValue:  unmarshalTestStruct{Foo: \"v2p\", Bar: 2},\n\t\t\t\toutputPtr:   new(*unmarshalTestStruct),\n\t\t\t\texpectedVal: &unmarshalTestStruct{Foo: \"v2p\", Bar: 2},\n\t\t\t},\n\t\t\t{\n\t\t\t\tname:        \"unmarshal nil pointer\",\n\t\t\t\tinputValue:  (*unmarshalTestStruct)(nil),\n\t\t\t\toutputPtr:   &struct{ v *unmarshalTestStruct }{v: &unmarshalTestStruct{}}, // placeholder to be replaced\n\t\t\t\texpectedVal: (*unmarshalTestStruct)(nil),\n\t\t\t},\n\t\t\t{\n\t\t\t\tname:        \"convertible types\",\n\t\t\t\tinputValue:  int32(42),\n\t\t\t\toutputPtr:   new(int64),\n\t\t\t\texpectedVal: int64(42),\n\t\t\t},\n\t\t\t{\n\t\t\t\tname:        \"pointer to pointer destination\",\n\t\t\t\tinputValue:  12345,\n\t\t\t\toutputPtr:   new(*int),\n\t\t\t\texpectedVal: ptr(12345),\n\t\t\t},\n\t\t\t{\n\t\t\t\tname:        \"unmarshal to any\",\n\t\t\t\tinputValue:  unmarshalTestStruct{Foo: \"any\", Bar: 101},\n\t\t\t\toutputPtr:   new(any),\n\t\t\t\texpectedVal: unmarshalTestStruct{Foo: \"any\", Bar: 101},\n\t\t\t},\n\t\t}\n\n\t\tfor _, tc := range testCases {\n\t\t\tt.Run(tc.name, func(t *testing.T) {\n\t\t\t\tdata, err := s.Marshal(tc.inputValue)\n\t\t\t\trequire.NoError(t, err)\n\n\t\t\t\t// Special handling for the nil test case to correctly pass the pointer.\n\t\t\t\tif tc.name == \"unmarshal nil pointer\" {\n\t\t\t\t\ttarget := tc.outputPtr.(*struct{ v *unmarshalTestStruct })\n\t\t\t\t\terr = s.Unmarshal(data, &target.v)\n\t\t\t\t\trequire.NoError(t, err)\n\t\t\t\t\tassert.Nil(t, target.v)\n\t\t\t\t\treturn\n\t\t\t\t}\n\n\t\t\t\terr = s.Unmarshal(data, tc.outputPtr)\n\t\t\t\trequire.NoError(t, err)\n\n\t\t\t\t// Dereference the pointer to get the actual value for comparison.\n\t\t\t\tactualVal := reflect.ValueOf(tc.outputPtr).Elem().Interface()\n\t\t\t\tassert.Equal(t, tc.expectedVal, actualVal)\n\t\t\t})\n\t\t}\n\t})\n\n\tt.Run(\"error cases\", func(t *testing.T) {\n\t\tdata, err := s.Marshal(123)\n\t\trequire.NoError(t, err)\n\n\t\tt.Run(\"destination not a pointer\", func(t *testing.T) {\n\t\t\tvar output int\n\t\t\terr := s.Unmarshal(data, output)\n\t\t\trequire.Error(t, err)\n\t\t\tassert.Contains(t, err.Error(), \"value must be a non-nil pointer\")\n\t\t})\n\n\t\tt.Run(\"destination is a nil pointer\", func(t *testing.T) {\n\t\t\tvar output *int // nil\n\t\t\terr := s.Unmarshal(data, output)\n\t\t\trequire.Error(t, err)\n\t\t\tassert.Contains(t, err.Error(), \"value must be a non-nil pointer\")\n\t\t})\n\n\t\tt.Run(\"type mismatch\", func(t *testing.T) {\n\t\t\tstrData, mErr := s.Marshal(\"i am a string\")\n\t\t\trequire.NoError(t, mErr)\n\n\t\t\tvar output int\n\t\t\terr := s.Unmarshal(strData, &output)\n\t\t\trequire.Error(t, err)\n\t\t\tassert.Contains(t, err.Error(), \"cannot assign\")\n\t\t})\n\n\t\tt.Run(\"unconvertible types\", func(t *testing.T) {\n\t\t\tintData, mErr := s.Marshal(123)\n\t\t\trequire.NoError(t, mErr)\n\n\t\t\tvar output bool\n\t\t\terr := s.Unmarshal(intData, &output)\n\t\t\trequire.Error(t, err)\n\t\t\tassert.Contains(t, err.Error(), \"cannot assign\")\n\t\t})\n\t})\n}\n"
  },
  {
    "path": "llms.txt",
    "content": "# Eino\n\n> Eino is a Go-based LLM application development framework by ByteDance.\n> It provides component abstractions, a graph/chain orchestration engine,\n> streaming primitives, a callback system, and an Agent Development Kit (ADK)\n> for building production-grade LLM applications.\n\n## Repositories\n\n- [eino](https://github.com/cloudwego/eino) — core framework (this repo)\n- [eino-ext](https://github.com/cloudwego/eino-ext) — component integrations (OpenAI, Ark, Ollama, Redis, S3, …)\n- [eino-examples](https://github.com/cloudwego/eino-examples) — runnable example applications\n\n## Overview & Background\n\n- [Overview](https://www.cloudwego.io/docs/eino/overview/)\n- [ByteDance Eino Practice](https://www.cloudwego.io/docs/eino/overview/bytedance_eino_practice/)\n- [Eino Open Source](https://www.cloudwego.io/docs/eino/overview/eino_open_source/)\n- [Graph or Agent — when to use which](https://www.cloudwego.io/docs/eino/overview/graph_or_agent/)\n\n## Quick Start\n\n- [Simple LLM Application](https://www.cloudwego.io/docs/eino/quick_start/simple_llm_application/)\n- [Agent with Tools](https://www.cloudwego.io/docs/eino/quick_start/agent_llm_with_tools/)\n- [Eino Cookbook](https://www.cloudwego.io/docs/eino/eino-cookbook/)\n\n## Core Concepts — Components\n\nComponents are the typed building blocks of eino pipelines. Each has a defined\ninterface in the core repo; implementations live in eino-ext.\n\n- [Components overview](https://www.cloudwego.io/docs/eino/core_modules/components/)\n- [ChatModel](https://www.cloudwego.io/docs/eino/core_modules/components/chat_model_guide/)\n- [ChatTemplate](https://www.cloudwego.io/docs/eino/core_modules/components/chat_template_guide/)\n- [ToolsNode](https://www.cloudwego.io/docs/eino/core_modules/components/tools_node_guide/)\n- [How to create a Tool](https://www.cloudwego.io/docs/eino/core_modules/components/tools_node_guide/how_to_create_a_tool/)\n- [Retriever](https://www.cloudwego.io/docs/eino/core_modules/components/retriever_guide/)\n- [Indexer](https://www.cloudwego.io/docs/eino/core_modules/components/indexer_guide/)\n- [Embedding](https://www.cloudwego.io/docs/eino/core_modules/components/embedding_guide/)\n- [DocumentLoader](https://www.cloudwego.io/docs/eino/core_modules/components/document_loader_guide/)\n- [DocumentParser](https://www.cloudwego.io/docs/eino/core_modules/components/document_loader_guide/document_parser_interface_guide/)\n- [DocumentTransformer](https://www.cloudwego.io/docs/eino/core_modules/components/document_transformer_guide/)\n- [Lambda](https://www.cloudwego.io/docs/eino/core_modules/components/lambda_guide/)\n- [AgenticChatModel](https://www.cloudwego.io/docs/eino/core_modules/components/agentic_chat_model_guide/)\n- [AgenticChatTemplate](https://www.cloudwego.io/docs/eino/core_modules/components/agentic_chat_template_guide/)\n- [AgenticToolsNode](https://www.cloudwego.io/docs/eino/core_modules/components/agentic_tools_node_guide/)\n\n## Core Concepts — Orchestration\n\nThe orchestration layer composes components into executable pipelines.\nChain is a linear sequence; Graph is a DAG with conditional edges;\nWorkflow is a higher-level structured abstraction over Graph.\n\n- [Chain & Graph introduction](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/chain_graph_introduction/)\n- [Orchestration design principles](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/)\n- [Workflow orchestration framework](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/workflow_orchestration_framework/)\n- [Stream programming essentials](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/)\n- [Callback system](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/callback_manual/)\n- [CallOption capabilities](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/call_option_capabilities/)\n- [Checkpoint & interrupt/resume](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/checkpoint_interrupt/)\n\n## Core Concepts — Flow Integration\n\n- [ReAct agent](https://www.cloudwego.io/docs/eino/core_modules/flow_integration_components/react_agent_manual/)\n- [Multi-agent hosting](https://www.cloudwego.io/docs/eino/core_modules/flow_integration_components/multi_agent_hosting/)\n\n## Agent Development Kit (ADK)\n\nThe ADK provides a higher-level runtime for building, composing, and deploying\nagents. It sits above the Graph layer and introduces Agent, Skill, and\nMiddleware abstractions.\n\n- [ADK overview](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/)\n- [Agent quickstart](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_quickstart/)\n- [Agent interface](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_interface/)\n- [Agent collaboration](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_collaboration/)\n- [Agent implementations](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/)\n  - [ChatModel agent](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/chat_model/)\n  - [Workflow agent](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/workflow/)\n  - [Supervisor agent](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/supervisor/)\n  - [Plan-and-execute agent](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/plan_execute/)\n  - [Deep agents](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/deepagents/)\n- [Agent extension](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_extension/)\n- [Human-in-the-loop (HITL)](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_hitl/)\n- [ADK callbacks](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/adk_agent_callback/)\n- [ChatModelAgent middleware](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/)\n  - [Filesystem middleware](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_filesystem/)\n  - [Skill middleware](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_skill/)\n  - [Summarization middleware](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_summarization/)\n  - [Plan-task middleware](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_plantask/)\n  - [Tool-search middleware](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_toolsearch/)\n  - [Tool-reduction middleware](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_toolreduction/)\n  - [Patch-toolcalls middleware](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_patchtoolcalls/)\n\n## DevOps Tooling\n\n- [IDE plugin guide](https://www.cloudwego.io/docs/eino/core_modules/devops/ide_plugin_guide/)\n- [Visual orchestration plugin](https://www.cloudwego.io/docs/eino/core_modules/devops/visual_orchestration_plugin_guide/)\n- [Visual debug plugin](https://www.cloudwego.io/docs/eino/core_modules/devops/visual_debug_plugin_guide/)\n\n## Ecosystem Integrations (eino-ext)\n\n- [ChatModel integrations](https://www.cloudwego.io/docs/eino/ecosystem_integration/chat_model/)\n- [Document integrations](https://www.cloudwego.io/docs/eino/ecosystem_integration/document/)\n- [Embedding integrations](https://www.cloudwego.io/docs/eino/ecosystem_integration/embedding/)\n- [Tool integrations](https://www.cloudwego.io/docs/eino/ecosystem_integration/tool/)\n- [Callback integrations](https://www.cloudwego.io/docs/eino/ecosystem_integration/callbacks/)\n- [Indexer integrations](https://www.cloudwego.io/docs/eino/ecosystem_integration/indexer/)\n- [Retriever integrations](https://www.cloudwego.io/docs/eino/ecosystem_integration/retriever/)\n- [ChatTemplate integrations](https://www.cloudwego.io/docs/eino/ecosystem_integration/chat_template/)\n\n## Release Notes & Migration\n\n- [v0.1](https://www.cloudwego.io/docs/eino/release_notes_and_migration/v01_first_release/)\n- [v0.2](https://www.cloudwego.io/docs/eino/release_notes_and_migration/v02_second_release/)\n- [v0.3 — breaking changes](https://www.cloudwego.io/docs/eino/release_notes_and_migration/v03_tiny_break_change/)\n- [v0.4 — compose optimization](https://www.cloudwego.io/docs/eino/release_notes_and_migration/eino_v0.4._-compose_optimization/)\n- [v0.5 — ADK implementation](https://www.cloudwego.io/docs/eino/release_notes_and_migration/eino_v0.5._-adk_implementation/)\n- [v0.6 — JSON schema optimization](https://www.cloudwego.io/docs/eino/release_notes_and_migration/eino_v0.6._-jsonschema_optimization/)\n- [v0.7 — interrupt/resume refactor](https://www.cloudwego.io/docs/eino/release_notes_and_migration/eino_v0.7._-interrupt_resume_refactor/)\n- [v0.8 — ADK middlewares](https://www.cloudwego.io/docs/eino/release_notes_and_migration/eino_v0.8._-adk_middlewares/)\n\n## FAQ\n\n- [FAQ](https://www.cloudwego.io/docs/eino/faq/)\n"
  },
  {
    "path": "schema/doc.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package schema defines the core data structures and utilities shared across\n// all Eino components.\n//\n// # Key Types\n//\n// [Message] is the universal unit of communication between users, models, and\n// tools. It carries role, text content, multimodal media, tool calls, and\n// response metadata. Helper constructors — [UserMessage], [SystemMessage],\n// [AssistantMessage], [ToolMessage] — cover the most common cases.\n//\n// [Document] represents a piece of text with a metadata map. Typed accessors\n// (Score, SubIndexes, DenseVector, SparseVector, DSLInfo, ExtraInfo) read and\n// write well-known metadata keys so pipeline stages can pass structured data\n// without coupling to specific struct types.\n//\n// [ToolInfo] describes a tool's name, description, and parameter schema.\n// Parameters can be declared either as a [ParameterInfo] map (simple, struct-\n// like) or as a raw [jsonschema.Schema] (full JSON Schema 2020-12 expressiveness).\n// [ToolChoice] controls whether the model must, may, or must not call tools.\n//\n// # Streaming\n//\n// [StreamReader] and [StreamWriter] are the building blocks for streaming data\n// through Eino pipelines. Create a linked pair with [Pipe]:\n//\n//\tsr, sw := schema.Pipe[*schema.Message](10)\n//\tgo func() {\n//\t\tdefer sw.Close()\n//\t\tsw.Send(chunk, nil)\n//\t}()\n//\tdefer sr.Close()\n//\tfor {\n//\t\tchunk, err := sr.Recv()\n//\t\tif errors.Is(err, io.EOF) { break }\n//\t}\n//\n// Important constraints:\n//   - A StreamReader is read-once: only one goroutine may call Recv.\n//   - Always call Close, even when the loop ends on io.EOF, to release resources.\n//   - To give the same stream to multiple consumers, call [StreamReader.Copy].\n//\n// # Four Streaming Paradigms\n//\n// Eino components and Lambda functions are classified by their input/output\n// streaming shape. The framework automatically bridges mismatches:\n//\n//   - Invoke: non-streaming in, non-streaming out (ping-pong).\n//   - Stream: non-streaming in, StreamReader out (server-streaming). ChatModel\n//     and Tool support this.\n//   - Collect: StreamReader in, non-streaming out (client-streaming). Useful\n//     for branch conditions that decide after the first chunk.\n//   - Transform: StreamReader in, StreamReader out (bidirectional).\n//\n// When an upstream node outputs T but a downstream node only accepts\n// StreamReader[T], the framework wraps T in a single-chunk StreamReader —\n// this is called a \"fake stream\". It satisfies the interface but does NOT\n// reduce time-to-first-chunk. Conversely, when a downstream node only accepts\n// T but the upstream outputs StreamReader[T], the framework automatically\n// concatenates the stream into a complete T.\n//\n// Utility functions:\n//   - [StreamReaderFromArray] wraps a slice as a stream (useful in tests).\n//   - [MergeStreamReaders] fans-in multiple streams into one.\n//   - [MergeNamedStreamReaders] like MergeStreamReaders but emits [SourceEOF]\n//     when each named source ends, useful for tracking per-source completion.\n//   - [StreamReaderWithConvert] transforms element types; return [ErrNoValue]\n//     from the convert function to skip an element.\n//\n// See https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/\npackage schema\n"
  },
  {
    "path": "schema/document.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nconst (\n\tdocMetaDataKeySubIndexes   = \"_sub_indexes\"\n\tdocMetaDataKeyScore        = \"_score\"\n\tdocMetaDataKeyExtraInfo    = \"_extra_info\"\n\tdocMetaDataKeyDSL          = \"_dsl\"\n\tdocMetaDataKeyDenseVector  = \"_dense_vector\"\n\tdocMetaDataKeySparseVector = \"_sparse_vector\"\n)\n\n// Document is a piece of text with a metadata map. It is the shared currency\n// between Loader, Transformer, Indexer, and Retriever components.\n//\n// Metadata is an open map[string]any that lets pipeline stages attach typed\n// values to a document without creating a new struct. Well-known keys are\n// managed through typed accessor methods — Score, SubIndexes, DenseVector,\n// SparseVector, DSLInfo, ExtraInfo — so callers never need to reference the\n// raw key strings.\n//\n// Transformer implementations should preserve existing metadata and merge new\n// keys rather than replacing the map outright, so provenance information\n// accumulated by earlier stages is not lost.\ntype Document struct {\n\t// ID is the unique identifier of the document.\n\tID string `json:\"id\"`\n\t// Content is the content of the document.\n\tContent string `json:\"content\"`\n\t// MetaData is the metadata of the document, can be used to store extra information.\n\tMetaData map[string]any `json:\"meta_data\"`\n}\n\n// String returns the content of the document.\nfunc (d *Document) String() string {\n\treturn d.Content\n}\n\n// WithSubIndexes sets the sub-indexes on the document metadata and returns the\n// document for chaining. Sub-indexes let an Indexer route a document into\n// multiple logical partitions of a vector store simultaneously.\n// Use [Document.SubIndexes] to retrieve them.\nfunc (d *Document) WithSubIndexes(indexes []string) *Document {\n\tif d.MetaData == nil {\n\t\td.MetaData = make(map[string]any)\n\t}\n\n\td.MetaData[docMetaDataKeySubIndexes] = indexes\n\n\treturn d\n}\n\n// SubIndexes returns the sub indexes of the document.\n// can use doc.WithSubIndexes() to set the sub indexes.\nfunc (d *Document) SubIndexes() []string {\n\tif d.MetaData == nil {\n\t\treturn nil\n\t}\n\n\tindexes, ok := d.MetaData[docMetaDataKeySubIndexes].([]string)\n\tif ok {\n\t\treturn indexes\n\t}\n\n\treturn nil\n}\n\n// WithScore sets the relevance score on the document, typically written by a\n// Retriever after ranking results. A higher score means higher relevance.\n// Note: [retriever.WithScoreThreshold] filters by this value, not sort order.\n// Use [Document.Score] to retrieve it.\nfunc (d *Document) WithScore(score float64) *Document {\n\tif d.MetaData == nil {\n\t\td.MetaData = make(map[string]any)\n\t}\n\n\td.MetaData[docMetaDataKeyScore] = score\n\n\treturn d\n}\n\n// Score returns the score of the document.\n// can use doc.WithScore() to set the score.\nfunc (d *Document) Score() float64 {\n\tif d.MetaData == nil {\n\t\treturn 0\n\t}\n\n\tscore, ok := d.MetaData[docMetaDataKeyScore].(float64)\n\tif ok {\n\t\treturn score\n\t}\n\n\treturn 0\n}\n\n// WithExtraInfo sets the extra info of the document.\n// can use doc.ExtraInfo() to get the extra info.\nfunc (d *Document) WithExtraInfo(extraInfo string) *Document {\n\tif d.MetaData == nil {\n\t\td.MetaData = make(map[string]any)\n\t}\n\n\td.MetaData[docMetaDataKeyExtraInfo] = extraInfo\n\n\treturn d\n}\n\n// ExtraInfo returns the extra info of the document.\n// can use doc.WithExtraInfo() to set the extra info.\nfunc (d *Document) ExtraInfo() string {\n\tif d.MetaData == nil {\n\t\treturn \"\"\n\t}\n\n\textraInfo, ok := d.MetaData[docMetaDataKeyExtraInfo].(string)\n\tif ok {\n\t\treturn extraInfo\n\t}\n\n\treturn \"\"\n}\n\n// WithDSLInfo attaches a domain-specific-language query description to the\n// document. This is consumed by Retriever implementations that support\n// structured queries (e.g., filter expressions) alongside vector search.\n// Use [Document.DSLInfo] to retrieve it.\nfunc (d *Document) WithDSLInfo(dslInfo map[string]any) *Document {\n\tif d.MetaData == nil {\n\t\td.MetaData = make(map[string]any)\n\t}\n\n\td.MetaData[docMetaDataKeyDSL] = dslInfo\n\n\treturn d\n}\n\n// DSLInfo returns the dsl info of the document.\n// can use doc.WithDSLInfo() to set the dsl info.\nfunc (d *Document) DSLInfo() map[string]any {\n\tif d.MetaData == nil {\n\t\treturn nil\n\t}\n\n\tdslInfo, ok := d.MetaData[docMetaDataKeyDSL].(map[string]any)\n\tif ok {\n\t\treturn dslInfo\n\t}\n\n\treturn nil\n}\n\n// WithDenseVector sets the dense vector of the document.\n// can use doc.DenseVector() to get the dense vector.\nfunc (d *Document) WithDenseVector(vector []float64) *Document {\n\tif d.MetaData == nil {\n\t\td.MetaData = make(map[string]any)\n\t}\n\n\td.MetaData[docMetaDataKeyDenseVector] = vector\n\n\treturn d\n}\n\n// DenseVector returns the dense vector of the document.\n// can use doc.WithDenseVector() to set the dense vector.\nfunc (d *Document) DenseVector() []float64 {\n\tif d.MetaData == nil {\n\t\treturn nil\n\t}\n\n\tvector, ok := d.MetaData[docMetaDataKeyDenseVector].([]float64)\n\tif ok {\n\t\treturn vector\n\t}\n\n\treturn nil\n}\n\n// WithSparseVector sets the sparse vector of the document, key indices -> value vector.\n// can use doc.SparseVector() to get the sparse vector.\nfunc (d *Document) WithSparseVector(sparse map[int]float64) *Document {\n\tif d.MetaData == nil {\n\t\td.MetaData = make(map[string]any)\n\t}\n\n\td.MetaData[docMetaDataKeySparseVector] = sparse\n\n\treturn d\n}\n\n// SparseVector returns the sparse vector of the document, key indices -> value vector.\n// can use doc.WithSparseVector() to set the sparse vector.\nfunc (d *Document) SparseVector() map[int]float64 {\n\tif d.MetaData == nil {\n\t\treturn nil\n\t}\n\n\tsparse, ok := d.MetaData[docMetaDataKeySparseVector].(map[int]float64)\n\tif ok {\n\t\treturn sparse\n\t}\n\n\treturn nil\n}\n"
  },
  {
    "path": "schema/document_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"testing\"\n\n\t\"github.com/smartystreets/goconvey/convey\"\n)\n\nfunc TestDocument(t *testing.T) {\n\tconvey.Convey(\"test document\", t, func() {\n\t\tvar (\n\t\t\tsubIndexes = []string{\"hello\", \"bye\"}\n\t\t\tscore      = 1.1\n\t\t\textraInfo  = \"asd\"\n\t\t\tdslInfo    = map[string]any{\"hello\": true}\n\t\t\tvector     = []float64{1.1, 2.2}\n\t\t)\n\n\t\td := &Document{\n\t\t\tID:       \"asd\",\n\t\t\tContent:  \"qwe\",\n\t\t\tMetaData: nil,\n\t\t}\n\n\t\td.WithSubIndexes(subIndexes).\n\t\t\tWithDenseVector(vector).\n\t\t\tWithScore(score).\n\t\t\tWithExtraInfo(extraInfo).\n\t\t\tWithDSLInfo(dslInfo)\n\n\t\tconvey.So(d.SubIndexes(), convey.ShouldEqual, subIndexes)\n\t\tconvey.So(d.Score(), convey.ShouldEqual, score)\n\t\tconvey.So(d.ExtraInfo(), convey.ShouldEqual, extraInfo)\n\t\tconvey.So(d.DSLInfo(), convey.ShouldEqual, dslInfo)\n\t\tconvey.So(d.DenseVector(), convey.ShouldEqual, vector)\n\t})\n}\n"
  },
  {
    "path": "schema/message.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"io\"\n\t\"reflect\"\n\t\"sort\"\n\t\"strings\"\n\t\"sync\"\n\t\"text/template\"\n\n\t\"github.com/nikolalohinski/gonja\"\n\t\"github.com/nikolalohinski/gonja/config\"\n\t\"github.com/nikolalohinski/gonja/nodes\"\n\t\"github.com/nikolalohinski/gonja/parser\"\n\t\"github.com/slongfield/pyfmt\"\n\n\t\"github.com/cloudwego/eino/internal\"\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\nfunc init() {\n\tinternal.RegisterStreamChunkConcatFunc(ConcatMessages)\n\tinternal.RegisterStreamChunkConcatFunc(ConcatMessageArray)\n\n\tinternal.RegisterStreamChunkConcatFunc(ConcatToolResults)\n}\n\n// ConcatMessageArray merges aligned slices of messages into a single slice,\n// concatenating messages at the same index across the input arrays.\nfunc ConcatMessageArray(mas [][]*Message) ([]*Message, error) {\n\tarrayLen := len(mas[0])\n\n\tret := make([]*Message, arrayLen)\n\tslicesToConcat := make([][]*Message, arrayLen)\n\n\tfor _, ma := range mas {\n\t\tif len(ma) != arrayLen {\n\t\t\treturn nil, fmt.Errorf(\"unexpected array length. \"+\n\t\t\t\t\"Got %d, expected %d\", len(ma), arrayLen)\n\t\t}\n\n\t\tfor i := 0; i < arrayLen; i++ {\n\t\t\tm := ma[i]\n\t\t\tif m != nil {\n\t\t\t\tslicesToConcat[i] = append(slicesToConcat[i], m)\n\t\t\t}\n\t\t}\n\t}\n\n\tfor i, slice := range slicesToConcat {\n\t\tif len(slice) == 0 {\n\t\t\tret[i] = nil\n\t\t} else if len(slice) == 1 {\n\t\t\tret[i] = slice[0]\n\t\t} else {\n\t\t\tcm, err := ConcatMessages(slice)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\tret[i] = cm\n\t\t}\n\t}\n\n\treturn ret, nil\n}\n\n// FormatType used by MessageTemplate.Format\ntype FormatType uint8\n\nconst (\n\t// FString Supported by pyfmt(github.com/slongfield/pyfmt), which is an implementation of https://peps.python.org/pep-3101/.\n\tFString FormatType = 0\n\t// GoTemplate https://pkg.go.dev/text/template.\n\tGoTemplate FormatType = 1\n\t// Jinja2 Supported by gonja(github.com/nikolalohinski/gonja), which is a implementation of https://jinja.palletsprojects.com/en/3.1.x/templates/.\n\tJinja2 FormatType = 2\n)\n\n// RoleType is the type of the role of a message.\ntype RoleType string\n\nconst (\n\t// Assistant is the role of an assistant, means the message is returned by ChatModel.\n\tAssistant RoleType = \"assistant\"\n\t// User is the role of a user, means the message is a user message.\n\tUser RoleType = \"user\"\n\t// System is the role of a system, means the message is a system message.\n\tSystem RoleType = \"system\"\n\t// Tool is the role of a tool, means the message is a tool call output.\n\tTool RoleType = \"tool\"\n)\n\n// FunctionCall is the function call in a message.\n// It's used in Assistant Message.\ntype FunctionCall struct {\n\t// Name is the name of the function to call, it can be used to identify the specific function.\n\tName string `json:\"name,omitempty\"`\n\t// Arguments is the arguments to call the function with, in JSON format.\n\tArguments string `json:\"arguments,omitempty\"`\n}\n\n// ToolCall is the tool call in a message.\n// It's used in Assistant Message when there are tool calls should be made.\ntype ToolCall struct {\n\t// Index is used when there are multiple tool calls in a message.\n\t// In stream mode, it's used to identify the chunk of the tool call for merging.\n\tIndex *int `json:\"index,omitempty\"`\n\t// ID is the id of the tool call, it can be used to identify the specific tool call.\n\tID string `json:\"id\"`\n\t// Type is the type of the tool call, default is \"function\".\n\tType string `json:\"type\"`\n\t// Function is the function call to be made.\n\tFunction FunctionCall `json:\"function\"`\n\n\t// Extra is used to store extra information for the tool call.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// ImageURLDetail is the detail of the image url.\ntype ImageURLDetail string\n\nconst (\n\t// ImageURLDetailHigh means the high quality image url.\n\tImageURLDetailHigh ImageURLDetail = \"high\"\n\t// ImageURLDetailLow means the low quality image url.\n\tImageURLDetailLow ImageURLDetail = \"low\"\n\t// ImageURLDetailAuto means the auto quality image url.\n\tImageURLDetailAuto ImageURLDetail = \"auto\"\n)\n\n// MessagePartCommon represents the common abstract components for input and output of multi-modal types.\ntype MessagePartCommon struct {\n\t// URL is primarily used for HTTP or HTTPS access links.\n\t// For data in the format 'data:[<mediatype>][;base64],<data>' (the 'data' URL Schema of RFC-2397 (https://www.rfc-editor.org/rfc/rfc2397)),\n\t// it is recommended to use Base64Data and MIMEType fields separately instead.\n\tURL *string `json:\"url,omitempty\"`\n\n\t// Base64Data represents the binary data in Base64 encoded string format.\n\tBase64Data *string `json:\"base64data,omitempty\"`\n\n\t// MIMEType is the mime type , eg.\"image/png\",\"\"audio/wav\" etc.\n\tMIMEType string `json:\"mime_type,omitempty\"`\n\n\t// Deprecated: Use MessageOutputPart.Extra or MessageInputPart.Extra to set additional metadata instead.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// MessageInputImage is used to represent an image part in message.\n// Choose either URL or Base64Data.\ntype MessageInputImage struct {\n\tMessagePartCommon\n\n\t// Detail is the quality of the image url.\n\tDetail ImageURLDetail `json:\"detail,omitempty\"`\n}\n\n// MessageInputAudio is used to represent an audio part in message.\n// Choose either URL or Base64Data.\ntype MessageInputAudio struct {\n\tMessagePartCommon\n}\n\n// MessageInputVideo is used to represent a video part in message.\n// Choose either URL or Base64Data.\ntype MessageInputVideo struct {\n\tMessagePartCommon\n}\n\n// MessageInputFile is used to represent a file part in message.\n// Choose either URL or Base64Data.\ntype MessageInputFile struct {\n\tMessagePartCommon\n\n\t// Name represents the filename.\n\t// Optional.\n\tName string `json:\"name,omitempty\"`\n}\n\n// MessageInputPart represents the input part of message.\ntype MessageInputPart struct {\n\tType ChatMessagePartType `json:\"type\"`\n\n\tText string `json:\"text,omitempty\"`\n\n\t// Image is the image input of the part, it's used when Type is \"image_url\".\n\tImage *MessageInputImage `json:\"image,omitempty\"`\n\n\t// Audio  is the audio input of the part, it's used when Type is \"audio_url\".\n\tAudio *MessageInputAudio `json:\"audio,omitempty\"`\n\n\t// Video is the video input of the part, it's used when Type is \"video_url\".\n\tVideo *MessageInputVideo `json:\"video,omitempty\"`\n\n\t// File is the file input of the part, it's used when Type is \"file_url\".\n\tFile *MessageInputFile `json:\"file,omitempty\"`\n\n\t// Extra is used to store extra information.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// MessageOutputImage is used to represent an image part in message.\ntype MessageOutputImage struct {\n\tMessagePartCommon\n}\n\n// MessageOutputAudio is used to represent an audio part in message.\ntype MessageOutputAudio struct {\n\tMessagePartCommon\n}\n\n// MessageOutputVideo is used to represent a video part in message.\ntype MessageOutputVideo struct {\n\tMessagePartCommon\n}\n\n// MessageOutputReasoning represents the reasoning content generated by reasoning models.\n// Some models produce reasoning steps before generating the final response.\n// This struct captures that reasoning output.\ntype MessageOutputReasoning struct {\n\t// Text is either the thought summary or the raw reasoning text itself.\n\tText string `json:\"text,omitempty\"`\n\n\t// Signature contains encrypted reasoning tokens.\n\t// Required by some models when passing reasoning context back in subsequent requests.\n\tSignature string `json:\"signature,omitempty\"`\n}\n\n// MessageStreamingMeta contains metadata for streaming responses.\n// It is used to track position of part when the model outputs multiple parts in a single response.\ntype MessageStreamingMeta struct {\n\t// Index specifies the index position of this part in the final response.\n\t// This is useful for reassembling multiple reasoning/content parts in correct order.\n\tIndex int `json:\"index,omitempty\"`\n}\n\n// MessageOutputPart represents a part of an assistant-generated message.\n// It can contain text, or multimedia content like images, audio, or video.\ntype MessageOutputPart struct {\n\t// Type is the type of the part, e.g. \"text\", \"image_url\", \"audio_url\", \"video_url\".\n\tType ChatMessagePartType `json:\"type\"`\n\n\t// Text is the text of the part, it's used when Type is \"text\".\n\tText string `json:\"text,omitempty\"`\n\n\t// Image is the image output of the part, used when Type is ChatMessagePartTypeImageURL.\n\tImage *MessageOutputImage `json:\"image,omitempty\"`\n\n\t// Audio is the audio output of the part, used when Type is ChatMessagePartTypeAudioURL.\n\tAudio *MessageOutputAudio `json:\"audio,omitempty\"`\n\n\t// Video is the video output of the part, used when Type is ChatMessagePartTypeVideoURL.\n\tVideo *MessageOutputVideo `json:\"video,omitempty\"`\n\n\t// Reasoning contains the reasoning content generated by the model.\n\t// Used when Type is ChatMessagePartTypeReasoning.\n\tReasoning *MessageOutputReasoning `json:\"reasoning,omitempty\"`\n\n\t// Extra is used to store extra information.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n\n\t// StreamingMeta contains metadata for streaming responses.\n\t// This field is typically used at runtime and not serialized.\n\tStreamingMeta *MessageStreamingMeta `json:\"-\"`\n}\n\n// ToolPartType defines the type of content in a tool output part.\n// It is used to distinguish between different types of multimodal content returned by tools.\ntype ToolPartType string\n\nconst (\n\t// ToolPartTypeText means the part is a text.\n\tToolPartTypeText ToolPartType = \"text\"\n\n\t// ToolPartTypeImage means the part is an image url.\n\tToolPartTypeImage ToolPartType = \"image\"\n\n\t// ToolPartTypeAudio means the part is an audio url.\n\tToolPartTypeAudio ToolPartType = \"audio\"\n\n\t// ToolPartTypeVideo means the part is a video url.\n\tToolPartTypeVideo ToolPartType = \"video\"\n\n\t// ToolPartTypeFile means the part is a file url.\n\tToolPartTypeFile ToolPartType = \"file\"\n)\n\n// ToolOutputImage represents an image in tool output.\n// It contains URL or Base64-encoded data along with MIME type information.\ntype ToolOutputImage struct {\n\tMessagePartCommon\n}\n\n// ToolOutputAudio represents an audio file in tool output.\n// It contains URL or Base64-encoded data along with MIME type information.\ntype ToolOutputAudio struct {\n\tMessagePartCommon\n}\n\n// ToolOutputVideo represents a video file in tool output.\n// It contains URL or Base64-encoded data along with MIME type information.\ntype ToolOutputVideo struct {\n\tMessagePartCommon\n}\n\n// ToolOutputFile represents a generic file in tool output.\n// It contains URL or Base64-encoded data along with MIME type information.\ntype ToolOutputFile struct {\n\tMessagePartCommon\n}\n\n// ToolOutputPart represents a part of tool execution output.\n// It supports streaming scenarios through the Index field for chunk merging.\ntype ToolOutputPart struct {\n\n\t// Type is the type of the part, e.g., \"text\", \"image_url\", \"audio_url\", \"video_url\".\n\tType ToolPartType `json:\"type\"`\n\n\t// Text is the text content, used when Type is \"text\".\n\tText string `json:\"text,omitempty\"`\n\n\t// Image is the image content, used when Type is ToolPartTypeImage.\n\tImage *ToolOutputImage `json:\"image,omitempty\"`\n\n\t// Audio is the audio content, used when Type is ToolPartTypeAudio.\n\tAudio *ToolOutputAudio `json:\"audio,omitempty\"`\n\n\t// Video is the video content, used when Type is ToolPartTypeVideo.\n\tVideo *ToolOutputVideo `json:\"video,omitempty\"`\n\n\t// File is the file content, used when Type is ToolPartTypeFile.\n\tFile *ToolOutputFile `json:\"file,omitempty\"`\n\n\t// Extra is used to store extra information.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// ToolArgument contains the input information for a tool call.\n// It is used to pass tool call arguments to enhanced tools.\ntype ToolArgument struct {\n\t// Text contains the arguments for the tool call in JSON format.\n\tText string `json:\"text,omitempty\"`\n}\n\n// ToolResult represents the structured multimodal output from a tool execution.\n// It is used when a tool needs to return more than just a simple string,\n// such as images, files, or other structured data.\ntype ToolResult struct {\n\t// Parts contains the multimodal output parts. Each part can be a different\n\t// type of content, like text, an image, or a file.\n\tParts []ToolOutputPart `json:\"parts,omitempty\"`\n}\n\nfunc convToolOutputPartToMessageInputPart(toolPart ToolOutputPart) (MessageInputPart, error) {\n\tswitch toolPart.Type {\n\tcase ToolPartTypeText:\n\t\treturn MessageInputPart{\n\t\t\tType:  ChatMessagePartTypeText,\n\t\t\tText:  toolPart.Text,\n\t\t\tExtra: toolPart.Extra,\n\t\t}, nil\n\tcase ToolPartTypeImage:\n\t\tif toolPart.Image == nil {\n\t\t\treturn MessageInputPart{}, fmt.Errorf(\"image content is nil for tool part type %v\", toolPart.Type)\n\t\t}\n\t\treturn MessageInputPart{\n\t\t\tType:  ChatMessagePartTypeImageURL,\n\t\t\tImage: &MessageInputImage{MessagePartCommon: toolPart.Image.MessagePartCommon},\n\t\t\tExtra: toolPart.Extra,\n\t\t}, nil\n\tcase ToolPartTypeAudio:\n\t\tif toolPart.Audio == nil {\n\t\t\treturn MessageInputPart{}, fmt.Errorf(\"audio content is nil for tool part type %v\", toolPart.Type)\n\t\t}\n\t\treturn MessageInputPart{\n\t\t\tType:  ChatMessagePartTypeAudioURL,\n\t\t\tAudio: &MessageInputAudio{MessagePartCommon: toolPart.Audio.MessagePartCommon},\n\t\t\tExtra: toolPart.Extra,\n\t\t}, nil\n\tcase ToolPartTypeVideo:\n\t\tif toolPart.Video == nil {\n\t\t\treturn MessageInputPart{}, fmt.Errorf(\"video content is nil for tool part type %v\", toolPart.Type)\n\t\t}\n\t\treturn MessageInputPart{\n\t\t\tType:  ChatMessagePartTypeVideoURL,\n\t\t\tVideo: &MessageInputVideo{MessagePartCommon: toolPart.Video.MessagePartCommon},\n\t\t\tExtra: toolPart.Extra,\n\t\t}, nil\n\tcase ToolPartTypeFile:\n\t\tif toolPart.File == nil {\n\t\t\treturn MessageInputPart{}, fmt.Errorf(\"file content is nil for tool part type %v\", toolPart.Type)\n\t\t}\n\t\treturn MessageInputPart{\n\t\t\tType:  ChatMessagePartTypeFileURL,\n\t\t\tFile:  &MessageInputFile{MessagePartCommon: toolPart.File.MessagePartCommon},\n\t\t\tExtra: toolPart.Extra,\n\t\t}, nil\n\tdefault:\n\t\treturn MessageInputPart{}, fmt.Errorf(\"unknown tool part type: %v\", toolPart.Type)\n\t}\n}\n\n// ToMessageInputParts converts ToolOutputPart slice to MessageInputPart slice.\n// This is used when passing tool results as input to the model.\n//\n// Parameters:\n//   - None (method receiver is *ToolResult)\n//\n// Returns:\n//   - []MessageInputPart: The converted message input parts that can be used in a Message.\n//   - error: An error if conversion fails due to unknown part types or nil content fields.\n//\n// Example:\n//\n//\ttoolResult := &schema.ToolResult{\n//\t    Parts: []schema.ToolOutputPart{\n//\t        {Type: schema.ToolPartTypeText, Text: \"Result text\"},\n//\t        {Type: schema.ToolPartTypeImage, Image: &schema.ToolOutputImage{...}},\n//\t    },\n//\t}\n//\tinputParts, err := toolResult.ToMessageInputParts()\nfunc (tr *ToolResult) ToMessageInputParts() ([]MessageInputPart, error) {\n\tif tr == nil || len(tr.Parts) == 0 {\n\t\treturn nil, nil\n\t}\n\tresult := make([]MessageInputPart, len(tr.Parts))\n\tfor i, part := range tr.Parts {\n\t\tvar err error\n\t\tresult[i], err = convToolOutputPartToMessageInputPart(part)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\treturn result, nil\n}\n\n// Deprecated: This struct is deprecated as the MultiContent field is deprecated.\n// For the image input part of the model, use MessageInputImage.\n// For the image output part of the model, use MessageOutputImage.\n// Choose either URL or URI.\n// If your model implementation supports it, URL could embed inline image data\n// as defined in RFC-2397.\ntype ChatMessageImageURL struct {\n\t// URL can either be a traditional URL or a special URL conforming to RFC-2397 (https://www.rfc-editor.org/rfc/rfc2397).\n\t// double check with model implementations for detailed instructions on how to use this.\n\tURL string `json:\"url,omitempty\"`\n\n\tURI string `json:\"uri,omitempty\"`\n\t// Detail is the quality of the image url.\n\tDetail ImageURLDetail `json:\"detail,omitempty\"`\n\n\t// MIMEType is the mime type of the image, eg. \"image/png\".\n\tMIMEType string `json:\"mime_type,omitempty\"`\n\t// Extra is used to store extra information for the image url.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// ChatMessagePartType is the type of the part in a chat message.\ntype ChatMessagePartType string\n\nconst (\n\t// ChatMessagePartTypeText means the part is a text.\n\tChatMessagePartTypeText ChatMessagePartType = \"text\"\n\t// ChatMessagePartTypeImageURL means the part is an image url.\n\tChatMessagePartTypeImageURL ChatMessagePartType = \"image_url\"\n\t// ChatMessagePartTypeAudioURL means the part is an audio url.\n\tChatMessagePartTypeAudioURL ChatMessagePartType = \"audio_url\"\n\t// ChatMessagePartTypeVideoURL means the part is a video url.\n\tChatMessagePartTypeVideoURL ChatMessagePartType = \"video_url\"\n\t// ChatMessagePartTypeFileURL means the part is a file url.\n\tChatMessagePartTypeFileURL ChatMessagePartType = \"file_url\"\n\t// ChatMessagePartTypeReasoning means the part is a reasoning block.\n\tChatMessagePartTypeReasoning ChatMessagePartType = \"reasoning\"\n)\n\n// Deprecated: This struct is deprecated as the MultiContent field is deprecated.\n// For the audio input part of the model, use MessageInputAudio.\n// For the audio output part of the model, use MessageOutputAudio.\n// Choose either URL or URI.\n// If supported, URL may embed inline audio data per RFC-2397.\ntype ChatMessageAudioURL struct {\n\t// URL can either be a traditional URL or a special URL conforming to RFC-2397 (https://www.rfc-editor.org/rfc/rfc2397).\n\t// double check with model implementations for detailed instructions on how to use this.\n\tURL string `json:\"url,omitempty\"`\n\tURI string `json:\"uri,omitempty\"`\n\n\t// MIMEType is the mime type of the audio, eg. \"audio/wav\" or \"audio/ogg\".\n\tMIMEType string `json:\"mime_type,omitempty\"`\n\t// Extra is used to store extra information for the audio url.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// Deprecated: This struct is deprecated as the MultiContent field is deprecated.\n// For the video input part of the model, use MessageInputVideo.\n// For the video output part of the model, use MessageOutputVideo.\n// Choose either URL or URI.\n// If supported, URL may embed inline video data per RFC-2397.\ntype ChatMessageVideoURL struct {\n\t// URL can either be a traditional URL or a special URL conforming to RFC-2397 (https://www.rfc-editor.org/rfc/rfc2397).\n\t// double check with model implementations for detailed instructions on how to use this.\n\tURL string `json:\"url,omitempty\"`\n\tURI string `json:\"uri,omitempty\"`\n\n\t// MIMEType is the mime type of the video, eg. \"video/mp4\".\n\tMIMEType string `json:\"mime_type,omitempty\"`\n\t// Extra is used to store extra information for the video url.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// Deprecated: This struct is deprecated as the MultiContent field is deprecated.\n// For the file input part of the model, use MessageInputFile.\n// Choose either URL or URI.\ntype ChatMessageFileURL struct {\n\tURL string `json:\"url,omitempty\"`\n\tURI string `json:\"uri,omitempty\"`\n\n\t// MIMEType is the mime type of the file, eg. \"application/pdf\", \"text/plain\".\n\tMIMEType string `json:\"mime_type,omitempty\"`\n\t// Name is the name of the file.\n\tName string `json:\"name,omitempty\"`\n\n\t// Extra is used to store extra information for the file url.\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// Deprecated: This struct is deprecated as the MultiContent field is deprecated.\n// For model input, use MessageInputPart. For model output, use MessageOutputPart.\ntype ChatMessagePart struct {\n\t// Type is the type of the part, eg. \"text\", \"image_url\", \"audio_url\", \"video_url\", \"file_url\".\n\tType ChatMessagePartType `json:\"type,omitempty\"`\n\n\t// Text is the text of the part, it's used when Type is \"text\".\n\tText string `json:\"text,omitempty\"`\n\n\t// ImageURL is the image url of the part, it's used when Type is \"image_url\".\n\tImageURL *ChatMessageImageURL `json:\"image_url,omitempty\"`\n\t// AudioURL is the audio url of the part, it's used when Type is \"audio_url\".\n\tAudioURL *ChatMessageAudioURL `json:\"audio_url,omitempty\"`\n\t// VideoURL is the video url of the part, it's used when Type is \"video_url\".\n\tVideoURL *ChatMessageVideoURL `json:\"video_url,omitempty\"`\n\t// FileURL is the file url of the part, it's used when Type is \"file_url\".\n\tFileURL *ChatMessageFileURL `json:\"file_url,omitempty\"`\n}\n\n// LogProbs is the top-level structure containing the log probability information.\ntype LogProbs struct {\n\t// Content is a list of message content tokens with log probability information.\n\tContent []LogProb `json:\"content\"`\n}\n\n// LogProb represents the probability information for a token.\ntype LogProb struct {\n\t// Token represents the text of the token, which is a contiguous sequence of characters\n\t// (e.g., a word, part of a word, or punctuation) as understood by the tokenization process used by the language model.\n\tToken string `json:\"token\"`\n\t// LogProb is the log probability of this token, if it is within the top 20 most likely tokens.\n\t// Otherwise, the value `-9999.0` is used to signify that the token is very unlikely.\n\tLogProb float64 `json:\"logprob\"`\n\t// Bytes is a list of integers representing the UTF-8 bytes representation of the token.\n\t// Useful in instances where characters are represented by multiple tokens and\n\t// their byte representations must be combined to generate the correct text\n\t// representation. Can be `null` if there is no bytes representation for the token.\n\tBytes []int64 `json:\"bytes,omitempty\"` // Omitting the field if it is null\n\t// TopLogProbs is a list of the most likely tokens and their log probability, at this token position.\n\t// In rare cases, there may be fewer than the number of requested top_logprobs returned.\n\tTopLogProbs []TopLogProb `json:\"top_logprobs\"`\n}\n\n// TopLogProb describes a likely token and its log probability at a position.\ntype TopLogProb struct {\n\t// Token represents the text of the token, which is a contiguous sequence of characters\n\t// (e.g., a word, part of a word, or punctuation) as understood by the tokenization process used by the language model.\n\tToken string `json:\"token\"`\n\t// LogProb is the log probability of this token, if it is within the top 20 most likely tokens.\n\t// Otherwise, the value `-9999.0` is used to signify that the token is very unlikely.\n\tLogProb float64 `json:\"logprob\"`\n\t// Bytes is a list of integers representing the UTF-8 bytes representation of the token.\n\t// Useful in instances where characters are represented by multiple tokens and\n\t// their byte representations must be combined to generate the correct text\n\t// representation. Can be `null` if there is no bytes representation for the token.\n\tBytes []int64 `json:\"bytes,omitempty\"`\n}\n\n// ResponseMeta collects meta information about a chat response.\ntype ResponseMeta struct {\n\t// FinishReason is the reason why the chat response is finished.\n\t// It's usually \"stop\", \"length\", \"tool_calls\", \"content_filter\", \"null\". This is defined by chat model implementation.\n\tFinishReason string `json:\"finish_reason,omitempty\"`\n\t// Usage is the token usage of the chat response, whether usage exists depends on whether the chat model implementation returns.\n\tUsage *TokenUsage `json:\"usage,omitempty\"`\n\t// LogProbs is Log probability information.\n\tLogProbs *LogProbs `json:\"logprobs,omitempty\"`\n}\n\n// Message denotes the data structure for model input and output, originating from either user input or model return.\n// It supports both text-only and multimodal content.\n//\n// For text-only input from a user, use the Content field:\n//\n//\t&schema.Message{\n//\t\tRole:    schema.User,\n//\t\tContent: \"What is the capital of France?\",\n//\t}\n//\n// For multimodal input from a user, use the UserInputMultiContent field.\n// This allows combining text with other media like images:\n//\n//\t&schema.Message{\n//\t\tRole: schema.User,\n//\t\tUserInputMultiContent: []schema.MessageInputPart{\n//\t\t\t{Type: schema.ChatMessagePartTypeText, Text: \"What is in this image?\"},\n//\t\t\t{Type: schema.ChatMessagePartTypeImageURL, Image: &schema.MessageInputImage{\n//\t\t\t\tMessagePartCommon: schema.MessagePartCommon{\n//\t\t\t\t\tURL: toPtr(\"https://example.com/cat.jpg\"),\n//\t\t\t\t},\n//\t\t\t\tDetail: schema.ImageURLDetailHigh,\n//\t\t\t}},\n//\t\t},\n//\t}\n//\n// When the model returns multimodal content, it is available in the AssistantGenMultiContent field:\n//\n//\t&schema.Message{\n//\t\tRole: schema.Assistant,\n//\t\tAssistantGenMultiContent: []schema.MessageOutputPart{\n//\t\t\t{Type: schema.ChatMessagePartTypeText, Text: \"Here is the generated image:\"},\n//\t\t\t{Type: schema.ChatMessagePartTypeImage, Image: &schema.MessageOutputImage{\n//\t\t\t\tMessagePartCommon: schema.MessagePartCommon{\n//\t\t\t\t\tBase64Data: toPtr(\"base64_image_binary\"),\n//\t\t\t\t\tMIMEType:   \"image/png\",\n//\t\t\t\t},\n//\t\t\t}},\n//\t\t},\n//\t}\ntype Message struct {\n\tRole RoleType `json:\"role\"`\n\n\t// Content is for user text input and model text output.\n\tContent string `json:\"content\"`\n\n\t// if MultiContent is not empty, use this instead of Content\n\t// if MultiContent is empty, use Content\n\t// Deprecated: Use UserInputMultiContent for user multimodal inputs and AssistantGenMultiContent for model multimodal outputs.\n\tMultiContent []ChatMessagePart `json:\"multi_content,omitempty\"`\n\n\t// UserInputMultiContent passes multimodal content provided by the user to the model.\n\tUserInputMultiContent []MessageInputPart `json:\"user_input_multi_content,omitempty\"`\n\n\t// AssistantGenMultiContent is for receiving multimodal output from the model.\n\tAssistantGenMultiContent []MessageOutputPart `json:\"assistant_output_multi_content,omitempty\"`\n\n\tName string `json:\"name,omitempty\"`\n\n\t// only for AssistantMessage\n\tToolCalls []ToolCall `json:\"tool_calls,omitempty\"`\n\n\t// only for ToolMessage\n\tToolCallID string `json:\"tool_call_id,omitempty\"`\n\t// only for ToolMessage\n\tToolName string `json:\"tool_name,omitempty\"`\n\n\tResponseMeta *ResponseMeta `json:\"response_meta,omitempty\"`\n\n\t// ReasoningContent is the thinking process of the model, which will be included when the model returns reasoning content.\n\tReasoningContent string `json:\"reasoning_content,omitempty\"`\n\n\t// customized information for model implementation\n\tExtra map[string]any `json:\"extra,omitempty\"`\n}\n\n// TokenUsage Represents the token usage of chat model request.\ntype TokenUsage struct {\n\t// PromptTokens is the number of prompt tokens, including all the input tokens of this request.\n\tPromptTokens int `json:\"prompt_tokens\"`\n\t// PromptTokenDetails is a breakdown of the prompt tokens.\n\tPromptTokenDetails PromptTokenDetails `json:\"prompt_token_details\"`\n\t// CompletionTokens is the number of completion tokens.\n\tCompletionTokens int `json:\"completion_tokens\"`\n\t// TotalTokens is the total number of tokens.\n\tTotalTokens int `json:\"total_tokens\"`\n\t// CompletionTokensDetails is breakdown of completion tokens.\n\tCompletionTokensDetails CompletionTokensDetails `json:\"completion_token_details\"`\n}\n\ntype CompletionTokensDetails struct {\n\t// ReasoningTokens tokens generated by the model for reasoning.\n\t// This is currently supported by OpenAI, Gemini, ARK and Qwen  chat models.\n\t// For other models, this field will be 0.\n\tReasoningTokens int `json:\"reasoning_tokens,omitempty\"`\n}\n\n// PromptTokenDetails provides a breakdown of prompt token usage.\ntype PromptTokenDetails struct {\n\t// Cached tokens present in the prompt.\n\tCachedTokens int `json:\"cached_tokens\"`\n}\n\nvar _ MessagesTemplate = &Message{}\nvar _ MessagesTemplate = MessagesPlaceholder(\"\", false)\n\n// MessagesTemplate is the interface for messages template.\n// It's used to render a template to a list of messages.\n// e.g.\n//\n//\tchatTemplate := prompt.FromMessages(\n//\t\tschema.SystemMessage(\"you are eino helper\"),\n//\t\tschema.MessagesPlaceholder(\"history\", false), // <= this will use the value of \"history\" in params\n//\t)\n//\tmsgs, err := chatTemplate.Format(ctx, params)\ntype MessagesTemplate interface {\n\tFormat(ctx context.Context, vs map[string]any, formatType FormatType) ([]*Message, error)\n}\n\ntype messagesPlaceholder struct {\n\tkey      string\n\toptional bool\n}\n\n// MessagesPlaceholder can render a placeholder to a list of messages in params.\n// e.g.\n//\n//\tplaceholder := MessagesPlaceholder(\"history\", false)\n//\tparams := map[string]any{\n//\t\t\"history\": []*schema.Message{{Role: \"user\", Content: \"what is eino?\"}, {Role: \"assistant\", Content: \"eino is a great freamwork to build llm apps\"}},\n//\t\t\"query\": \"how to use eino?\",\n//\t}\n//\tchatTemplate := chatTpl := prompt.FromMessages(\n//\t\tschema.SystemMessage(\"you are eino helper\"),\n//\t\tschema.MessagesPlaceholder(\"history\", false), // <= this will use the value of \"history\" in params\n//\t)\n//\tmsgs, err := chatTemplate.Format(ctx, params)\nfunc MessagesPlaceholder(key string, optional bool) MessagesTemplate {\n\treturn &messagesPlaceholder{\n\t\tkey:      key,\n\t\toptional: optional,\n\t}\n}\n\n// Format just return the messages of specified key.\n// because it's a placeholder.\n// e.g.\n//\n//\tplaceholder := MessagesPlaceholder(\"history\", false)\n//\tparams := map[string]any{\n//\t\t\"history\": []*schema.Message{{Role: \"user\", Content: \"what is eino?\"}, {Role: \"assistant\", Content: \"eino is a great freamwork to build llm apps\"}},\n//\t\t\"query\": \"how to use eino?\",\n//\t}\n//\tmsgs, err := placeholder.Format(ctx, params) // <= this will return the value of \"history\" in params\nfunc (p *messagesPlaceholder) Format(_ context.Context, vs map[string]any, _ FormatType) ([]*Message, error) {\n\tv, ok := vs[p.key]\n\tif !ok {\n\t\tif p.optional {\n\t\t\treturn []*Message{}, nil\n\t\t}\n\n\t\treturn nil, fmt.Errorf(\"message placeholder format: %s not found\", p.key)\n\t}\n\n\tmsgs, ok := v.([]*Message)\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"only messages can be used to format message placeholder, key: %v, actual type: %v\", p.key, reflect.TypeOf(v))\n\t}\n\n\treturn msgs, nil\n}\n\nfunc formatContent(content string, vs map[string]any, formatType FormatType) (string, error) {\n\tswitch formatType {\n\tcase FString:\n\t\treturn pyfmt.Fmt(content, vs)\n\tcase GoTemplate:\n\t\tparsedTmpl, err := template.New(\"template\").\n\t\t\tOption(\"missingkey=error\").\n\t\t\tParse(content)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tsb := new(strings.Builder)\n\t\terr = parsedTmpl.Execute(sb, vs)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn sb.String(), nil\n\tcase Jinja2:\n\t\tenv, err := getJinjaEnv()\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\ttpl, err := env.FromString(content)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tout, err := tpl.Execute(vs)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn out, nil\n\tdefault:\n\t\treturn \"\", fmt.Errorf(\"unknown format type: %v\", formatType)\n\t}\n}\n\n// Format returns the messages after rendering by the given formatType.\n// e.g.\n//\n//\tmsg := schema.UserMessage(\"hello world, {name}\")\n//\tmsgs, err := msg.Format(ctx, map[string]any{\"name\": \"eino\"}, schema.FString) // <= this will render the content of msg by pyfmt\n//\t// msgs[0].Content will be \"hello world, eino\"\nfunc (m *Message) Format(_ context.Context, vs map[string]any, formatType FormatType) ([]*Message, error) {\n\tc, err := formatContent(m.Content, vs, formatType)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tcopied := *m\n\tcopied.Content = c\n\n\tif len(m.MultiContent) > 0 {\n\t\tcopied.MultiContent, err = formatMultiContent(m.MultiContent, vs, formatType)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tif len(m.UserInputMultiContent) > 0 {\n\t\tcopied.UserInputMultiContent, err = formatUserInputMultiContent(m.UserInputMultiContent, vs, formatType)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\treturn []*Message{&copied}, nil\n}\n\nfunc formatMultiContent(multiContent []ChatMessagePart, vs map[string]any, formatType FormatType) ([]ChatMessagePart, error) {\n\tcopiedMC := make([]ChatMessagePart, len(multiContent))\n\tcopy(copiedMC, multiContent)\n\n\tfor i, mc := range copiedMC {\n\t\tswitch mc.Type {\n\t\tcase ChatMessagePartTypeText:\n\t\t\tnmc, err := formatContent(mc.Text, vs, formatType)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tcopiedMC[i].Text = nmc\n\t\tcase ChatMessagePartTypeImageURL:\n\t\t\tif mc.ImageURL == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\turl, err := formatContent(mc.ImageURL.URL, vs, formatType)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tcopiedMC[i].ImageURL.URL = url\n\t\tcase ChatMessagePartTypeAudioURL:\n\t\t\tif mc.AudioURL == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\turl, err := formatContent(mc.AudioURL.URL, vs, formatType)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tcopiedMC[i].AudioURL.URL = url\n\t\tcase ChatMessagePartTypeVideoURL:\n\t\t\tif mc.VideoURL == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\turl, err := formatContent(mc.VideoURL.URL, vs, formatType)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tcopiedMC[i].VideoURL.URL = url\n\t\tcase ChatMessagePartTypeFileURL:\n\t\t\tif mc.FileURL == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\turl, err := formatContent(mc.FileURL.URL, vs, formatType)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tcopiedMC[i].FileURL.URL = url\n\t\t}\n\t}\n\n\treturn copiedMC, nil\n}\n\nfunc formatUserInputMultiContent(userInputMultiContent []MessageInputPart, vs map[string]any, formatType FormatType) ([]MessageInputPart, error) {\n\tcopiedUIMC := make([]MessageInputPart, len(userInputMultiContent))\n\tcopy(copiedUIMC, userInputMultiContent)\n\n\tfor i, uimc := range copiedUIMC {\n\t\tswitch uimc.Type {\n\t\tcase ChatMessagePartTypeText:\n\t\t\ttext, err := formatContent(uimc.Text, vs, formatType)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tcopiedUIMC[i].Text = text\n\t\tcase ChatMessagePartTypeImageURL:\n\t\t\tif uimc.Image == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif uimc.Image.URL != nil && *uimc.Image.URL != \"\" {\n\t\t\t\turl, err := formatContent(*uimc.Image.URL, vs, formatType)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcopiedUIMC[i].Image.URL = &url\n\t\t\t}\n\t\t\tif uimc.Image.Base64Data != nil && *uimc.Image.Base64Data != \"\" {\n\t\t\t\tbase64data, err := formatContent(*uimc.Image.Base64Data, vs, formatType)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcopiedUIMC[i].Image.Base64Data = &base64data\n\t\t\t}\n\t\tcase ChatMessagePartTypeAudioURL:\n\t\t\tif uimc.Audio == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif uimc.Audio.URL != nil && *uimc.Audio.URL != \"\" {\n\t\t\t\turl, err := formatContent(*uimc.Audio.URL, vs, formatType)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcopiedUIMC[i].Audio.URL = &url\n\t\t\t}\n\t\t\tif uimc.Audio.Base64Data != nil && *uimc.Audio.Base64Data != \"\" {\n\t\t\t\tbase64data, err := formatContent(*uimc.Audio.Base64Data, vs, formatType)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcopiedUIMC[i].Audio.Base64Data = &base64data\n\t\t\t}\n\t\tcase ChatMessagePartTypeVideoURL:\n\t\t\tif uimc.Video == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif uimc.Video.URL != nil && *uimc.Video.URL != \"\" {\n\t\t\t\turl, err := formatContent(*uimc.Video.URL, vs, formatType)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcopiedUIMC[i].Video.URL = &url\n\t\t\t}\n\t\t\tif uimc.Video.Base64Data != nil && *uimc.Video.Base64Data != \"\" {\n\t\t\t\tbase64data, err := formatContent(*uimc.Video.Base64Data, vs, formatType)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcopiedUIMC[i].Video.Base64Data = &base64data\n\t\t\t}\n\t\tcase ChatMessagePartTypeFileURL:\n\t\t\tif uimc.File == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif uimc.File.URL != nil && *uimc.File.URL != \"\" {\n\t\t\t\turl, err := formatContent(*uimc.File.URL, vs, formatType)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcopiedUIMC[i].File.URL = &url\n\t\t\t}\n\t\t\tif uimc.File.Base64Data != nil && *uimc.File.Base64Data != \"\" {\n\t\t\t\tbase64data, err := formatContent(*uimc.File.Base64Data, vs, formatType)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tcopiedUIMC[i].File.Base64Data = &base64data\n\t\t\t}\n\t\t}\n\t}\n\n\treturn copiedUIMC, nil\n}\n\n// String returns the string representation of the message.\n// e.g.\n//\n//\tmsg := schema.UserMessage(\"hello world\")\n//\tfmt.Println(msg.String()) // Output will be: `user: hello world``\n//\n//\tmsg := schema.Message{\n//\t\tRole:    schema.Tool,\n//\t\tContent: \"{...}\",\n//\t\tToolCallID: \"callxxxx\"\n//\t}\n//\tfmt.Println(msg.String())\n//\tOutput will be:\n//\t\ttool: {...}\n//\t\tcall_id: callxxxx\nfunc (m *Message) String() string {\n\tsb := &strings.Builder{}\n\tsb.WriteString(fmt.Sprintf(\"%s: %s\", m.Role, m.Content))\n\n\tif len(m.UserInputMultiContent) > 0 {\n\t\tsb.WriteString(\"\\nuser_input_multi_content:\")\n\t\tfor i, part := range m.UserInputMultiContent {\n\t\t\tsb.WriteString(fmt.Sprintf(\"\\n  [%d] %s\", i, formatInputPart(part)))\n\t\t}\n\t}\n\n\tif len(m.AssistantGenMultiContent) > 0 {\n\t\tsb.WriteString(\"\\nassistant_gen_multi_content:\")\n\t\tfor i, part := range m.AssistantGenMultiContent {\n\t\t\tsb.WriteString(fmt.Sprintf(\"\\n  [%d] %s\", i, formatOutputPart(part)))\n\t\t}\n\t}\n\n\tif len(m.MultiContent) > 0 {\n\t\tsb.WriteString(\"\\nmulti_content:\")\n\t\tfor i, part := range m.MultiContent {\n\t\t\tsb.WriteString(fmt.Sprintf(\"\\n  [%d] %s\", i, formatChatMessagePart(part)))\n\t\t}\n\t}\n\n\tif len(m.ReasoningContent) > 0 {\n\t\tsb.WriteString(\"\\nreasoning content:\\n\")\n\t\tsb.WriteString(m.ReasoningContent)\n\t}\n\tif len(m.ToolCalls) > 0 {\n\t\tsb.WriteString(\"\\ntool_calls:\\n\")\n\t\tfor _, tc := range m.ToolCalls {\n\t\t\tif tc.Index != nil {\n\t\t\t\tsb.WriteString(fmt.Sprintf(\"index[%d]:\", *tc.Index))\n\t\t\t}\n\t\t\tsb.WriteString(fmt.Sprintf(\"%+v\\n\", tc))\n\t\t}\n\t}\n\tif m.ToolCallID != \"\" {\n\t\tsb.WriteString(fmt.Sprintf(\"\\ntool_call_id: %s\", m.ToolCallID))\n\t}\n\tif m.ToolName != \"\" {\n\t\tsb.WriteString(fmt.Sprintf(\"\\ntool_call_name: %s\", m.ToolName))\n\t}\n\tif m.ResponseMeta != nil {\n\t\tsb.WriteString(fmt.Sprintf(\"\\nfinish_reason: %s\", m.ResponseMeta.FinishReason))\n\t\tif m.ResponseMeta.Usage != nil {\n\t\t\tsb.WriteString(fmt.Sprintf(\"\\nusage: %v\", m.ResponseMeta.Usage))\n\t\t}\n\t}\n\n\treturn sb.String()\n}\n\nfunc formatInputPart(part MessageInputPart) string {\n\tswitch part.Type {\n\tcase ChatMessagePartTypeText:\n\t\treturn fmt.Sprintf(\"text: %s\", part.Text)\n\tcase ChatMessagePartTypeImageURL:\n\t\treturn fmt.Sprintf(\"image: %s\", formatMessageInputMedia(part.Image))\n\tcase ChatMessagePartTypeAudioURL:\n\t\treturn fmt.Sprintf(\"audio: %s\", formatMessageInputMedia(part.Audio))\n\tcase ChatMessagePartTypeVideoURL:\n\t\treturn fmt.Sprintf(\"video: %s\", formatMessageInputMedia(part.Video))\n\tcase ChatMessagePartTypeFileURL:\n\t\treturn fmt.Sprintf(\"file: %s\", formatMessageInputFile(part.File))\n\tdefault:\n\t\treturn fmt.Sprintf(\"unknown type: %s\", part.Type)\n\t}\n}\n\nfunc formatMessageInputMedia[T MessageInputImage | MessageInputAudio | MessageInputVideo](media *T) string {\n\tif media == nil {\n\t\treturn \"<nil>\"\n\t}\n\tvar parts []string\n\tswitch v := any(media).(type) {\n\tcase *MessageInputImage:\n\t\tif v.URL != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"url=%s\", *v.URL))\n\t\t}\n\t\tif v.Base64Data != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"base64[%d bytes]\", len(*v.Base64Data)))\n\t\t}\n\t\tif v.MIMEType != \"\" {\n\t\t\tparts = append(parts, fmt.Sprintf(\"mime=%s\", v.MIMEType))\n\t\t}\n\t\tif v.Detail != \"\" {\n\t\t\tparts = append(parts, fmt.Sprintf(\"detail=%s\", v.Detail))\n\t\t}\n\t\tif len(v.Extra) > 0 {\n\t\t\tparts = append(parts, fmt.Sprintf(\"extra=%v\", v.Extra))\n\t\t}\n\tcase *MessageInputAudio:\n\t\tif v.URL != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"url=%s\", *v.URL))\n\t\t}\n\t\tif v.Base64Data != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"base64[%d bytes]\", len(*v.Base64Data)))\n\t\t}\n\t\tif v.MIMEType != \"\" {\n\t\t\tparts = append(parts, fmt.Sprintf(\"mime=%s\", v.MIMEType))\n\t\t}\n\t\tif len(v.Extra) > 0 {\n\t\t\tparts = append(parts, fmt.Sprintf(\"extra=%v\", v.Extra))\n\t\t}\n\tcase *MessageInputVideo:\n\t\tif v.URL != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"url=%s\", *v.URL))\n\t\t}\n\t\tif v.Base64Data != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"base64[%d bytes]\", len(*v.Base64Data)))\n\t\t}\n\t\tif v.MIMEType != \"\" {\n\t\t\tparts = append(parts, fmt.Sprintf(\"mime=%s\", v.MIMEType))\n\t\t}\n\t\tif len(v.Extra) > 0 {\n\t\t\tparts = append(parts, fmt.Sprintf(\"extra=%v\", v.Extra))\n\t\t}\n\t}\n\tif len(parts) == 0 {\n\t\treturn \"<empty>\"\n\t}\n\treturn strings.Join(parts, \", \")\n}\n\nfunc formatMessageInputFile(file *MessageInputFile) string {\n\tif file == nil {\n\t\treturn \"<nil>\"\n\t}\n\tvar parts []string\n\tif file.URL != nil {\n\t\tparts = append(parts, fmt.Sprintf(\"url=%s\", *file.URL))\n\t}\n\tif file.Base64Data != nil {\n\t\tparts = append(parts, fmt.Sprintf(\"base64[%d bytes]\", len(*file.Base64Data)))\n\t}\n\tif file.MIMEType != \"\" {\n\t\tparts = append(parts, fmt.Sprintf(\"mime=%s\", file.MIMEType))\n\t}\n\tif file.Name != \"\" {\n\t\tparts = append(parts, fmt.Sprintf(\"name=%s\", file.Name))\n\t}\n\tif len(file.Extra) > 0 {\n\t\tparts = append(parts, fmt.Sprintf(\"extra=%v\", file.Extra))\n\t}\n\tif len(parts) == 0 {\n\t\treturn \"<empty>\"\n\t}\n\treturn strings.Join(parts, \", \")\n}\n\nfunc formatOutputPart(part MessageOutputPart) string {\n\tswitch part.Type {\n\tcase ChatMessagePartTypeText:\n\t\treturn fmt.Sprintf(\"text: %s\", part.Text)\n\tcase ChatMessagePartTypeImageURL:\n\t\treturn fmt.Sprintf(\"image: %s\", formatMessageOutputMedia(part.Image))\n\tcase ChatMessagePartTypeAudioURL:\n\t\treturn fmt.Sprintf(\"audio: %s\", formatMessageOutputMedia(part.Audio))\n\tcase ChatMessagePartTypeVideoURL:\n\t\treturn fmt.Sprintf(\"video: %s\", formatMessageOutputMedia(part.Video))\n\tdefault:\n\t\treturn fmt.Sprintf(\"unknown type: %s\", part.Type)\n\t}\n}\n\nfunc formatMessageOutputMedia[T MessageOutputImage | MessageOutputAudio | MessageOutputVideo](media *T) string {\n\tif media == nil {\n\t\treturn \"<nil>\"\n\t}\n\tvar parts []string\n\tswitch v := any(media).(type) {\n\tcase *MessageOutputImage:\n\t\tif v.URL != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"url=%s\", *v.URL))\n\t\t}\n\t\tif v.Base64Data != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"base64[%d bytes]\", len(*v.Base64Data)))\n\t\t}\n\t\tif v.MIMEType != \"\" {\n\t\t\tparts = append(parts, fmt.Sprintf(\"mime=%s\", v.MIMEType))\n\t\t}\n\t\tif len(v.Extra) > 0 {\n\t\t\tparts = append(parts, fmt.Sprintf(\"extra=%v\", v.Extra))\n\t\t}\n\tcase *MessageOutputAudio:\n\t\tif v.URL != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"url=%s\", *v.URL))\n\t\t}\n\t\tif v.Base64Data != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"base64[%d bytes]\", len(*v.Base64Data)))\n\t\t}\n\t\tif v.MIMEType != \"\" {\n\t\t\tparts = append(parts, fmt.Sprintf(\"mime=%s\", v.MIMEType))\n\t\t}\n\t\tif len(v.Extra) > 0 {\n\t\t\tparts = append(parts, fmt.Sprintf(\"extra=%v\", v.Extra))\n\t\t}\n\tcase *MessageOutputVideo:\n\t\tif v.URL != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"url=%s\", *v.URL))\n\t\t}\n\t\tif v.Base64Data != nil {\n\t\t\tparts = append(parts, fmt.Sprintf(\"base64[%d bytes]\", len(*v.Base64Data)))\n\t\t}\n\t\tif v.MIMEType != \"\" {\n\t\t\tparts = append(parts, fmt.Sprintf(\"mime=%s\", v.MIMEType))\n\t\t}\n\t\tif len(v.Extra) > 0 {\n\t\t\tparts = append(parts, fmt.Sprintf(\"extra=%v\", v.Extra))\n\t\t}\n\t}\n\tif len(parts) == 0 {\n\t\treturn \"<empty>\"\n\t}\n\treturn strings.Join(parts, \", \")\n}\n\nfunc formatChatMessagePart(part ChatMessagePart) string {\n\tswitch part.Type {\n\tcase ChatMessagePartTypeText:\n\t\treturn fmt.Sprintf(\"text: %s\", part.Text)\n\tcase ChatMessagePartTypeImageURL:\n\t\tif part.ImageURL != nil {\n\t\t\treturn fmt.Sprintf(\"image_url: %s\", part.ImageURL.URL)\n\t\t}\n\t\treturn \"image_url: <nil>\"\n\tcase ChatMessagePartTypeAudioURL:\n\t\tif part.AudioURL != nil {\n\t\t\treturn fmt.Sprintf(\"audio_url: %s\", part.AudioURL.URL)\n\t\t}\n\t\treturn \"audio_url: <nil>\"\n\tcase ChatMessagePartTypeVideoURL:\n\t\tif part.VideoURL != nil {\n\t\t\treturn fmt.Sprintf(\"video_url: %s\", part.VideoURL.URL)\n\t\t}\n\t\treturn \"video_url: <nil>\"\n\tcase ChatMessagePartTypeFileURL:\n\t\tif part.FileURL != nil {\n\t\t\treturn fmt.Sprintf(\"file_url: %s\", part.FileURL.URL)\n\t\t}\n\t\treturn \"file_url: <nil>\"\n\tdefault:\n\t\treturn fmt.Sprintf(\"unknown type: %s\", part.Type)\n\t}\n}\n\n// SystemMessage represents a message with Role \"system\".\nfunc SystemMessage(content string) *Message {\n\treturn &Message{\n\t\tRole:    System,\n\t\tContent: content,\n\t}\n}\n\n// AssistantMessage represents a message with Role \"assistant\".\nfunc AssistantMessage(content string, toolCalls []ToolCall) *Message {\n\treturn &Message{\n\t\tRole:      Assistant,\n\t\tContent:   content,\n\t\tToolCalls: toolCalls,\n\t}\n}\n\n// UserMessage represents a message with Role \"user\".\nfunc UserMessage(content string) *Message {\n\treturn &Message{\n\t\tRole:    User,\n\t\tContent: content,\n\t}\n\n}\n\ntype toolMessageOptions struct {\n\ttoolName string\n}\n\n// ToolMessageOption defines a option for ToolMessage\ntype ToolMessageOption func(*toolMessageOptions)\n\n// WithToolName returns a ToolMessageOption that sets the tool call name.\nfunc WithToolName(name string) ToolMessageOption {\n\treturn func(o *toolMessageOptions) {\n\t\to.toolName = name\n\t}\n}\n\n// ToolMessage represents a message with Role \"tool\".\nfunc ToolMessage(content string, toolCallID string, opts ...ToolMessageOption) *Message {\n\to := &toolMessageOptions{}\n\tfor _, opt := range opts {\n\t\topt(o)\n\t}\n\treturn &Message{\n\t\tRole:       Tool,\n\t\tContent:    content,\n\t\tToolCallID: toolCallID,\n\t\tToolName:   o.toolName,\n\t}\n}\n\n// ConcatToolResults merges multiple ToolResult chunks into a single ToolResult.\n// It collects all ToolOutputParts from the input chunks and merges contiguous text parts within each chunk.\n//\n// Merge rules:\n//   - Text parts: Contiguous text parts within each chunk are concatenated into a single text part.\n//   - Non-text parts (image, audio, video, file): These parts are kept as-is without merging.\n//     Each non-text part type can only appear in one chunk; if the same non-text type appears\n//     in multiple chunks, an error is returned.\n//\n// This function is primarily used in streaming scenarios where tool output is delivered\n// in multiple chunks that need to be merged into a complete result.\n//\n// Parameters:\n//   - chunks: A slice of ToolResult pointers representing sequential chunks from a stream.\n//     Nil chunks and chunks with empty Parts are safely ignored.\n//\n// Returns:\n//   - *ToolResult: The merged ToolResult containing all content from the chunks.\n//     Returns an empty ToolResult if chunks is empty or all chunks are nil/empty.\n//   - error: An error if the same non-text part type appears in multiple chunks.\nfunc ConcatToolResults(chunks []*ToolResult) (*ToolResult, error) {\n\tif len(chunks) == 0 {\n\t\treturn &ToolResult{}, nil\n\t}\n\n\tnonTextPartTypes := make(map[ToolPartType]int)\n\n\tvar allParts []ToolOutputPart\n\tfor chunkIdx, chunk := range chunks {\n\t\tif chunk == nil || len(chunk.Parts) == 0 {\n\t\t\tcontinue\n\t\t}\n\n\t\tfor _, part := range chunk.Parts {\n\t\t\tif part.Type != ToolPartTypeText {\n\t\t\t\tif prevChunkIdx, exists := nonTextPartTypes[part.Type]; exists {\n\t\t\t\t\treturn nil, fmt.Errorf(\"conflicting %s parts found in chunk %d and chunk %d: \"+\n\t\t\t\t\t\t\"non-text modality parts cannot appear in multiple chunks\", part.Type, prevChunkIdx, chunkIdx)\n\t\t\t\t}\n\t\t\t\tnonTextPartTypes[part.Type] = chunkIdx\n\t\t\t}\n\t\t}\n\n\t\tmergedChunkParts, err := concatToolOutputParts(chunk.Parts)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to merge text parts in chunk %d: %w\", chunkIdx, err)\n\t\t}\n\t\tallParts = append(allParts, mergedChunkParts...)\n\t}\n\n\tif len(allParts) == 0 {\n\t\treturn &ToolResult{}, nil\n\t}\n\n\treturn &ToolResult{Parts: allParts}, nil\n}\n\nfunc concatToolOutputParts(parts []ToolOutputPart) ([]ToolOutputPart, error) {\n\tif len(parts) == 0 {\n\t\treturn nil, nil\n\t}\n\n\tgroups := groupToolOutputParts(parts)\n\n\tmerged := make([]ToolOutputPart, 0, len(groups))\n\tfor _, group := range groups {\n\t\tif len(group) == 1 {\n\t\t\tmerged = append(merged, group...)\n\t\t\tcontinue\n\t\t}\n\t\tswitch group[0].Type {\n\t\tcase ToolPartTypeText:\n\t\t\tmergedPart, err := mergeToolTextParts(group)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tmerged = append(merged, mergedPart)\n\t\tdefault:\n\t\t\tmerged = append(merged, group...)\n\t\t}\n\t}\n\n\treturn merged, nil\n}\n\nfunc groupToolOutputParts(parts []ToolOutputPart) [][]ToolOutputPart {\n\tgroups := make([][]ToolOutputPart, 0)\n\ti := 0\n\tfor i < len(parts) {\n\t\tif parts[i].Type == ToolPartTypeText {\n\t\t\tend := i + 1\n\t\t\tfor end < len(parts) && parts[end].Type == ToolPartTypeText {\n\t\t\t\tend++\n\t\t\t}\n\t\t\tgroups = append(groups, parts[i:end])\n\t\t\ti = end\n\t\t} else {\n\t\t\tgroups = append(groups, parts[i:i+1])\n\t\t\ti++\n\t\t}\n\t}\n\treturn groups\n}\n\nfunc mergeToolTextParts(group []ToolOutputPart) (ToolOutputPart, error) {\n\tvar sb strings.Builder\n\textraList := make([]map[string]any, 0, len(group))\n\tfor _, part := range group {\n\t\tsb.WriteString(part.Text)\n\t\tif len(part.Extra) > 0 {\n\t\t\textraList = append(extraList, part.Extra)\n\t\t}\n\t}\n\tvar mergedExtra map[string]any\n\tif len(extraList) > 0 {\n\t\tvar err error\n\t\tmergedExtra, err = concatExtra(extraList)\n\t\tif err != nil {\n\t\t\treturn ToolOutputPart{}, fmt.Errorf(\"failed to concat tool output text part extra: %w\", err)\n\t\t}\n\t}\n\treturn ToolOutputPart{\n\t\tType:  ToolPartTypeText,\n\t\tText:  sb.String(),\n\t\tExtra: mergedExtra,\n\t}, nil\n}\n\nfunc concatToolCalls(chunks []ToolCall) ([]ToolCall, error) {\n\tvar merged []ToolCall\n\tm := make(map[int][]int)\n\tfor i := range chunks {\n\t\tindex := chunks[i].Index\n\t\tif index == nil {\n\t\t\tmerged = append(merged, chunks[i])\n\t\t} else {\n\t\t\tm[*index] = append(m[*index], i)\n\t\t}\n\t}\n\n\tvar args strings.Builder\n\tfor k, v := range m {\n\t\tindex := k\n\t\ttoolCall := ToolCall{Index: &index}\n\t\tif len(v) > 0 {\n\t\t\ttoolCall = chunks[v[0]]\n\t\t}\n\n\t\targs.Reset()\n\t\ttoolID, toolType, toolName := \"\", \"\", \"\" // these field will output atomically in any chunk\n\n\t\tfor _, n := range v {\n\t\t\tchunk := chunks[n]\n\t\t\tif chunk.ID != \"\" {\n\t\t\t\tif toolID == \"\" {\n\t\t\t\t\ttoolID = chunk.ID\n\t\t\t\t} else if toolID != chunk.ID {\n\t\t\t\t\treturn nil, fmt.Errorf(\"cannot concat ToolCalls with different tool id: '%s' '%s'\", toolID, chunk.ID)\n\t\t\t\t}\n\n\t\t\t}\n\n\t\t\tif chunk.Type != \"\" {\n\t\t\t\tif toolType == \"\" {\n\t\t\t\t\ttoolType = chunk.Type\n\t\t\t\t} else if toolType != chunk.Type {\n\t\t\t\t\treturn nil, fmt.Errorf(\"cannot concat ToolCalls with different tool type: '%s' '%s'\", toolType, chunk.Type)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tif chunk.Function.Name != \"\" {\n\t\t\t\tif toolName == \"\" {\n\t\t\t\t\ttoolName = chunk.Function.Name\n\t\t\t\t} else if toolName != chunk.Function.Name {\n\t\t\t\t\treturn nil, fmt.Errorf(\"cannot concat ToolCalls with different tool name: '%s' '%s'\", toolName, chunk.Function.Name)\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tif chunk.Function.Arguments != \"\" {\n\t\t\t\t_, err := args.WriteString(chunk.Function.Arguments)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\ttoolCall.ID = toolID\n\t\ttoolCall.Type = toolType\n\t\ttoolCall.Function.Name = toolName\n\t\ttoolCall.Function.Arguments = args.String()\n\n\t\tmerged = append(merged, toolCall)\n\t}\n\n\tif len(merged) > 1 {\n\t\tsort.SliceStable(merged, func(i, j int) bool {\n\t\t\tiVal, jVal := merged[i].Index, merged[j].Index\n\t\t\tif iVal == nil && jVal == nil {\n\t\t\t\treturn false\n\t\t\t} else if iVal == nil && jVal != nil {\n\t\t\t\treturn true\n\t\t\t} else if iVal != nil && jVal == nil {\n\t\t\t\treturn false\n\t\t\t}\n\n\t\t\treturn *iVal < *jVal\n\t\t})\n\t}\n\n\treturn merged, nil\n}\n\nfunc concatAssistantMultiContent(parts []MessageOutputPart) ([]MessageOutputPart, error) {\n\tif len(parts) == 0 {\n\t\treturn parts, nil\n\t}\n\n\tgroups := groupOutputParts(parts)\n\n\tmerged := make([]MessageOutputPart, 0, len(groups))\n\tfor _, group := range groups {\n\t\tmergedPart, err := mergeOutputPartGroup(group)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tmerged = append(merged, mergedPart)\n\t}\n\n\treturn merged, nil\n}\n\nfunc groupOutputParts(parts []MessageOutputPart) [][]MessageOutputPart {\n\tif len(parts) == 0 {\n\t\treturn nil\n\t}\n\n\tgroups := make([][]MessageOutputPart, 0)\n\tcurrentGroup := []MessageOutputPart{parts[0]}\n\n\tfor i := 1; i < len(parts); i++ {\n\t\tif canMergeOutputParts(currentGroup[0], parts[i]) {\n\t\t\tcurrentGroup = append(currentGroup, parts[i])\n\t\t} else {\n\t\t\tgroups = append(groups, currentGroup)\n\t\t\tcurrentGroup = []MessageOutputPart{parts[i]}\n\t\t}\n\t}\n\tgroups = append(groups, currentGroup)\n\n\treturn groups\n}\n\nfunc canMergeOutputParts(current, next MessageOutputPart) bool {\n\tif current.Type != next.Type {\n\t\treturn false\n\t}\n\n\tif !isMergeableOutputPartType(current) {\n\t\treturn false\n\t}\n\n\tif current.StreamingMeta != nil && next.StreamingMeta != nil {\n\t\treturn current.StreamingMeta.Index == next.StreamingMeta.Index\n\t}\n\n\treturn current.StreamingMeta == nil && next.StreamingMeta == nil\n}\n\nfunc isMergeableOutputPartType(part MessageOutputPart) bool {\n\tswitch part.Type {\n\tcase ChatMessagePartTypeText, ChatMessagePartTypeReasoning:\n\t\treturn true\n\tcase ChatMessagePartTypeAudioURL:\n\t\treturn isBase64MessageOutputAudioPart(part)\n\tdefault:\n\t\treturn false\n\t}\n}\n\nfunc mergeOutputPartGroup(group []MessageOutputPart) (MessageOutputPart, error) {\n\tif len(group) == 0 {\n\t\treturn MessageOutputPart{}, nil\n\t}\n\n\tif len(group) == 1 {\n\t\treturn group[0], nil\n\t}\n\n\tfirst := group[0]\n\tswitch first.Type {\n\tcase ChatMessagePartTypeText:\n\t\treturn mergeTextParts(group)\n\tcase ChatMessagePartTypeReasoning:\n\t\treturn mergeReasoningParts(group)\n\tcase ChatMessagePartTypeAudioURL:\n\t\tif isBase64MessageOutputAudioPart(first) {\n\t\t\treturn mergeAudioParts(group)\n\t\t}\n\t}\n\n\treturn first, nil\n}\n\nfunc mergeTextParts(group []MessageOutputPart) (MessageOutputPart, error) {\n\tvar sb strings.Builder\n\textraList := make([]map[string]any, 0, len(group))\n\tfor _, part := range group {\n\t\tsb.WriteString(part.Text)\n\t\tif len(part.Extra) > 0 {\n\t\t\textraList = append(extraList, part.Extra)\n\t\t}\n\t}\n\tvar mergedExtra map[string]any\n\tif len(extraList) > 0 {\n\t\tvar err error\n\t\tmergedExtra, err = concatExtra(extraList)\n\t\tif err != nil {\n\t\t\treturn MessageOutputPart{}, fmt.Errorf(\"failed to concat text part extra: %w\", err)\n\t\t}\n\t}\n\treturn MessageOutputPart{\n\t\tType:          ChatMessagePartTypeText,\n\t\tText:          sb.String(),\n\t\tExtra:         mergedExtra,\n\t\tStreamingMeta: group[0].StreamingMeta,\n\t}, nil\n}\n\nfunc mergeReasoningParts(group []MessageOutputPart) (MessageOutputPart, error) {\n\tvar textBuilder strings.Builder\n\tvar signature string\n\textraList := make([]map[string]any, 0, len(group))\n\tfor _, part := range group {\n\t\tif part.Reasoning != nil {\n\t\t\ttextBuilder.WriteString(part.Reasoning.Text)\n\t\t\tif part.Reasoning.Signature != \"\" {\n\t\t\t\tsignature = part.Reasoning.Signature\n\t\t\t}\n\t\t}\n\t\tif len(part.Extra) > 0 {\n\t\t\textraList = append(extraList, part.Extra)\n\t\t}\n\t}\n\tvar mergedExtra map[string]any\n\tif len(extraList) > 0 {\n\t\tvar err error\n\t\tmergedExtra, err = concatExtra(extraList)\n\t\tif err != nil {\n\t\t\treturn MessageOutputPart{}, fmt.Errorf(\"failed to concat reasoning part extra: %w\", err)\n\t\t}\n\t}\n\treturn MessageOutputPart{\n\t\tType: ChatMessagePartTypeReasoning,\n\t\tReasoning: &MessageOutputReasoning{\n\t\t\tText:      textBuilder.String(),\n\t\t\tSignature: signature,\n\t\t},\n\t\tExtra:         mergedExtra,\n\t\tStreamingMeta: group[0].StreamingMeta,\n\t}, nil\n}\n\nfunc mergeAudioParts(group []MessageOutputPart) (MessageOutputPart, error) {\n\tvar b64Builder strings.Builder\n\tvar mimeType string\n\taudioExtraList := make([]map[string]any, 0, len(group))\n\tpartExtraList := make([]map[string]any, 0, len(group))\n\n\tfor _, part := range group {\n\t\taudioPart := part.Audio\n\t\tif audioPart.Base64Data != nil {\n\t\t\tb64Builder.WriteString(*audioPart.Base64Data)\n\t\t}\n\t\tif mimeType == \"\" {\n\t\t\tmimeType = audioPart.MIMEType\n\t\t}\n\t\tif len(audioPart.Extra) > 0 {\n\t\t\taudioExtraList = append(audioExtraList, audioPart.Extra)\n\t\t}\n\t\tif len(part.Extra) > 0 {\n\t\t\tpartExtraList = append(partExtraList, part.Extra)\n\t\t}\n\t}\n\n\tvar mergedAudioExtra map[string]any\n\tvar err error\n\tif len(audioExtraList) > 0 {\n\t\tmergedAudioExtra, err = concatExtra(audioExtraList)\n\t\tif err != nil {\n\t\t\treturn MessageOutputPart{}, fmt.Errorf(\"failed to concat audio extra: %w\", err)\n\t\t}\n\t}\n\n\tvar mergedPartExtra map[string]any\n\tif len(partExtraList) > 0 {\n\t\tmergedPartExtra, err = concatExtra(partExtraList)\n\t\tif err != nil {\n\t\t\treturn MessageOutputPart{}, fmt.Errorf(\"failed to concat audio part extra: %w\", err)\n\t\t}\n\t}\n\n\tmergedB64 := b64Builder.String()\n\treturn MessageOutputPart{\n\t\tType: ChatMessagePartTypeAudioURL,\n\t\tAudio: &MessageOutputAudio{\n\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\tBase64Data: &mergedB64,\n\t\t\t\tMIMEType:   mimeType,\n\t\t\t\tExtra:      mergedAudioExtra,\n\t\t\t},\n\t\t},\n\t\tExtra:         mergedPartExtra,\n\t\tStreamingMeta: group[0].StreamingMeta,\n\t}, nil\n}\n\nfunc isBase64MessageOutputAudioPart(part MessageOutputPart) bool {\n\treturn part.Type == ChatMessagePartTypeAudioURL &&\n\t\tpart.Audio != nil &&\n\t\tpart.Audio.Base64Data != nil &&\n\t\tpart.Audio.URL == nil\n}\n\nfunc concatUserMultiContent(parts []MessageInputPart) ([]MessageInputPart, error) {\n\tif len(parts) == 0 {\n\t\treturn parts, nil\n\t}\n\n\tmerged := make([]MessageInputPart, 0, len(parts))\n\ti := 0\n\tfor i < len(parts) {\n\t\tcurrentPart := parts[i]\n\n\t\tif currentPart.Type == ChatMessagePartTypeText {\n\t\t\tend := i + 1\n\t\t\tfor end < len(parts) && parts[end].Type == ChatMessagePartTypeText {\n\t\t\t\tend++\n\t\t\t}\n\n\t\t\tif end == i+1 {\n\t\t\t\tmerged = append(merged, currentPart)\n\t\t\t} else {\n\t\t\t\tvar sb strings.Builder\n\t\t\t\tfor k := i; k < end; k++ {\n\t\t\t\t\tsb.WriteString(parts[k].Text)\n\t\t\t\t}\n\t\t\t\tmergedPart := MessageInputPart{\n\t\t\t\t\tType: ChatMessagePartTypeText,\n\t\t\t\t\tText: sb.String(),\n\t\t\t\t}\n\t\t\t\tmerged = append(merged, mergedPart)\n\t\t\t}\n\t\t\ti = end\n\t\t} else {\n\n\t\t\tmerged = append(merged, currentPart)\n\t\t\ti++\n\t\t}\n\t}\n\n\treturn merged, nil\n}\n\nfunc concatExtra(extraList []map[string]any) (map[string]any, error) {\n\tif len(extraList) == 1 {\n\t\treturn generic.CopyMap(extraList[0]), nil\n\t}\n\n\treturn internal.ConcatItems(extraList)\n}\n\n// ConcatMessages concat messages with the same role and name.\n// It will concat tool calls with the same index.\n// It will return an error if the messages have different roles or names.\n// It's useful for concatenating messages from a stream.\n// e.g.\n//\n//\tmsgs := []*Message{}\n//\tfor {\n//\t\tmsg, err := stream.Recv()\n//\t\tif errors.Is(err, io.EOF) {\n//\t\t\tbreak\n//\t\t}\n//\t\tif err != nil {...}\n//\t\tmsgs = append(msgs, msg)\n//\t}\n//\n// concatedMsg, err := ConcatMessages(msgs) // concatedMsg.Content will be full content of all messages\nfunc ConcatMessages(msgs []*Message) (*Message, error) {\n\tvar (\n\t\tcontents                      []string\n\t\tcontentLen                    int\n\t\treasoningContents             []string\n\t\treasoningContentLen           int\n\t\ttoolCalls                     []ToolCall\n\t\tmultiContentParts             []ChatMessagePart\n\t\tassistantGenMultiContentParts []MessageOutputPart\n\t\tuserInputMultiContentParts    []MessageInputPart\n\t\tret                           = Message{}\n\t\textraList                     = make([]map[string]any, 0, len(msgs))\n\t)\n\n\tfor idx, msg := range msgs {\n\t\tif msg == nil {\n\t\t\treturn nil, fmt.Errorf(\"unexpected nil chunk in message stream, index: %d\", idx)\n\t\t}\n\n\t\tif msg.Role != \"\" {\n\t\t\tif ret.Role == \"\" {\n\t\t\t\tret.Role = msg.Role\n\t\t\t} else if ret.Role != msg.Role {\n\t\t\t\treturn nil, fmt.Errorf(\"cannot concat messages with \"+\n\t\t\t\t\t\"different roles: '%s' '%s'\", ret.Role, msg.Role)\n\t\t\t}\n\t\t}\n\n\t\tif msg.Name != \"\" {\n\t\t\tif ret.Name == \"\" {\n\t\t\t\tret.Name = msg.Name\n\t\t\t} else if ret.Name != msg.Name {\n\t\t\t\treturn nil, fmt.Errorf(\"cannot concat messages with\"+\n\t\t\t\t\t\" different names: '%s' '%s'\", ret.Name, msg.Name)\n\t\t\t}\n\t\t}\n\n\t\tif msg.ToolCallID != \"\" {\n\t\t\tif ret.ToolCallID == \"\" {\n\t\t\t\tret.ToolCallID = msg.ToolCallID\n\t\t\t} else if ret.ToolCallID != msg.ToolCallID {\n\t\t\t\treturn nil, fmt.Errorf(\"cannot concat messages with\"+\n\t\t\t\t\t\" different toolCallIDs: '%s' '%s'\", ret.ToolCallID, msg.ToolCallID)\n\t\t\t}\n\t\t}\n\t\tif msg.ToolName != \"\" {\n\t\t\tif ret.ToolName == \"\" {\n\t\t\t\tret.ToolName = msg.ToolName\n\t\t\t} else if ret.ToolName != msg.ToolName {\n\t\t\t\treturn nil, fmt.Errorf(\"cannot concat messages with\"+\n\t\t\t\t\t\" different toolNames: '%s' '%s'\", ret.ToolCallID, msg.ToolCallID)\n\t\t\t}\n\t\t}\n\n\t\tif msg.Content != \"\" {\n\t\t\tcontents = append(contents, msg.Content)\n\t\t\tcontentLen += len(msg.Content)\n\t\t}\n\t\tif msg.ReasoningContent != \"\" {\n\t\t\treasoningContents = append(reasoningContents, msg.ReasoningContent)\n\t\t\treasoningContentLen += len(msg.ReasoningContent)\n\t\t}\n\n\t\tif len(msg.ToolCalls) > 0 {\n\t\t\ttoolCalls = append(toolCalls, msg.ToolCalls...)\n\t\t}\n\n\t\tif len(msg.Extra) > 0 {\n\t\t\textraList = append(extraList, msg.Extra)\n\t\t}\n\n\t\t// The 'MultiContent' field is deprecated but is kept for backward compatibility.\n\t\tif len(msg.MultiContent) > 0 {\n\t\t\tmultiContentParts = append(multiContentParts, msg.MultiContent...)\n\t\t}\n\n\t\tif len(msg.AssistantGenMultiContent) > 0 {\n\t\t\tassistantGenMultiContentParts = append(assistantGenMultiContentParts, msg.AssistantGenMultiContent...)\n\t\t}\n\t\tif len(msg.UserInputMultiContent) > 0 {\n\t\t\tuserInputMultiContentParts = append(userInputMultiContentParts, msg.UserInputMultiContent...)\n\t\t}\n\t\tif msg.ResponseMeta != nil && ret.ResponseMeta == nil {\n\t\t\tret.ResponseMeta = &ResponseMeta{}\n\t\t}\n\n\t\tif msg.ResponseMeta != nil && ret.ResponseMeta != nil {\n\t\t\t// keep the last FinishReason with a valid value.\n\t\t\tif msg.ResponseMeta.FinishReason != \"\" {\n\t\t\t\tret.ResponseMeta.FinishReason = msg.ResponseMeta.FinishReason\n\t\t\t}\n\n\t\t\tif msg.ResponseMeta.Usage != nil {\n\t\t\t\tif ret.ResponseMeta.Usage == nil {\n\t\t\t\t\tret.ResponseMeta.Usage = &TokenUsage{}\n\t\t\t\t}\n\n\t\t\t\tif msg.ResponseMeta.Usage.PromptTokens > ret.ResponseMeta.Usage.PromptTokens {\n\t\t\t\t\tret.ResponseMeta.Usage.PromptTokens = msg.ResponseMeta.Usage.PromptTokens\n\t\t\t\t}\n\t\t\t\tif msg.ResponseMeta.Usage.CompletionTokens > ret.ResponseMeta.Usage.CompletionTokens {\n\t\t\t\t\tret.ResponseMeta.Usage.CompletionTokens = msg.ResponseMeta.Usage.CompletionTokens\n\t\t\t\t}\n\n\t\t\t\tif msg.ResponseMeta.Usage.TotalTokens > ret.ResponseMeta.Usage.TotalTokens {\n\t\t\t\t\tret.ResponseMeta.Usage.TotalTokens = msg.ResponseMeta.Usage.TotalTokens\n\t\t\t\t}\n\n\t\t\t\tif msg.ResponseMeta.Usage.PromptTokenDetails.CachedTokens > ret.ResponseMeta.Usage.PromptTokenDetails.CachedTokens {\n\t\t\t\t\tret.ResponseMeta.Usage.PromptTokenDetails.CachedTokens = msg.ResponseMeta.Usage.PromptTokenDetails.CachedTokens\n\t\t\t\t}\n\n\t\t\t\tif msg.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens > ret.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens {\n\t\t\t\t\tret.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens = msg.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tif msg.ResponseMeta.LogProbs != nil {\n\t\t\t\tif ret.ResponseMeta.LogProbs == nil {\n\t\t\t\t\tret.ResponseMeta.LogProbs = &LogProbs{}\n\t\t\t\t}\n\n\t\t\t\tret.ResponseMeta.LogProbs.Content = append(ret.ResponseMeta.LogProbs.Content, msg.ResponseMeta.LogProbs.Content...)\n\t\t\t}\n\n\t\t}\n\t}\n\n\tif len(contents) > 0 {\n\t\tvar sb strings.Builder\n\t\tsb.Grow(contentLen)\n\t\tfor _, content := range contents {\n\t\t\t_, err := sb.WriteString(content)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\n\t\tret.Content = sb.String()\n\t}\n\tif len(reasoningContents) > 0 {\n\t\tvar sb strings.Builder\n\t\tsb.Grow(reasoningContentLen)\n\t\tfor _, rc := range reasoningContents {\n\t\t\t_, err := sb.WriteString(rc)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\n\t\tret.ReasoningContent = sb.String()\n\t}\n\n\tif len(toolCalls) > 0 {\n\t\tmerged, err := concatToolCalls(toolCalls)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tret.ToolCalls = merged\n\t}\n\n\tif len(extraList) > 0 {\n\t\textra, err := concatExtra(extraList)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to concat message's extra: %w\", err)\n\t\t}\n\n\t\tif len(extra) > 0 {\n\t\t\tret.Extra = extra\n\t\t}\n\t}\n\n\tif len(multiContentParts) > 0 {\n\t\tret.MultiContent = multiContentParts\n\t}\n\n\tif len(assistantGenMultiContentParts) > 0 {\n\t\tmerged, err := concatAssistantMultiContent(assistantGenMultiContentParts)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to concat message's assistant multicontent: %w\", err)\n\t\t}\n\t\tret.AssistantGenMultiContent = merged\n\t}\n\n\tif len(userInputMultiContentParts) > 0 {\n\t\tmerged, err := concatUserMultiContent(userInputMultiContentParts)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to concat message's user multicontent: %w\", err)\n\t\t}\n\t\tret.UserInputMultiContent = merged\n\t}\n\n\treturn &ret, nil\n}\n\n// ConcatMessageStream drains a stream of messages and returns a single\n// concatenated message representing the merged content.\nfunc ConcatMessageStream(s *StreamReader[*Message]) (*Message, error) {\n\tdefer s.Close()\n\n\tvar msgs []*Message\n\tfor {\n\t\tmsg, err := s.Recv()\n\t\tif err != nil {\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\treturn nil, err\n\t\t}\n\n\t\tmsgs = append(msgs, msg)\n\t}\n\n\treturn ConcatMessages(msgs)\n}\n\n// custom jinja env\nvar jinjaEnvOnce sync.Once\nvar jinjaEnv *gonja.Environment\nvar envInitErr error\n\nconst (\n\tjinjaInclude = \"include\"\n\tjinjaExtends = \"extends\"\n\tjinjaImport  = \"import\"\n\tjinjaFrom    = \"from\"\n)\n\nfunc getJinjaEnv() (*gonja.Environment, error) {\n\tjinjaEnvOnce.Do(func() {\n\t\tjinjaEnv = gonja.NewEnvironment(config.DefaultConfig, gonja.DefaultLoader)\n\t\tformatInitError := \"init jinja env fail: %w\"\n\t\tvar err error\n\t\tif jinjaEnv.Statements.Exists(jinjaInclude) {\n\t\t\terr = jinjaEnv.Statements.Replace(jinjaInclude, func(parser *parser.Parser, args *parser.Parser) (nodes.Statement, error) {\n\t\t\t\treturn nil, fmt.Errorf(\"keyword[include] has been disabled\")\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tenvInitErr = fmt.Errorf(formatInitError, err)\n\t\t\t\treturn\n\t\t\t}\n\t\t}\n\t\tif jinjaEnv.Statements.Exists(jinjaExtends) {\n\t\t\terr = jinjaEnv.Statements.Replace(jinjaExtends, func(parser *parser.Parser, args *parser.Parser) (nodes.Statement, error) {\n\t\t\t\treturn nil, fmt.Errorf(\"keyword[extends] has been disabled\")\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tenvInitErr = fmt.Errorf(formatInitError, err)\n\t\t\t\treturn\n\t\t\t}\n\t\t}\n\t\tif jinjaEnv.Statements.Exists(jinjaFrom) {\n\t\t\terr = jinjaEnv.Statements.Replace(jinjaFrom, func(parser *parser.Parser, args *parser.Parser) (nodes.Statement, error) {\n\t\t\t\treturn nil, fmt.Errorf(\"keyword[from] has been disabled\")\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tenvInitErr = fmt.Errorf(formatInitError, err)\n\t\t\t\treturn\n\t\t\t}\n\t\t}\n\t\tif jinjaEnv.Statements.Exists(jinjaImport) {\n\t\t\terr = jinjaEnv.Statements.Replace(jinjaImport, func(parser *parser.Parser, args *parser.Parser) (nodes.Statement, error) {\n\t\t\t\treturn nil, fmt.Errorf(\"keyword[import] has been disabled\")\n\t\t\t})\n\t\t\tif err != nil {\n\t\t\t\tenvInitErr = fmt.Errorf(formatInitError, err)\n\t\t\t\treturn\n\t\t\t}\n\t\t}\n\t})\n\treturn jinjaEnv, envInitErr\n}\n"
  },
  {
    "path": "schema/message_parser.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"strings\"\n\n\t\"github.com/bytedance/sonic\"\n)\n\n// MessageParser parses a Message into a strongly typed value.\ntype MessageParser[T any] interface {\n\tParse(ctx context.Context, m *Message) (T, error)\n}\n\n// MessageParseFrom determines the source of the data to be parsed. default is content (Message.Content).\ntype MessageParseFrom string\n\n// MessageParseFrom indicates the source data used by the parser.\nconst (\n\tMessageParseFromContent  MessageParseFrom = \"content\"\n\tMessageParseFromToolCall MessageParseFrom = \"tool_call\"\n)\n\n// MessageJSONParseConfig configures JSON parsing behavior for Message.\ntype MessageJSONParseConfig struct {\n\t// parse from content or tool call, default is content.\n\tParseFrom MessageParseFrom `json:\"parse_from,omitempty\"`\n\n\t// parse key path, default is empty.\n\t// must be a valid json path expression, eg: field.sub_field\n\tParseKeyPath string `json:\"parse_key_path,omitempty\"`\n}\n\n// NewMessageJSONParser creates a new MessageJSONParser.\nfunc NewMessageJSONParser[T any](config *MessageJSONParseConfig) MessageParser[T] {\n\tif config == nil {\n\t\tconfig = &MessageJSONParseConfig{}\n\t}\n\n\tif config.ParseFrom == \"\" {\n\t\tconfig.ParseFrom = MessageParseFromContent\n\t}\n\n\treturn &MessageJSONParser[T]{\n\t\tParseFrom:    config.ParseFrom,\n\t\tParseKeyPath: config.ParseKeyPath,\n\t}\n}\n\n// MessageJSONParser is a parser that parses a message into an object T, using json unmarshal.\n// eg of parse to single struct:\n//\n//\tconfig := &MessageJSONParseConfig{\n//\t\tParseFrom: MessageParseFromToolCall,\n//\t}\n//\tparser := NewMessageJSONParser[GetUserParam](config)\n//\tparam, err := parser.Parse(ctx, message)\n//\n//\teg of parse to slice of struct:\n//\n//\tconfig := &MessageJSONParseConfig{\n//\t\tParseFrom: MessageParseFromToolCall,\n//\t}\n//\n// parser := NewMessageJSONParser[GetUserParam](config)\n// param, err := parser.Parse(ctx, message)\ntype MessageJSONParser[T any] struct {\n\tParseFrom    MessageParseFrom\n\tParseKeyPath string\n}\n\n// Parse parses a message into an object T.\nfunc (p *MessageJSONParser[T]) Parse(ctx context.Context, m *Message) (parsed T, err error) {\n\tif p.ParseFrom == MessageParseFromContent {\n\t\treturn p.parse(m.Content)\n\t} else if p.ParseFrom == MessageParseFromToolCall {\n\t\tif len(m.ToolCalls) == 0 {\n\t\t\treturn parsed, fmt.Errorf(\"no tool call found\")\n\t\t}\n\n\t\treturn p.parse(m.ToolCalls[0].Function.Arguments)\n\t}\n\n\treturn parsed, fmt.Errorf(\"invalid parse from type: %s\", p.ParseFrom)\n}\n\n// extractData extracts data from a string using the parse key path.\nfunc (p *MessageJSONParser[T]) extractData(data string) (string, error) {\n\tif p.ParseKeyPath == \"\" {\n\t\treturn data, nil\n\t}\n\n\tkeys := strings.Split(p.ParseKeyPath, \".\")\n\tinterfaceKeys := make([]any, len(keys))\n\tfor i, key := range keys {\n\t\tinterfaceKeys[i] = key\n\t}\n\n\tnode, err := sonic.GetFromString(data, interfaceKeys...)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to get parse key path: %w\", err)\n\t}\n\n\tbytes, err := node.MarshalJSON()\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to marshal node: %w\", err)\n\t}\n\n\treturn string(bytes), nil\n}\n\n// parse parses a string into an object T.\nfunc (p *MessageJSONParser[T]) parse(data string) (parsed T, err error) {\n\tparsedData, err := p.extractData(data)\n\tif err != nil {\n\t\treturn parsed, err\n\t}\n\n\tif err := sonic.UnmarshalString(parsedData, &parsed); err != nil {\n\t\treturn parsed, fmt.Errorf(\"failed to unmarshal content: %w\", err)\n\t}\n\n\treturn parsed, nil\n}\n"
  },
  {
    "path": "schema/message_parser_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\ntype TestStructForParse struct {\n\tID   int    `json:\"id\"`\n\tName string `json:\"name\"`\n\tXX   struct {\n\t\tYY int `json:\"yy\"`\n\t} `json:\"xx\"`\n}\n\nfunc TestMessageJSONParser(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"parse from content\", func(t *testing.T) {\n\t\tparser := NewMessageJSONParser[TestStructForParse](&MessageJSONParseConfig{\n\t\t\tParseFrom: MessageParseFromContent,\n\t\t})\n\n\t\tparsed, err := parser.Parse(ctx, &Message{\n\t\t\tContent: `{\"id\": 1, \"name\": \"test\", \"xx\": {\"yy\": 2}}`,\n\t\t})\n\t\tassert.Nil(t, err)\n\t\tassert.Equal(t, 1, parsed.ID)\n\t})\n\n\tt.Run(\"parse from tool call\", func(t *testing.T) {\n\t\tt.Run(\"only one tool call, default use first tool call\", func(t *testing.T) {\n\t\t\tparser := NewMessageJSONParser[TestStructForParse](&MessageJSONParseConfig{\n\t\t\t\tParseFrom: MessageParseFromToolCall,\n\t\t\t})\n\n\t\t\tparsed, err := parser.Parse(ctx, &Message{\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{Function: FunctionCall{Arguments: `{\"id\": 1, \"name\": \"test\", \"xx\": {\"yy\": 2}}`}},\n\t\t\t\t},\n\t\t\t})\n\t\t\tassert.Nil(t, err)\n\t\t\tassert.Equal(t, 1, parsed.ID)\n\t\t})\n\n\t\tt.Run(\"parse key path\", func(t *testing.T) {\n\t\t\ttype TestStructForParse2 struct {\n\t\t\t\tYY int `json:\"yy\"`\n\t\t\t}\n\n\t\t\tparser := NewMessageJSONParser[TestStructForParse2](&MessageJSONParseConfig{\n\t\t\t\tParseFrom:    MessageParseFromToolCall,\n\t\t\t\tParseKeyPath: \"xx\",\n\t\t\t})\n\n\t\t\tparsed, err := parser.Parse(ctx, &Message{\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{Function: FunctionCall{Arguments: `{\"id\": 1, \"name\": \"test\", \"xx\": {\"yy\": 2}}`}},\n\t\t\t\t},\n\t\t\t})\n\t\t\tassert.Nil(t, err)\n\t\t\tassert.Equal(t, 2, parsed.YY)\n\t\t})\n\n\t\tt.Run(\"parse key path, deep level\", func(t *testing.T) {\n\t\t\ttype TestStructForParse3 struct {\n\t\t\t\tZZ int `json:\"zz\"`\n\t\t\t}\n\n\t\t\tparser := NewMessageJSONParser[TestStructForParse3](&MessageJSONParseConfig{\n\t\t\t\tParseFrom:    MessageParseFromToolCall,\n\t\t\t\tParseKeyPath: \"xx.yy\",\n\t\t\t})\n\n\t\t\tparsed, err := parser.Parse(ctx, &Message{\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{Function: FunctionCall{Arguments: `{\"id\": 1, \"name\": \"test\", \"xx\": {\"yy\": {\"zz\": 3}}}`}},\n\t\t\t\t},\n\t\t\t})\n\t\t\tassert.Nil(t, err)\n\t\t\tassert.Equal(t, 3, parsed.ZZ)\n\t\t})\n\n\t\tt.Run(\"parse key with pointer\", func(t *testing.T) {\n\t\t\ttype TestStructForParse4 struct {\n\t\t\t\tZZ *int `json:\"zz\"`\n\t\t\t}\n\n\t\t\tparser := NewMessageJSONParser[**TestStructForParse4](&MessageJSONParseConfig{\n\t\t\t\tParseFrom: MessageParseFromToolCall,\n\t\t\t})\n\n\t\t\tparsed, err := parser.Parse(ctx, &Message{\n\t\t\t\tToolCalls: []ToolCall{{Function: FunctionCall{Arguments: `{\"zz\": 3}`}}},\n\t\t\t})\n\t\t\tassert.Nil(t, err)\n\t\t\tassert.Equal(t, 3, *((**parsed).ZZ))\n\t\t})\n\t})\n\n\tt.Run(\"parse of slice\", func(t *testing.T) {\n\t\tt.Run(\"valid slice string, not multiple tool calls\", func(t *testing.T) {\n\t\t\tparser := NewMessageJSONParser[[]map[string]any](&MessageJSONParseConfig{\n\t\t\t\tParseFrom: MessageParseFromToolCall,\n\t\t\t})\n\n\t\t\tparsed, err := parser.Parse(ctx, &Message{\n\t\t\t\tToolCalls: []ToolCall{{Function: FunctionCall{Arguments: `[{\"id\": 1}, {\"id\": 2}]`}}},\n\t\t\t})\n\t\t\tassert.Nil(t, err)\n\t\t\tassert.Equal(t, 2, len(parsed))\n\t\t})\n\n\t\tt.Run(\"invalid slice string, not multiple tool calls\", func(t *testing.T) {\n\t\t\tparser := NewMessageJSONParser[[]map[string]any](&MessageJSONParseConfig{\n\t\t\t\tParseFrom: MessageParseFromToolCall,\n\t\t\t})\n\n\t\t\t_, err := parser.Parse(ctx, &Message{\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{Function: FunctionCall{Arguments: `{\"id\": 1}`}},\n\t\t\t\t\t{Function: FunctionCall{Arguments: `{\"id\": 2}`}},\n\t\t\t\t},\n\t\t\t})\n\t\t\tassert.NotNil(t, err)\n\t\t})\n\t})\n\n\tt.Run(\"invalid configs\", func(t *testing.T) {\n\t\tparser := NewMessageJSONParser[TestStructForParse](nil)\n\t\t_, err := parser.Parse(ctx, &Message{\n\t\t\tContent: \"\",\n\t\t})\n\t\tassert.NotNil(t, err)\n\t})\n\n\tt.Run(\"invalid parse key path\", func(t *testing.T) {\n\t\tparser := NewMessageJSONParser[TestStructForParse](&MessageJSONParseConfig{\n\t\t\tParseKeyPath: \"...invalid\",\n\t\t})\n\t\t_, err := parser.Parse(ctx, &Message{})\n\t\tassert.NotNil(t, err)\n\t})\n\n\tt.Run(\"invalid parse from\", func(t *testing.T) {\n\t\tparser := NewMessageJSONParser[TestStructForParse](&MessageJSONParseConfig{\n\t\t\tParseFrom: \"invalid\",\n\t\t})\n\t\t_, err := parser.Parse(ctx, &Message{})\n\t\tassert.NotNil(t, err)\n\t})\n\n\tt.Run(\"invalid parse from type\", func(t *testing.T) {\n\t\tparser := NewMessageJSONParser[int](&MessageJSONParseConfig{\n\t\t\tParseFrom: MessageParseFrom(\"invalid\"),\n\t\t})\n\t\t_, err := parser.Parse(ctx, &Message{})\n\t\tassert.NotNil(t, err)\n\t})\n\n}\n"
  },
  {
    "path": "schema/message_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"context\"\n\t\"reflect\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n)\n\nfunc TestMessageTemplate(t *testing.T) {\n\tpyFmtMessage := UserMessage(\"input: {question}\")\n\tjinja2Message := UserMessage(\"input: {{question}}\")\n\tgoTemplateMessage := UserMessage(\"input: {{.question}}\")\n\tctx := context.Background()\n\tquestion := \"what's the weather today\"\n\texpected := []*Message{UserMessage(\"input: \" + question)}\n\n\tms, err := pyFmtMessage.Format(ctx, map[string]any{\"question\": question}, FString)\n\tassert.Nil(t, err)\n\tassert.True(t, reflect.DeepEqual(expected, ms))\n\tms, err = jinja2Message.Format(ctx, map[string]any{\"question\": question}, Jinja2)\n\tassert.Nil(t, err)\n\tassert.True(t, reflect.DeepEqual(expected, ms))\n\tms, err = goTemplateMessage.Format(ctx, map[string]any{\"question\": question}, GoTemplate)\n\tassert.Nil(t, err)\n\tassert.True(t, reflect.DeepEqual(expected, ms))\n\n\tmp := MessagesPlaceholder(\"chat_history\", false)\n\tm1 := UserMessage(\"how are you?\")\n\tm2 := AssistantMessage(\"I'm good. how about you?\", nil)\n\tms, err = mp.Format(ctx, map[string]any{\"chat_history\": []*Message{m1, m2}}, FString)\n\tassert.Nil(t, err)\n\n\t// len(ms) == 2\n\tassert.Equal(t, 2, len(ms))\n\tassert.Equal(t, ms[0], m1)\n\tassert.Equal(t, ms[1], m2)\n}\n\nfunc TestConcatMessage(t *testing.T) {\n\tt.Run(\"tool_call_normal_append\", func(t *testing.T) {\n\t\texpectMsg := &Message{\n\t\t\tRole:    \"assistant\",\n\t\t\tContent: \"\",\n\t\t\tToolCalls: []ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\tID:    \"i_am_a_too_call_id\",\n\t\t\t\t\tType:  \"function\",\n\t\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\t\tName:      \"i_am_a_tool_name\",\n\t\t\t\t\t\tArguments: \"{}\",\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tgivenMsgList := []*Message{\n\t\t\t{\n\t\t\t\tRole:    \"\",\n\t\t\t\tContent: \"\",\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\t\tID:    \"\",\n\t\t\t\t\t\tType:  \"\",\n\t\t\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\t\t\tName: \"\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole:    \"assistant\",\n\t\t\t\tContent: \"\",\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\t\tID:    \"i_am_a_too_call_id\",\n\t\t\t\t\t\tType:  \"function\",\n\t\t\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\t\t\tName: \"i_am_a_tool_name\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole:    \"\",\n\t\t\t\tContent: \"\",\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\t\tID:    \"\",\n\t\t\t\t\t\tType:  \"\",\n\t\t\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\t\t\tName:      \"\",\n\t\t\t\t\t\t\tArguments: \"{}\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmsg, err := ConcatMessages(givenMsgList)\n\t\tassert.NoError(t, err)\n\t\tassert.EqualValues(t, expectMsg, msg)\n\t})\n\n\tt.Run(\"exist_nil_message\", func(t *testing.T) {\n\t\tgivenMsgList := []*Message{\n\t\t\tnil,\n\t\t\t{\n\t\t\t\tRole:    \"assistant\",\n\t\t\t\tContent: \"\",\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\t\tID:    \"i_am_a_too_call_id\",\n\t\t\t\t\t\tType:  \"function\",\n\t\t\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\t\t\tName: \"i_am_a_tool_name\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := ConcatMessages(givenMsgList)\n\t\tassert.ErrorContains(t, err, \"unexpected nil chunk in message stream\")\n\t})\n\n\tt.Run(\"response_meta\", func(t *testing.T) {\n\t\texpectedMsg := &Message{\n\t\t\tRole: \"assistant\",\n\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\tFinishReason: \"stop\",\n\t\t\t\tUsage: &TokenUsage{\n\t\t\t\t\tCompletionTokens: 15,\n\t\t\t\t\tPromptTokens:     30,\n\t\t\t\t\tPromptTokenDetails: PromptTokenDetails{\n\t\t\t\t\t\tCachedTokens: 15,\n\t\t\t\t\t},\n\t\t\t\t\tCompletionTokensDetails: CompletionTokensDetails{\n\t\t\t\t\t\tReasoningTokens: 8,\n\t\t\t\t\t},\n\t\t\t\t\tTotalTokens: 45,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tgivenMsgList := []*Message{\n\t\t\t{\n\t\t\t\tRole: \"assistant\",\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: \"assistant\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tFinishReason: \"\",\n\t\t\t\t\tUsage: &TokenUsage{\n\t\t\t\t\t\tCompletionTokens: 10,\n\t\t\t\t\t\tPromptTokens:     20,\n\t\t\t\t\t\tPromptTokenDetails: PromptTokenDetails{\n\t\t\t\t\t\t\tCachedTokens: 10,\n\t\t\t\t\t\t},\n\t\t\t\t\t\tCompletionTokensDetails: CompletionTokensDetails{\n\t\t\t\t\t\t\tReasoningTokens: 5,\n\t\t\t\t\t\t},\n\t\t\t\t\t\tTotalTokens: 30,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: \"assistant\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tFinishReason: \"stop\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: \"assistant\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tUsage: &TokenUsage{\n\t\t\t\t\t\tCompletionTokens: 15,\n\t\t\t\t\t\tPromptTokens:     30,\n\t\t\t\t\t\tPromptTokenDetails: PromptTokenDetails{\n\t\t\t\t\t\t\tCachedTokens: 15,\n\t\t\t\t\t\t},\n\t\t\t\t\t\tCompletionTokensDetails: CompletionTokensDetails{\n\t\t\t\t\t\t\tReasoningTokens: 8,\n\t\t\t\t\t\t},\n\t\t\t\t\t\tTotalTokens: 45,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmsg, err := ConcatMessages(givenMsgList)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, expectedMsg, msg)\n\n\t\tgivenMsgList = append(givenMsgList, &Message{\n\t\t\tRole: \"assistant\",\n\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\tFinishReason: \"tool_calls\",\n\t\t\t},\n\t\t})\n\t\tmsg, err = ConcatMessages(givenMsgList)\n\t\tassert.NoError(t, err)\n\t\texpectedMsg.ResponseMeta.FinishReason = \"tool_calls\"\n\t\tassert.Equal(t, expectedMsg, msg)\n\n\t})\n\n\tt.Run(\"err: different roles\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{Role: User},\n\t\t\t{Role: Assistant},\n\t\t}\n\n\t\tmsg, err := ConcatMessages(msgs)\n\t\tif assert.Error(t, err) {\n\t\t\tassert.ErrorContains(t, err, \"cannot concat messages with different roles\")\n\t\t\tassert.Nil(t, msg)\n\t\t}\n\t})\n\n\tt.Run(\"err: different name\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{Role: Assistant, Name: \"n\", Content: \"1\"},\n\t\t\t{Role: Assistant, Name: \"a\", Content: \"2\"},\n\t\t}\n\n\t\tmsg, err := ConcatMessages(msgs)\n\t\tif assert.Error(t, err) {\n\t\t\tassert.ErrorContains(t, err, \"cannot concat messages with different names\")\n\t\t\tassert.Nil(t, msg)\n\t\t}\n\t})\n\n\tt.Run(\"err: different tool name\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole:       \"\",\n\t\t\t\tContent:    \"\",\n\t\t\t\tToolCallID: \"123\",\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\t\tID:    \"abc\",\n\t\t\t\t\t\tType:  \"\",\n\t\t\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\t\t\tName: \"\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole:       \"assistant\",\n\t\t\t\tContent:    \"\",\n\t\t\t\tToolCallID: \"321\",\n\t\t\t\tToolCalls: []ToolCall{\n\t\t\t\t\t{\n\t\t\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\t\t\tID:    \"abc\",\n\t\t\t\t\t\tType:  \"function\",\n\t\t\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\t\t\tName: \"i_am_a_tool_name\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmsg, err := ConcatMessages(msgs)\n\t\tif assert.Error(t, err) {\n\t\t\tassert.ErrorContains(t, err, \"cannot concat messages with different toolCallIDs\")\n\t\t\tassert.Nil(t, msg)\n\t\t}\n\t})\n\n\tt.Run(\"first response meta usage is nil\", func(t *testing.T) {\n\t\texp := &Message{\n\t\t\tRole: \"assistant\",\n\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\tFinishReason: \"stop\",\n\t\t\t\tUsage: &TokenUsage{\n\t\t\t\t\tCompletionTokens: 15,\n\t\t\t\t\tPromptTokens:     30,\n\t\t\t\t\tTotalTokens:      45,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: \"assistant\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tFinishReason: \"\",\n\t\t\t\t\tUsage:        nil,\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: \"assistant\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tFinishReason: \"stop\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: \"assistant\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tUsage: &TokenUsage{\n\t\t\t\t\t\tCompletionTokens: 15,\n\t\t\t\t\t\tPromptTokens:     30,\n\t\t\t\t\t\tTotalTokens:      45,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, exp, msg)\n\t})\n\n\tt.Run(\"concurrent concat\", func(t *testing.T) {\n\t\tcontent := \"i_am_a_good_concat_message\"\n\t\texp := &Message{Role: Assistant, Content: content}\n\t\tvar msgs []*Message\n\t\tfor i := 0; i < len(content); i++ {\n\t\t\tmsgs = append(msgs, &Message{Role: Assistant, Content: content[i : i+1]})\n\t\t}\n\n\t\twg := sync.WaitGroup{}\n\t\tsize := 100\n\t\twg.Add(size)\n\t\tfor i := 0; i < size; i++ {\n\t\t\tgo func() {\n\t\t\t\tdefer wg.Done()\n\t\t\t\tmsg, err := ConcatMessages(msgs)\n\t\t\t\tassert.NoError(t, err)\n\t\t\t\tassert.Equal(t, exp, msg)\n\t\t\t}()\n\t\t}\n\n\t\twg.Wait()\n\t})\n\n\tt.Run(\"concat logprobs\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole:    Assistant,\n\t\t\t\tContent: \"🚀\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tLogProbs: &LogProbs{\n\t\t\t\t\t\tContent: []LogProb{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tToken:   \"\\\\xf0\\\\x9f\\\\x9a\",\n\t\t\t\t\t\t\t\tLogProb: -0.0000073458323,\n\t\t\t\t\t\t\t\tBytes:   []int64{240, 159, 154},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tToken:   \"\\\\x80\",\n\t\t\t\t\t\t\t\tLogProb: 0,\n\t\t\t\t\t\t\t\tBytes:   []int64{128},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole:    \"\",\n\t\t\t\tContent: \"❤️\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tLogProbs: &LogProbs{\n\t\t\t\t\t\tContent: []LogProb{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tToken:   \"❤️\",\n\t\t\t\t\t\t\t\tLogProb: -0.0011431955,\n\t\t\t\t\t\t\t\tBytes:   []int64{226, 157, 164, 239, 184, 143},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: \"\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tFinishReason: \"stop\",\n\t\t\t\t\tUsage: &TokenUsage{\n\t\t\t\t\t\tPromptTokens:     7,\n\t\t\t\t\t\tCompletionTokens: 3,\n\t\t\t\t\t\tTotalTokens:      10,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, 3, len(msg.ResponseMeta.LogProbs.Content))\n\t\tassert.Equal(t, msgs[0].ResponseMeta.LogProbs.Content[0], msg.ResponseMeta.LogProbs.Content[0])\n\t\tassert.Equal(t, msgs[0].ResponseMeta.LogProbs.Content[1], msg.ResponseMeta.LogProbs.Content[1])\n\t\tassert.Equal(t, msgs[1].ResponseMeta.LogProbs.Content[0], msg.ResponseMeta.LogProbs.Content[2])\n\t})\n\n\tt.Run(\"fix unexpected setting ResponseMeta of the first element in slice after ConcatMessages\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole:    Assistant,\n\t\t\t\tContent: \"🚀\",\n\t\t\t\t//ResponseMeta: &ResponseMeta{},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole:         \"\",\n\t\t\t\tContent:      \"❤️\",\n\t\t\t\tResponseMeta: &ResponseMeta{},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: \"\",\n\t\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\t\tFinishReason: \"stop\",\n\t\t\t\t\tUsage: &TokenUsage{\n\t\t\t\t\t\tPromptTokens:     7,\n\t\t\t\t\t\tCompletionTokens: 3,\n\t\t\t\t\t\tTotalTokens:      10,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, msgs[2].ResponseMeta, msg.ResponseMeta)\n\t\tassert.Nil(t, msgs[0].ResponseMeta)\n\t})\n\n\tt.Run(\"concat assistant multi content\", func(t *testing.T) {\n\t\tbase64Audio1 := \"dGVzdF9hdWRpb18x\"\n\t\tbase64Audio2 := \"dGVzdF9hdWRpb18y\"\n\t\timageURL1 := \"https://example.com/image1.png\"\n\t\timageURL2 := \"https://example.com/image2.png\"\n\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Hello, \"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"world!\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio1}}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio2, MIMEType: \"audio/wav\"}}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{MessagePartCommon: MessagePartCommon{URL: &imageURL1}}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{MessagePartCommon: MessagePartCommon{URL: &imageURL2}}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\tmergedBase64Audio := base64Audio1 + base64Audio2\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"Hello, world!\"},\n\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &mergedBase64Audio, MIMEType: \"audio/wav\"}}},\n\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{MessagePartCommon: MessagePartCommon{URL: &imageURL1}}},\n\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{MessagePartCommon: MessagePartCommon{URL: &imageURL2}}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat assistant multi content with extra\", func(t *testing.T) {\n\t\tbase64Audio1 := \"dGVzdF9hdWRpb18x\"\n\t\tbase64Audio2 := \"dGVzdF9hdWRpb18y\"\n\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio1, Extra: map[string]any{\"key1\": \"val1\"}}}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio2, Extra: map[string]any{\"key2\": \"val2\"}}}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\tmergedBase64Audio := base64Audio1 + base64Audio2\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &mergedBase64Audio, Extra: map[string]any{\"key1\": \"val1\", \"key2\": \"val2\"}}}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat assistant multi content with single extra\", func(t *testing.T) {\n\t\tbase64Audio1 := \"dGVzdF9hdWRpb18x\"\n\t\tbase64Audio2 := \"dGVzdF9hdWRpb18y\"\n\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio1, Extra: map[string]any{\"key1\": \"val1\"}}}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio2}}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\tmergedBase64Audio := base64Audio1 + base64Audio2\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &mergedBase64Audio, Extra: map[string]any{\"key1\": \"val1\"}}}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat text parts with extra\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Hello \", Extra: map[string]any{\"key1\": \"val1\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"World\", Extra: map[string]any{\"key2\": \"val2\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"Hello World\", Extra: map[string]any{\"key1\": \"val1\", \"key2\": \"val2\"}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat text parts with single extra\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Hello \", Extra: map[string]any{\"key1\": \"val1\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"World\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"Hello World\", Extra: map[string]any{\"key1\": \"val1\"}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat reasoning parts with extra\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"First, \"}, Extra: map[string]any{\"key1\": \"val1\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"I need to think.\"}, Extra: map[string]any{\"key2\": \"val2\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"First, I need to think.\"}, Extra: map[string]any{\"key1\": \"val1\", \"key2\": \"val2\"}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat audio parts with outer extra\", func(t *testing.T) {\n\t\tbase64Audio1 := \"dGVzdF9hdWRpb18x\"\n\t\tbase64Audio2 := \"dGVzdF9hdWRpb18y\"\n\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio1}}, Extra: map[string]any{\"outer1\": \"val1\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio2}}, Extra: map[string]any{\"outer2\": \"val2\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\tmergedBase64Audio := base64Audio1 + base64Audio2\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &mergedBase64Audio}}, Extra: map[string]any{\"outer1\": \"val1\", \"outer2\": \"val2\"}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat reasoning parts\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"First, \"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"I need to think.\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Here is my answer.\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"First, I need to think.\"}},\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"Here is my answer.\"},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat reasoning parts with signature\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"Step 1: \"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"analyze.\", Signature: \"sig_abc\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \" Step 2: \", Signature: \"sig_xyz\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"conclude.\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"Step 1: analyze. Step 2: conclude.\", Signature: \"sig_xyz\"}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat with streaming meta index grouping\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"Think \"}, StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"more.\"}, StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Hello \", StreamingMeta: &MessageStreamingMeta{Index: 1}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"world!\", StreamingMeta: &MessageStreamingMeta{Index: 1}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: \"Think more.\"}, StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"Hello world!\", StreamingMeta: &MessageStreamingMeta{Index: 1}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat with different streaming meta index should not merge\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"First block \", StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Second block \", StreamingMeta: &MessageStreamingMeta{Index: 1}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"continues.\", StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"First block \", StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"Second block \", StreamingMeta: &MessageStreamingMeta{Index: 1}},\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"continues.\", StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t}\n\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat empty parts\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole:                     Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\t\tassert.Empty(t, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat single part no merge needed\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Single\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"Single\"},\n\t\t}\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat multiple consecutive text parts\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"One \"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Two \"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Three \"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Four\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"One Two Three Four\"},\n\t\t}\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat without streaming meta should not merge with streaming meta parts\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"No meta \"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"With meta\", StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedContent := []MessageOutputPart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"No meta \"},\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"With meta\", StreamingMeta: &MessageStreamingMeta{Index: 0}},\n\t\t}\n\t\tassert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent)\n\t})\n\n\tt.Run(\"concat multi content (deprecated)\", func(t *testing.T) {\n\t\tmsgs := []*Message{\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tMultiContent: []ChatMessagePart{\n\t\t\t\t\t{Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: \"image1.jpg\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tRole: Assistant,\n\t\t\t\tMultiContent: []ChatMessagePart{\n\t\t\t\t\t{Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: \"image2.jpg\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tmergedMsg, err := ConcatMessages(msgs)\n\t\tassert.NoError(t, err)\n\n\t\texpectedMultiContent := []ChatMessagePart{\n\t\t\t{Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: \"image1.jpg\"}},\n\t\t\t{Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: \"image2.jpg\"}},\n\t\t}\n\n\t\tassert.Equal(t, expectedMultiContent, mergedMsg.MultiContent)\n\t})\n}\n\nfunc TestConcatToolCalls(t *testing.T) {\n\tt.Run(\"atomic_field_in_first_chunk\", func(t *testing.T) {\n\t\tgivenToolCalls := []ToolCall{\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"tool_name\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\texpectedToolCall := ToolCall{\n\t\t\tIndex: generic.PtrOf(0),\n\t\t\tID:    \"tool_call_id\",\n\t\t\tType:  \"function\",\n\t\t\tFunction: FunctionCall{\n\t\t\t\tName:      \"tool_name\",\n\t\t\t\tArguments: \"call me please\",\n\t\t\t},\n\t\t}\n\n\t\ttc, err := concatToolCalls(givenToolCalls)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tc, 1)\n\t\tassert.EqualValues(t, expectedToolCall, tc[0])\n\t})\n\n\tt.Run(\"atomic_field_in_every_chunk\", func(t *testing.T) {\n\t\tgivenToolCalls := []ToolCall{\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"tool_name\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"tool_name\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\texpectedToolCall := ToolCall{\n\t\t\tIndex: generic.PtrOf(0),\n\t\t\tID:    \"tool_call_id\",\n\t\t\tType:  \"function\",\n\t\t\tFunction: FunctionCall{\n\t\t\t\tName:      \"tool_name\",\n\t\t\t\tArguments: \"call me please\",\n\t\t\t},\n\t\t}\n\n\t\ttc, err := concatToolCalls(givenToolCalls)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tc, 1)\n\t\tassert.EqualValues(t, expectedToolCall, tc[0])\n\t})\n\n\tt.Run(\"atomic_field_in_interval\", func(t *testing.T) {\n\t\tgivenToolCalls := []ToolCall{\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"\",\n\t\t\t\t\tArguments: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\texpectedToolCall := ToolCall{\n\t\t\tIndex: generic.PtrOf(0),\n\t\t\tID:    \"tool_call_id\",\n\t\t\tType:  \"function\",\n\t\t\tFunction: FunctionCall{\n\t\t\t\tName:      \"\",\n\t\t\t\tArguments: \"call me please\",\n\t\t\t},\n\t\t}\n\n\t\ttc, err := concatToolCalls(givenToolCalls)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, tc, 1)\n\t\tassert.EqualValues(t, expectedToolCall, tc[0])\n\t})\n\n\tt.Run(\"different_tool_id\", func(t *testing.T) {\n\t\tgivenToolCalls := []ToolCall{\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"tool_name\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id_1\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"tool_name\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := concatToolCalls(givenToolCalls)\n\t\tassert.ErrorContains(t, err, \"cannot concat ToolCalls with different tool id\")\n\t})\n\n\tt.Run(\"different_tool_type\", func(t *testing.T) {\n\t\tgivenToolCalls := []ToolCall{\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"tool_name\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function_1\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"tool_name\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := concatToolCalls(givenToolCalls)\n\t\tassert.ErrorContains(t, err, \"cannot concat ToolCalls with different tool type\")\n\t})\n\n\tt.Run(\"different_tool_name\", func(t *testing.T) {\n\t\tgivenToolCalls := []ToolCall{\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"tool_name\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"tool_name_1\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := concatToolCalls(givenToolCalls)\n\t\tassert.ErrorContains(t, err, \"cannot concat ToolCalls with different tool name\")\n\t})\n\n\tt.Run(\"multi_tool_call\", func(t *testing.T) {\n\t\tgivenToolCalls := []ToolCall{\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"\",\n\t\t\t\t\tArguments: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(1),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(1),\n\t\t\t\tID:    \"\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(1),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"\",\n\t\t\t\t\tArguments: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: nil,\n\t\t\t\tID:    \"22\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: nil,\n\t\t\t\tID:    \"44\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\texpectedToolCall := []ToolCall{\n\t\t\t{\n\t\t\t\tIndex: nil,\n\t\t\t\tID:    \"22\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: nil,\n\t\t\t\tID:    \"44\",\n\t\t\t\tType:  \"\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(0),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tIndex: generic.PtrOf(1),\n\t\t\t\tID:    \"tool_call_id\",\n\t\t\t\tType:  \"function\",\n\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\tName:      \"\",\n\t\t\t\t\tArguments: \"call me please\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\ttc, err := concatToolCalls(givenToolCalls)\n\t\tassert.NoError(t, err)\n\t\tassert.EqualValues(t, expectedToolCall, tc)\n\t})\n}\n\nfunc TestFormatMultiContent(t *testing.T) {\n\tvs := map[string]any{\n\t\t\"name\": \"eino\",\n\t\t\"url\":  \"https://example.com/img.png\",\n\t\t\"id\":   \"42\",\n\t}\n\n\tt.Run(\"empty input\", func(t *testing.T) {\n\t\tout, err := formatMultiContent(nil, vs, FString)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []ChatMessagePart{}, out)\n\t})\n\n\tt.Run(\"render text and urls with FString\", func(t *testing.T) {\n\t\tin := []ChatMessagePart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"hello {name}\"},\n\t\t\t{Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: \"{url}\"}},\n\t\t\t{Type: ChatMessagePartTypeAudioURL, AudioURL: &ChatMessageAudioURL{URL: \"http://audio/{id}.wav\"}},\n\t\t\t{Type: ChatMessagePartTypeVideoURL, VideoURL: &ChatMessageVideoURL{URL: \"http://video/{id}.mp4\"}},\n\t\t\t{Type: ChatMessagePartTypeFileURL, FileURL: &ChatMessageFileURL{URL: \"http://file/{id}.txt\"}},\n\t\t}\n\n\t\tout, err := formatMultiContent(in, vs, FString)\n\t\tassert.NoError(t, err)\n\t\tif assert.Len(t, out, len(in)) {\n\t\t\tassert.Equal(t, \"hello eino\", out[0].Text)\n\t\t\tassert.Equal(t, \"https://example.com/img.png\", out[1].ImageURL.URL)\n\t\t\tassert.Equal(t, \"http://audio/42.wav\", out[2].AudioURL.URL)\n\t\t\tassert.Equal(t, \"http://video/42.mp4\", out[3].VideoURL.URL)\n\t\t\tassert.Equal(t, \"http://file/42.txt\", out[4].FileURL.URL)\n\t\t}\n\t})\n\n\tt.Run(\"nil nested pointer should be skipped\", func(t *testing.T) {\n\t\tin := []ChatMessagePart{\n\t\t\t{Type: ChatMessagePartTypeImageURL, ImageURL: nil},\n\t\t\t{Type: ChatMessagePartTypeAudioURL, AudioURL: nil},\n\t\t\t{Type: ChatMessagePartTypeVideoURL, VideoURL: nil},\n\t\t\t{Type: ChatMessagePartTypeFileURL, FileURL: nil},\n\t\t}\n\t\tout, err := formatMultiContent(in, vs, FString)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, in, out)\n\t})\n\n\tt.Run(\"missing var should error in GoTemplate\", func(t *testing.T) {\n\t\tin := []ChatMessagePart{{Type: ChatMessagePartTypeText, Text: \"hi {{.who}}\"}}\n\t\t_, err := formatMultiContent(in, map[string]any{\"name\": \"x\"}, GoTemplate)\n\t\tassert.Error(t, err)\n\t})\n\n}\n\nfunc TestFormatUserInputMultiContent(t *testing.T) {\n\tmakeStrPtr := func(s string) *string { return &s }\n\n\tvs := map[string]any{\n\t\t\"x\":    \"world\",\n\t\t\"img\":  \"https://example.com/i.png\",\n\t\t\"b64\":  \"YmFzZTY0\",\n\t\t\"aid\":  \"99\",\n\t\t\"vid\":  \"77\",\n\t\t\"file\": \"abc\",\n\t}\n\n\tt.Run(\"empty input\", func(t *testing.T) {\n\t\tout, err := formatUserInputMultiContent(nil, vs, FString)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, []MessageInputPart{}, out)\n\t})\n\n\tt.Run(\"render text and both URL/Base64 for each type\", func(t *testing.T) {\n\t\tin := []MessageInputPart{\n\t\t\t{Type: ChatMessagePartTypeText, Text: \"hello {x}\"},\n\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageInputImage{MessagePartCommon: MessagePartCommon{URL: makeStrPtr(\"{img}\"), Base64Data: makeStrPtr(\"{b64}\")}}},\n\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageInputAudio{MessagePartCommon: MessagePartCommon{URL: makeStrPtr(\"http://a/{aid}.wav\"), Base64Data: makeStrPtr(\"{b64}\")}}},\n\t\t\t{Type: ChatMessagePartTypeVideoURL, Video: &MessageInputVideo{MessagePartCommon: MessagePartCommon{URL: makeStrPtr(\"http://v/{vid}.mp4\"), Base64Data: makeStrPtr(\"{b64}\")}}},\n\t\t\t{Type: ChatMessagePartTypeFileURL, File: &MessageInputFile{MessagePartCommon: MessagePartCommon{URL: makeStrPtr(\"/f/{file}.txt\"), Base64Data: makeStrPtr(\"{b64}\")}}},\n\t\t}\n\n\t\tout, err := formatUserInputMultiContent(in, vs, FString)\n\t\tassert.NoError(t, err)\n\t\tif assert.Len(t, out, len(in)) {\n\t\t\tassert.Equal(t, \"hello world\", out[0].Text)\n\t\t\tassert.Equal(t, \"https://example.com/i.png\", *out[1].Image.URL)\n\t\t\tassert.Equal(t, \"YmFzZTY0\", *out[1].Image.Base64Data)\n\t\t\tassert.Equal(t, \"http://a/99.wav\", *out[2].Audio.URL)\n\t\t\tassert.Equal(t, \"YmFzZTY0\", *out[2].Audio.Base64Data)\n\t\t\tassert.Equal(t, \"http://v/77.mp4\", *out[3].Video.URL)\n\t\t\tassert.Equal(t, \"YmFzZTY0\", *out[3].Video.Base64Data)\n\t\t\tassert.Equal(t, \"/f/abc.txt\", *out[4].File.URL)\n\t\t\tassert.Equal(t, \"YmFzZTY0\", *out[4].File.Base64Data)\n\t\t}\n\t})\n\n\tt.Run(\"empty string pointer should not be formatted\", func(t *testing.T) {\n\t\tempty := \"\"\n\t\tin := []MessageInputPart{\n\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageInputImage{MessagePartCommon: MessagePartCommon{URL: &empty, Base64Data: &empty}}},\n\t\t}\n\t\tout, err := formatUserInputMultiContent(in, vs, FString)\n\t\tassert.NoError(t, err)\n\t\tif assert.Len(t, out, 1) {\n\t\t\tassert.NotNil(t, out[0].Image.URL)\n\t\t\tassert.NotNil(t, out[0].Image.Base64Data)\n\t\t\tassert.Equal(t, \"\", *out[0].Image.URL)\n\t\t\tassert.Equal(t, \"\", *out[0].Image.Base64Data)\n\t\t}\n\t})\n}\n\nfunc TestConcatToolResults(t *testing.T) {\n\tt.Run(\"empty_chunks\", func(t *testing.T) {\n\t\tresult, err := ConcatToolResults([]*ToolResult{})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, result)\n\t\tassert.Empty(t, result.Parts)\n\t})\n\n\tt.Run(\"nil_chunks\", func(t *testing.T) {\n\t\tresult, err := ConcatToolResults([]*ToolResult{nil, nil})\n\t\tassert.NoError(t, err)\n\t\tassert.NotNil(t, result)\n\t\tassert.Empty(t, result.Parts)\n\t})\n\n\tt.Run(\"single_text_part\", func(t *testing.T) {\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Hello World\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 1)\n\t\tassert.Equal(t, ToolPartTypeText, result.Parts[0].Type)\n\t\tassert.Equal(t, \"Hello World\", result.Parts[0].Text)\n\t})\n\n\tt.Run(\"multiple_text_parts_merge\", func(t *testing.T) {\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Hello \"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"World\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"!\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 3)\n\n\t})\n\n\tt.Run(\"multiple_text_parts_merge_with_extra\", func(t *testing.T) {\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Hello \", Extra: map[string]any{\"k1\": \"v1\"}},\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"World\", Extra: map[string]any{\"k2\": \"v2\"}},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 1)\n\t\tassert.Equal(t, \"Hello World\", result.Parts[0].Text)\n\t\tassert.Equal(t, map[string]any{\"k1\": \"v1\", \"k2\": \"v2\"}, result.Parts[0].Extra)\n\t})\n\n\tt.Run(\"multiple_text_parts_merge_with_single_extra\", func(t *testing.T) {\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Hello \", Extra: map[string]any{\"k1\": \"v1\"}},\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"World\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 1)\n\t\tassert.Equal(t, \"Hello World\", result.Parts[0].Text)\n\t\tassert.Equal(t, map[string]any{\"k1\": \"v1\"}, result.Parts[0].Extra)\n\t})\n\n\tt.Run(\"cross_chunk_audio_conflict_error\", func(t *testing.T) {\n\t\tbase64Data1 := \"YXVkaW8x\"\n\t\tbase64Data2 := \"YXVkaW8y\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeAudio,\n\t\t\t\t\t\tAudio: &ToolOutputAudio{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tBase64Data: &base64Data1,\n\t\t\t\t\t\t\t\tMIMEType:   \"audio/wav\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeAudio,\n\t\t\t\t\t\tAudio: &ToolOutputAudio{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tBase64Data: &base64Data2,\n\t\t\t\t\t\t\t\tMIMEType:   \"audio/wav\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := ConcatToolResults(chunks)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"conflicting\")\n\t\tassert.Contains(t, err.Error(), \"audio\")\n\t})\n\n\tt.Run(\"mixed_types_no_merge\", func(t *testing.T) {\n\t\timageURL := \"https://example.com/image.png\"\n\t\tvideoURL := \"https://example.com/video.mp4\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Text part\"},\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeImage,\n\t\t\t\t\t\tImage: &ToolOutputImage{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &imageURL,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeVideo,\n\t\t\t\t\t\tVideo: &ToolOutputVideo{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &videoURL,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 3)\n\t\tassert.Equal(t, ToolPartTypeText, result.Parts[0].Type)\n\t\tassert.Equal(t, ToolPartTypeImage, result.Parts[1].Type)\n\t\tassert.Equal(t, ToolPartTypeVideo, result.Parts[2].Type)\n\t})\n\n\tt.Run(\"mixed_text_and_single_audio\", func(t *testing.T) {\n\t\tbase64Data1 := \"YXVkaW8x\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Part 1 \"},\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Part 2\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeAudio,\n\t\t\t\t\t\tAudio: &ToolOutputAudio{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tBase64Data: &base64Data1,\n\t\t\t\t\t\t\t\tMIMEType:   \"audio/wav\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \" Part 3\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 3)\n\n\t\tassert.Equal(t, ToolPartTypeText, result.Parts[0].Type)\n\t\tassert.Equal(t, \"Part 1 Part 2\", result.Parts[0].Text)\n\n\t\tassert.Equal(t, ToolPartTypeAudio, result.Parts[1].Type)\n\t\tassert.NotNil(t, result.Parts[1].Audio)\n\t\tassert.NotNil(t, result.Parts[1].Audio.Base64Data)\n\t\tassert.Equal(t, \"YXVkaW8x\", *result.Parts[1].Audio.Base64Data)\n\n\t\tassert.Equal(t, ToolPartTypeText, result.Parts[2].Type)\n\t\tassert.Equal(t, \" Part 3\", result.Parts[2].Text)\n\t})\n\n\tt.Run(\"cross_chunk_audio_url_and_base64_conflict_error\", func(t *testing.T) {\n\t\taudioURL := \"https://example.com/audio.wav\"\n\t\tbase64Data := \"YXVkaW8x\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeAudio,\n\t\t\t\t\t\tAudio: &ToolOutputAudio{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL:      &audioURL,\n\t\t\t\t\t\t\t\tMIMEType: \"audio/wav\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeAudio,\n\t\t\t\t\t\tAudio: &ToolOutputAudio{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tBase64Data: &base64Data,\n\t\t\t\t\t\t\t\tMIMEType:   \"audio/wav\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := ConcatToolResults(chunks)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"conflicting\")\n\t\tassert.Contains(t, err.Error(), \"audio\")\n\t})\n\n\tt.Run(\"single_audio_with_extra_fields\", func(t *testing.T) {\n\t\tbase64Data1 := \"YXVkaW8x\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeAudio,\n\t\t\t\t\t\tAudio: &ToolOutputAudio{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tBase64Data: &base64Data1,\n\t\t\t\t\t\t\t\tMIMEType:   \"audio/wav\",\n\t\t\t\t\t\t\t\tExtra: map[string]any{\n\t\t\t\t\t\t\t\t\t\"key1\": \"value1\",\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 1)\n\t\tassert.Equal(t, ToolPartTypeAudio, result.Parts[0].Type)\n\t\tassert.NotNil(t, result.Parts[0].Audio)\n\t\tassert.NotNil(t, result.Parts[0].Audio.Base64Data)\n\t\tassert.Equal(t, \"YXVkaW8x\", *result.Parts[0].Audio.Base64Data)\n\t\tassert.NotNil(t, result.Parts[0].Audio.Extra)\n\t\tassert.Equal(t, \"value1\", result.Parts[0].Audio.Extra[\"key1\"])\n\t})\n\n\tt.Run(\"cross_chunk_image_conflict_error\", func(t *testing.T) {\n\t\timageURL1 := \"https://example.com/image1.png\"\n\t\timageURL2 := \"https://example.com/image2.png\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeImage,\n\t\t\t\t\t\tImage: &ToolOutputImage{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &imageURL1,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeImage,\n\t\t\t\t\t\tImage: &ToolOutputImage{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &imageURL2,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := ConcatToolResults(chunks)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"conflicting\")\n\t\tassert.Contains(t, err.Error(), \"image\")\n\t})\n\n\tt.Run(\"cross_chunk_video_conflict_error\", func(t *testing.T) {\n\t\tvideoURL1 := \"https://example.com/video1.mp4\"\n\t\tvideoURL2 := \"https://example.com/video2.mp4\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeVideo,\n\t\t\t\t\t\tVideo: &ToolOutputVideo{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &videoURL1,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeVideo,\n\t\t\t\t\t\tVideo: &ToolOutputVideo{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &videoURL2,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := ConcatToolResults(chunks)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"conflicting\")\n\t\tassert.Contains(t, err.Error(), \"video\")\n\t})\n\n\tt.Run(\"cross_chunk_file_conflict_error\", func(t *testing.T) {\n\t\tfileURL1 := \"https://example.com/file1.pdf\"\n\t\tfileURL2 := \"https://example.com/file2.pdf\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeFile,\n\t\t\t\t\t\tFile: &ToolOutputFile{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &fileURL1,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeFile,\n\t\t\t\t\t\tFile: &ToolOutputFile{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &fileURL2,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t_, err := ConcatToolResults(chunks)\n\t\tassert.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"conflicting\")\n\t\tassert.Contains(t, err.Error(), \"file\")\n\t})\n\n\tt.Run(\"cross_chunk_text_not_merged\", func(t *testing.T) {\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Hello \"},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"World\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 2)\n\t\tassert.Equal(t, ToolPartTypeText, result.Parts[0].Type)\n\t\tassert.Equal(t, \"Hello \", result.Parts[0].Text)\n\t\tassert.Equal(t, ToolPartTypeText, result.Parts[1].Type)\n\t\tassert.Equal(t, \"World\", result.Parts[1].Text)\n\t})\n\n\tt.Run(\"same_chunk_text_merged\", func(t *testing.T) {\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"Hello \"},\n\t\t\t\t\t{Type: ToolPartTypeText, Text: \"World\"},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 1)\n\t\tassert.Equal(t, ToolPartTypeText, result.Parts[0].Type)\n\t\tassert.Equal(t, \"Hello World\", result.Parts[0].Text)\n\t})\n\n\tt.Run(\"different_non_text_types_across_chunks_allowed\", func(t *testing.T) {\n\t\timageURL := \"https://example.com/image.png\"\n\t\tvideoURL := \"https://example.com/video.mp4\"\n\t\tbase64Audio := \"YXVkaW8x\"\n\n\t\tchunks := []*ToolResult{\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeImage,\n\t\t\t\t\t\tImage: &ToolOutputImage{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &imageURL,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeVideo,\n\t\t\t\t\t\tVideo: &ToolOutputVideo{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tURL: &videoURL,\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t{\n\t\t\t\tParts: []ToolOutputPart{\n\t\t\t\t\t{\n\t\t\t\t\t\tType: ToolPartTypeAudio,\n\t\t\t\t\t\tAudio: &ToolOutputAudio{\n\t\t\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\t\t\tBase64Data: &base64Audio,\n\t\t\t\t\t\t\t\tMIMEType:   \"audio/wav\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tresult, err := ConcatToolResults(chunks)\n\t\tassert.NoError(t, err)\n\t\tassert.Len(t, result.Parts, 3)\n\t\tassert.Equal(t, ToolPartTypeImage, result.Parts[0].Type)\n\t\tassert.Equal(t, ToolPartTypeVideo, result.Parts[1].Type)\n\t\tassert.Equal(t, ToolPartTypeAudio, result.Parts[2].Type)\n\t})\n}\n\nfunc TestMessageString(t *testing.T) {\n\tt.Run(\"basic message\", func(t *testing.T) {\n\t\tmsg := &Message{\n\t\t\tRole:    User,\n\t\t\tContent: \"Hello, world!\",\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"user: Hello, world!\")\n\t})\n\n\tt.Run(\"message with UserInputMultiContent\", func(t *testing.T) {\n\t\timageURL := \"https://example.com/image.png\"\n\t\tmsg := &Message{\n\t\t\tRole:    User,\n\t\t\tContent: \"\",\n\t\t\tUserInputMultiContent: []MessageInputPart{\n\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Describe this image:\"},\n\t\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageInputImage{\n\t\t\t\t\tMessagePartCommon: MessagePartCommon{URL: &imageURL},\n\t\t\t\t}},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"user_input_multi_content:\")\n\t\tassert.Contains(t, result, \"[0] text: Describe this image:\")\n\t\tassert.Contains(t, result, \"[1] image: url=https://example.com/image.png\")\n\t})\n\n\tt.Run(\"message with AssistantGenMultiContent\", func(t *testing.T) {\n\t\tbase64Data := \"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==\"\n\t\tmsg := &Message{\n\t\t\tRole:    Assistant,\n\t\t\tContent: \"\",\n\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Here is the generated image:\"},\n\t\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{\n\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\tBase64Data: &base64Data,\n\t\t\t\t\t\tMIMEType:   \"image/png\",\n\t\t\t\t\t},\n\t\t\t\t}},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"assistant_gen_multi_content:\")\n\t\tassert.Contains(t, result, \"[0] text: Here is the generated image:\")\n\t\tassert.Contains(t, result, \"[1] image: base64[\")\n\t\tassert.Contains(t, result, \"mime=image/png\")\n\t})\n\n\tt.Run(\"message with MultiContent (deprecated)\", func(t *testing.T) {\n\t\tmsg := &Message{\n\t\t\tRole:    User,\n\t\t\tContent: \"\",\n\t\t\tMultiContent: []ChatMessagePart{\n\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"What is this?\"},\n\t\t\t\t{Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: \"https://example.com/photo.jpg\"}},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"multi_content:\")\n\t\tassert.Contains(t, result, \"[0] text: What is this?\")\n\t\tassert.Contains(t, result, \"[1] image_url: https://example.com/photo.jpg\")\n\t})\n\n\tt.Run(\"message with ToolCalls\", func(t *testing.T) {\n\t\tidx := 0\n\t\tmsg := &Message{\n\t\t\tRole:    Assistant,\n\t\t\tContent: \"\",\n\t\t\tToolCalls: []ToolCall{\n\t\t\t\t{\n\t\t\t\t\tIndex: &idx,\n\t\t\t\t\tID:    \"call_123\",\n\t\t\t\t\tType:  \"function\",\n\t\t\t\t\tFunction: FunctionCall{\n\t\t\t\t\t\tName:      \"get_weather\",\n\t\t\t\t\t\tArguments: `{\"location\": \"Beijing\"}`,\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"tool_calls:\")\n\t\tassert.Contains(t, result, \"index[0]:\")\n\t\tassert.Contains(t, result, \"get_weather\")\n\t})\n\n\tt.Run(\"tool message\", func(t *testing.T) {\n\t\tmsg := &Message{\n\t\t\tRole:       Tool,\n\t\t\tContent:    `{\"temperature\": 25}`,\n\t\t\tToolCallID: \"call_123\",\n\t\t\tToolName:   \"get_weather\",\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"tool: {\\\"temperature\\\": 25}\")\n\t\tassert.Contains(t, result, \"tool_call_id: call_123\")\n\t\tassert.Contains(t, result, \"tool_call_name: get_weather\")\n\t})\n\n\tt.Run(\"message with reasoning content\", func(t *testing.T) {\n\t\tmsg := &Message{\n\t\t\tRole:             Assistant,\n\t\t\tContent:          \"The answer is 42.\",\n\t\t\tReasoningContent: \"Let me think about this step by step...\",\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"reasoning content:\")\n\t\tassert.Contains(t, result, \"Let me think about this step by step...\")\n\t})\n\n\tt.Run(\"message with response meta\", func(t *testing.T) {\n\t\tmsg := &Message{\n\t\t\tRole:    Assistant,\n\t\t\tContent: \"Hello!\",\n\t\t\tResponseMeta: &ResponseMeta{\n\t\t\t\tFinishReason: \"stop\",\n\t\t\t\tUsage: &TokenUsage{\n\t\t\t\t\tPromptTokens:     10,\n\t\t\t\t\tCompletionTokens: 5,\n\t\t\t\t\tTotalTokens:      15,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"finish_reason: stop\")\n\t\tassert.Contains(t, result, \"usage:\")\n\t})\n\n\tt.Run(\"message with audio input\", func(t *testing.T) {\n\t\taudioURL := \"https://example.com/audio.wav\"\n\t\tmsg := &Message{\n\t\t\tRole: User,\n\t\t\tUserInputMultiContent: []MessageInputPart{\n\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageInputAudio{\n\t\t\t\t\tMessagePartCommon: MessagePartCommon{URL: &audioURL},\n\t\t\t\t}},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"[0] audio: url=https://example.com/audio.wav\")\n\t})\n\n\tt.Run(\"message with video input\", func(t *testing.T) {\n\t\tvideoURL := \"https://example.com/video.mp4\"\n\t\tmsg := &Message{\n\t\t\tRole: User,\n\t\t\tUserInputMultiContent: []MessageInputPart{\n\t\t\t\t{Type: ChatMessagePartTypeVideoURL, Video: &MessageInputVideo{\n\t\t\t\t\tMessagePartCommon: MessagePartCommon{URL: &videoURL},\n\t\t\t\t}},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"[0] video: url=https://example.com/video.mp4\")\n\t})\n\n\tt.Run(\"message with file input\", func(t *testing.T) {\n\t\tfileURL := \"https://example.com/document.pdf\"\n\t\tmsg := &Message{\n\t\t\tRole: User,\n\t\t\tUserInputMultiContent: []MessageInputPart{\n\t\t\t\t{Type: ChatMessagePartTypeFileURL, File: &MessageInputFile{\n\t\t\t\t\tMessagePartCommon: MessagePartCommon{URL: &fileURL},\n\t\t\t\t}},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"[0] file: url=https://example.com/document.pdf\")\n\t})\n\n\tt.Run(\"nil media parts\", func(t *testing.T) {\n\t\tmsg := &Message{\n\t\t\tRole: User,\n\t\t\tUserInputMultiContent: []MessageInputPart{\n\t\t\t\t{Type: ChatMessagePartTypeImageURL, Image: nil},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"[0] image: <nil>\")\n\t})\n\n\tt.Run(\"combined multi-content types\", func(t *testing.T) {\n\t\timageURL := \"https://example.com/image.png\"\n\t\tbase64Audio := \"YXVkaW9kYXRh\"\n\t\tmsg := &Message{\n\t\t\tRole:    User,\n\t\t\tContent: \"Main content\",\n\t\t\tUserInputMultiContent: []MessageInputPart{\n\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"User input text\"},\n\t\t\t\t{Type: ChatMessagePartTypeImageURL, Image: &MessageInputImage{\n\t\t\t\t\tMessagePartCommon: MessagePartCommon{URL: &imageURL},\n\t\t\t\t}},\n\t\t\t},\n\t\t\tAssistantGenMultiContent: []MessageOutputPart{\n\t\t\t\t{Type: ChatMessagePartTypeText, Text: \"Assistant output text\"},\n\t\t\t\t{Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{\n\t\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\t\tBase64Data: &base64Audio,\n\t\t\t\t\t\tMIMEType:   \"audio/wav\",\n\t\t\t\t\t},\n\t\t\t\t}},\n\t\t\t},\n\t\t}\n\t\tresult := msg.String()\n\t\tassert.Contains(t, result, \"user: Main content\")\n\t\tassert.Contains(t, result, \"user_input_multi_content:\")\n\t\tassert.Contains(t, result, \"assistant_gen_multi_content:\")\n\t})\n}\n\nfunc TestConvToolOutputPartToMessageInputPart(t *testing.T) {\n\tt.Run(\"text part\", func(t *testing.T) {\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType: ToolPartTypeText,\n\t\t\tText: \"test text\",\n\t\t\tExtra: map[string]any{\"key\": \"value\"},\n\t\t}\n\t\tresult, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, ChatMessagePartTypeText, result.Type)\n\t\tassert.Equal(t, \"test text\", result.Text)\n\t\tassert.Equal(t, map[string]any{\"key\": \"value\"}, result.Extra)\n\t})\n\n\tt.Run(\"image part\", func(t *testing.T) {\n\t\turl := \"https://example.com/image.png\"\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType: ToolPartTypeImage,\n\t\t\tImage: &ToolOutputImage{\n\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\tURL:      &url,\n\t\t\t\t\tMIMEType: \"image/png\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tExtra: map[string]any{\"img_key\": \"img_value\"},\n\t\t}\n\t\tresult, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, ChatMessagePartTypeImageURL, result.Type)\n\t\tassert.NotNil(t, result.Image)\n\t\tassert.Equal(t, url, *result.Image.URL)\n\t\tassert.Equal(t, \"image/png\", result.Image.MIMEType)\n\t\tassert.Equal(t, map[string]any{\"img_key\": \"img_value\"}, result.Extra)\n\t})\n\n\tt.Run(\"image part nil content\", func(t *testing.T) {\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType:  ToolPartTypeImage,\n\t\t\tImage: nil,\n\t\t}\n\t\tresult, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.Error(t, err)\n\t\tassert.ErrorContains(t, err, \"image content is nil\")\n\t\tassert.Equal(t, MessageInputPart{}, result)\n\t})\n\n\tt.Run(\"audio part\", func(t *testing.T) {\n\t\tbase64Data := \"dGVzdF9hdWRpbw==\"\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType: ToolPartTypeAudio,\n\t\t\tAudio: &ToolOutputAudio{\n\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\tBase64Data: &base64Data,\n\t\t\t\t\tMIMEType:   \"audio/wav\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, ChatMessagePartTypeAudioURL, result.Type)\n\t\tassert.NotNil(t, result.Audio)\n\t\tassert.Equal(t, base64Data, *result.Audio.Base64Data)\n\t\tassert.Equal(t, \"audio/wav\", result.Audio.MIMEType)\n\t})\n\n\tt.Run(\"audio part nil content\", func(t *testing.T) {\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType:  ToolPartTypeAudio,\n\t\t\tAudio: nil,\n\t\t}\n\t\t_, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.Error(t, err)\n\t\tassert.ErrorContains(t, err, \"audio content is nil\")\n\t})\n\n\tt.Run(\"video part\", func(t *testing.T) {\n\t\turl := \"https://example.com/video.mp4\"\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType: ToolPartTypeVideo,\n\t\t\tVideo: &ToolOutputVideo{\n\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\tURL:      &url,\n\t\t\t\t\tMIMEType: \"video/mp4\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\tresult, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, ChatMessagePartTypeVideoURL, result.Type)\n\t\tassert.NotNil(t, result.Video)\n\t\tassert.Equal(t, url, *result.Video.URL)\n\t\tassert.Equal(t, \"video/mp4\", result.Video.MIMEType)\n\t})\n\n\tt.Run(\"video part nil content\", func(t *testing.T) {\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType:  ToolPartTypeVideo,\n\t\t\tVideo: nil,\n\t\t}\n\t\t_, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.Error(t, err)\n\t\tassert.ErrorContains(t, err, \"video content is nil\")\n\t})\n\n\tt.Run(\"file part\", func(t *testing.T) {\n\t\turl := \"https://example.com/file.pdf\"\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType: ToolPartTypeFile,\n\t\t\tFile: &ToolOutputFile{\n\t\t\t\tMessagePartCommon: MessagePartCommon{\n\t\t\t\t\tURL:      &url,\n\t\t\t\t\tMIMEType: \"application/pdf\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tExtra: map[string]any{\"file_key\": \"file_value\"},\n\t\t}\n\t\tresult, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.NoError(t, err)\n\t\tassert.Equal(t, ChatMessagePartTypeFileURL, result.Type)\n\t\tassert.NotNil(t, result.File)\n\t\tassert.Equal(t, url, *result.File.URL)\n\t\tassert.Equal(t, \"application/pdf\", result.File.MIMEType)\n\t\tassert.Equal(t, map[string]any{\"file_key\": \"file_value\"}, result.Extra)\n\t})\n\n\tt.Run(\"file part nil content\", func(t *testing.T) {\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType: ToolPartTypeFile,\n\t\t\tFile: nil,\n\t\t}\n\t\t_, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.Error(t, err)\n\t\tassert.ErrorContains(t, err, \"file content is nil\")\n\t})\n\n\tt.Run(\"unknown type\", func(t *testing.T) {\n\t\ttoolPart := ToolOutputPart{\n\t\t\tType: \"unknown_type\",\n\t\t}\n\t\t_, err := convToolOutputPartToMessageInputPart(toolPart)\n\t\tassert.Error(t, err)\n\t\tassert.ErrorContains(t, err, \"unknown tool part type\")\n\t})\n}\n"
  },
  {
    "path": "schema/select.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nconst maxSelectNum = 5\n\nfunc receiveN[T any](chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {\n\treturn []func(chosenList []int, ss []*stream[T]) (index int, item *streamItem[T], ok bool){\n\t\tnil,\n\t\tfunc(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {\n\t\t\titem, ok := <-ss[chosenList[0]].items\n\t\t\treturn chosenList[0], &item, ok\n\t\t},\n\t\tfunc(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {\n\t\t\tselect {\n\t\t\tcase item, ok := <-ss[chosenList[0]].items:\n\t\t\t\treturn chosenList[0], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[1]].items:\n\t\t\t\treturn chosenList[1], &item, ok\n\t\t\t}\n\t\t},\n\t\tfunc(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {\n\t\t\tselect {\n\t\t\tcase item, ok := <-ss[chosenList[0]].items:\n\t\t\t\treturn chosenList[0], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[1]].items:\n\t\t\t\treturn chosenList[1], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[2]].items:\n\t\t\t\treturn chosenList[2], &item, ok\n\t\t\t}\n\t\t},\n\t\tfunc(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {\n\t\t\tselect {\n\t\t\tcase item, ok := <-ss[chosenList[0]].items:\n\t\t\t\treturn chosenList[0], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[1]].items:\n\t\t\t\treturn chosenList[1], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[2]].items:\n\t\t\t\treturn chosenList[2], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[3]].items:\n\t\t\t\treturn chosenList[3], &item, ok\n\t\t\t}\n\t\t},\n\t\tfunc(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {\n\t\t\tselect {\n\t\t\tcase item, ok := <-ss[chosenList[0]].items:\n\t\t\t\treturn chosenList[0], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[1]].items:\n\t\t\t\treturn chosenList[1], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[2]].items:\n\t\t\t\treturn chosenList[2], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[3]].items:\n\t\t\t\treturn chosenList[3], &item, ok\n\t\t\tcase item, ok := <-ss[chosenList[4]].items:\n\t\t\t\treturn chosenList[4], &item, ok\n\t\t\t}\n\t\t},\n\t}[len(chosenList)](chosenList, ss)\n}\n"
  },
  {
    "path": "schema/serialization.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"encoding/gob\"\n\t\"reflect\"\n\n\t\"github.com/cloudwego/eino/internal/generic\"\n\t\"github.com/cloudwego/eino/internal/serialization\"\n)\n\nfunc init() {\n\tRegisterName[Message](\"_eino_message\")\n\tRegisterName[[]*Message](\"_eino_message_slice\")\n\tRegisterName[Document](\"_eino_document\")\n\tRegisterName[RoleType](\"_eino_role_type\")\n\tRegisterName[ToolCall](\"_eino_tool_call\")\n\tRegisterName[FunctionCall](\"_eino_function_call\")\n\tRegisterName[ResponseMeta](\"_eino_response_meta\")\n\tRegisterName[TokenUsage](\"_eino_token_usage\")\n\tRegisterName[LogProbs](\"_eino_log_probs\")\n\tRegisterName[ChatMessagePart](\"_eino_chat_message_part\")\n\tRegisterName[ChatMessagePartType](\"_eino_chat_message_type\")\n\tRegisterName[ChatMessageImageURL](\"_eino_chat_message_image_url\")\n\tRegisterName[ChatMessageAudioURL](\"_eino_chat_message_audio_url\")\n\tRegisterName[ChatMessageVideoURL](\"_eino_chat_message_video_url\")\n\tRegisterName[ChatMessageFileURL](\"_eino_chat_message_file_url\")\n\tRegisterName[MessageInputPart](\"_eino_message_input_part\")\n\tRegisterName[MessageInputImage](\"_eino_message_input_image\")\n\tRegisterName[MessageInputAudio](\"_eino_message_input_audio\")\n\tRegisterName[MessageInputVideo](\"_eino_message_input_video\")\n\tRegisterName[MessageInputFile](\"_eino_message_input_file\")\n\tRegisterName[MessageOutputPart](\"_eino_message_output_part\")\n\tRegisterName[MessageOutputImage](\"_eino_message_output_image\")\n\tRegisterName[MessageOutputAudio](\"_eino_message_output_audio\")\n\tRegisterName[MessageOutputVideo](\"_eino_message_output_video\")\n\tRegisterName[MessagePartCommon](\"_eino_message_part_common\")\n\tRegisterName[ImageURLDetail](\"_eino_image_url_detail\")\n\tRegisterName[PromptTokenDetails](\"_eino_prompt_token_details\")\n}\n\n// RegisterName registers a type with a specific name for serialization. This is\n// required for any type you intend to persist in a graph or ADK checkpoint.\n// Use this function to maintain backward compatibility by mapping a type to a\n// previously used name. For new types, `Register` is preferred.\n//\n// It is recommended to call this in an `init()` function in the file where the\n// type is declared.\n//\n// What to Register:\n//   - Top-level types used as state (e.g., structs).\n//   - Concrete types that are assigned to interface fields.\n//\n// What NOT to Register:\n//   - Struct fields with concrete types (e.g., `string`, `int`, other structs).\n//     These are inferred via reflection.\n//\n// Serialization Rules:\n//\n// The serialization behavior is based on Go's standard `encoding/gob` package.\n// See https://pkg.go.dev/encoding/gob for detailed rules.\n//   - Only exported struct fields are serialized.\n//   - Functions and channels are not supported and will be ignored.\n//\n// This function panics if registration fails.\nfunc RegisterName[T any](name string) {\n\tgob.RegisterName(name, generic.NewInstance[T]())\n\n\terr := serialization.GenericRegister[T](name)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n}\n\nfunc getTypeName(rt reflect.Type) string {\n\tname := rt.String()\n\n\t// But for named types (or pointers to them), qualify with import path.\n\t// Dereference one pointer looking for a named type.\n\tstar := \"\"\n\tif rt.Name() == \"\" {\n\t\tif pt := rt; pt.Kind() == reflect.Pointer {\n\t\t\tstar = \"*\"\n\t\t\trt = pt.Elem()\n\t\t}\n\t}\n\tif rt.Name() != \"\" {\n\t\tif rt.PkgPath() == \"\" {\n\t\t\tname = star + rt.Name()\n\t\t} else {\n\t\t\tname = star + rt.PkgPath() + \".\" + rt.Name()\n\t\t}\n\t}\n\treturn name\n}\n\n// Register registers a type for serialization. This is required for any type\n// you intend to persist in a graph or ADK checkpoint. It automatically determines\n// the type name and is the recommended method for registering new types.\n//\n// It is recommended to call this in an `init()` function in the file where the\n// type is declared.\n//\n// What to Register:\n//   - Top-level types used as state (e.g., structs).\n//   - Concrete types that are assigned to interface fields.\n//\n// What NOT to Register:\n//   - Struct fields with concrete types (e.g., `string`, `int`, other structs).\n//     These are inferred via reflection.\n//\n// Serialization Rules:\n//\n// The serialization behavior is based on Go's standard `encoding/gob` package.\n// See https://pkg.go.dev/encoding/gob for detailed rules.\n//   - Only exported struct fields are serialized.\n//   - Functions and channels are not supported and will be ignored.\n//\n// This function panics if registration fails.\nfunc Register[T any]() {\n\tvalue := generic.NewInstance[T]()\n\n\tgob.Register(value)\n\n\tname := getTypeName(reflect.TypeOf(value))\n\n\terr := serialization.GenericRegister[T](name)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n}\n"
  },
  {
    "path": "schema/serialization_test.go",
    "content": "/*\n * Copyright 2025 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"bytes\"\n\t\"encoding/gob\"\n\t\"fmt\"\n\t\"reflect\"\n\t\"testing\"\n\n\t\"github.com/google/uuid\"\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/internal/serialization\"\n)\n\ntype testStruct struct{}\n\nfunc TestGetTypeName(t *testing.T) {\n\ttype localNamedType struct{}\n\n\ttestCases := []struct {\n\t\tname     string\n\t\tinput    reflect.Type\n\t\texpected string\n\t}{\n\t\t{\n\t\t\tname:     \"named type from current package\",\n\t\t\tinput:    reflect.TypeOf(testStruct{}),\n\t\t\texpected: \"github.com/cloudwego/eino/schema.testStruct\",\n\t\t},\n\t\t{\n\t\t\tname:     \"pointer to named type from current package\",\n\t\t\tinput:    reflect.TypeOf(&testStruct{}),\n\t\t\texpected: \"*github.com/cloudwego/eino/schema.testStruct\",\n\t\t},\n\t\t{\n\t\t\tname:     \"unnamed map type\",\n\t\t\tinput:    reflect.TypeOf(map[string]int{}),\n\t\t\texpected: \"map[string]int\",\n\t\t},\n\t\t{\n\t\t\tname:     \"pointer to unnamed map type\",\n\t\t\tinput:    reflect.TypeOf(new(map[string]int)),\n\t\t\texpected: \"*map[string]int\",\n\t\t},\n\t\t{\n\t\t\tname:     \"built-in type\",\n\t\t\tinput:    reflect.TypeOf(0),\n\t\t\texpected: \"int\",\n\t\t},\n\t\t{\n\t\t\tname:     \"pointer to built-in type\",\n\t\t\tinput:    reflect.TypeOf(new(int)),\n\t\t\texpected: \"*int\",\n\t\t},\n\t\t{\n\t\t\tname:     \"named type from standard library\",\n\t\t\tinput:    reflect.TypeOf(bytes.Buffer{}),\n\t\t\texpected: \"bytes.Buffer\",\n\t\t},\n\t\t{\n\t\t\tname:     \"pointer to named type from standard library\",\n\t\t\tinput:    reflect.TypeOf(&bytes.Buffer{}),\n\t\t\texpected: \"*bytes.Buffer\",\n\t\t},\n\t\t{\n\t\t\tname:     \"local named type\",\n\t\t\tinput:    reflect.TypeOf(localNamedType{}),\n\t\t\texpected: \"github.com/cloudwego/eino/schema.localNamedType\",\n\t\t},\n\t\t{\n\t\t\tname:     \"pointer to local named type\",\n\t\t\tinput:    reflect.TypeOf(&localNamedType{}),\n\t\t\texpected: \"*github.com/cloudwego/eino/schema.localNamedType\",\n\t\t},\n\t}\n\n\tfor _, tc := range testCases {\n\t\tt.Run(tc.name, func(t *testing.T) {\n\t\t\tactual := getTypeName(tc.input)\n\t\t\tif actual != tc.expected {\n\t\t\t\tt.Errorf(\"getTypeName() got %q, want %q\", actual, tc.expected)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestRegister(t *testing.T) {\n\ttype testStruct1 struct {\n\t\tA any\n\t\tB any\n\t\tC any\n\t\tD any\n\t\tE any\n\t\tF any\n\t}\n\n\ttype testStruct2 struct{}\n\n\tRegister[*testStruct1]()\n\tRegister[*testStruct2]()\n\tRegister[[]Message]()\n\tRegister[[]*testStruct2]()\n\tRegister[[]testStruct2]()\n\n\tt1 := testStruct1{A: []*Message{{}}, B: []Message{{}}, C: []*testStruct2{{}}, D: []testStruct2{{}},\n\t\tE: &testStruct1{}, F: []int{1}}\n\n\tin := &serialization.InternalSerializer{}\n\tmar, err := in.Marshal(t1)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\tvar t2 testStruct1\n\terr = in.Unmarshal(mar, &t2)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\tassert.Equal(t, t1, t2)\n\n\tbuf := new(bytes.Buffer)\n\terr = gob.NewEncoder(buf).Encode(t1)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\terr = gob.NewDecoder(buf).Decode(&t2)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\tassert.Equal(t, t1, t2)\n\n\tf := func() (err error) {\n\t\tdefer func() {\n\t\t\tif r := recover(); r != nil {\n\t\t\t\terr = fmt.Errorf(\"panic: %v\", r)\n\t\t\t}\n\t\t}()\n\n\t\tRegister[[]int]()\n\t\tRegister[map[string]any]()\n\t\tRegister[[]*testStruct1]()\n\t\tRegister[[]testStruct1]()\n\n\t\treturn nil\n\t}\n\n\terr = f()\n\tassert.NoError(t, err)\n}\n\n// TestRegisterStructWithUUIDField reproduces issue #607\n// uuid.UUID is a [16]byte array. Prior to the fix, calling schema.RegisterName on\n// a struct with a uuid.UUID field would panic during deserialization.\nfunc TestRegisterStructWithUUIDField(t *testing.T) {\n\ttype Item struct {\n\t\tID uuid.UUID\n\t}\n\n\tRegisterName[Item](\"test_item\")\n\n\toriginal := Item{\n\t\tID: uuid.MustParse(\"6ba7b810-9dad-11d1-80b4-00c04fd430c8\"),\n\t}\n\n\ts := &serialization.InternalSerializer{}\n\n\tdata, err := s.Marshal(original)\n\tassert.NoError(t, err)\n\n\tvar result Item\n\terr = s.Unmarshal(data, &result)\n\tassert.NoError(t, err)\n\n\tassert.Equal(t, original.ID, result.ID)\n}\n"
  },
  {
    "path": "schema/stream.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"reflect\"\n\t\"runtime\"\n\t\"runtime/debug\"\n\t\"sync\"\n\t\"sync/atomic\"\n\n\t\"github.com/cloudwego/eino/internal/safe\"\n)\n\n// ErrNoValue is a sentinel returned from the convert function passed to\n// [StreamReaderWithConvert] to skip a stream element — the element is dropped\n// and the next one is read without surfacing an error to the caller.\n//\n// Use it to filter out empty or irrelevant chunks:\n//\n//\toutStream = schema.StreamReaderWithConvert(s,\n//\t    func(src string) (string, error) {\n//\t        if len(src) == 0 {\n//\t            return \"\", schema.ErrNoValue // skip empty chunks\n//\t        }\n//\t        return src, nil\n//\t    })\n//\n// DO NOT use ErrNoValue in any other context.\nvar ErrNoValue = errors.New(\"no value\")\n\n// ErrRecvAfterClosed indicates that StreamReader.Recv was unexpectedly called after StreamReader.Close.\n// This error should not occur during normal use of StreamReader.Recv. If it does, please check your application code.\nvar ErrRecvAfterClosed = errors.New(\"recv after stream closed\")\n\n// SourceEOF represents an EOF error from a specific source stream.\n// It is only returned by the method Recv() of StreamReader created\n// with MergeNamedStreamReaders when one of its source streams reaches EOF.\ntype SourceEOF struct {\n\tsourceName string\n}\n\nfunc (e *SourceEOF) Error() string {\n\treturn fmt.Sprintf(\"EOF from source stream: %s\", e.sourceName)\n}\n\n// GetSourceName extracts the source stream name from a SourceEOF error.\n// It returns the source name and a boolean indicating whether the error was a SourceEOF.\n// If the error is not a SourceEOF, it returns an empty string and false.\nfunc GetSourceName(err error) (string, bool) {\n\tvar sErr *SourceEOF\n\tif errors.As(err, &sErr) {\n\t\treturn sErr.sourceName, true\n\t}\n\n\treturn \"\", false\n}\n\n// Pipe creates a new stream with the given capacity that represented with StreamWriter and StreamReader.\n// The capacity is the maximum number of items that can be buffered in the stream.\n// e.g.\n//\n//\tsr, sw := schema.Pipe[string](3)\n//\tgo func() { // send data\n//\t\tdefer sw.Close()\n//\t\tfor i := 0; i < 10; i++ {\n//\t\t\tsw.Send(i, nil)\n//\t\t}\n//\t}\n//\n//\tdefer sr.Close()\n//\tfor chunk, err := sr.Recv() {\n//\t\tif errors.Is(err, io.EOF) {\n//\t\t\tbreak\n//\t\t}\n//\t\tfmt.Println(chunk)\n//\t}\nfunc Pipe[T any](cap int) (*StreamReader[T], *StreamWriter[T]) {\n\tstm := newStream[T](cap)\n\treturn stm.asReader(), &StreamWriter[T]{stm: stm}\n}\n\n// StreamWriter the sender of a stream.\n// created by Pipe function.\n// eg.\n//\n//\tsr, sw := schema.Pipe[string](3)\n//\tgo func() { // send data\n//\t\tdefer sw.Close()\n//\t\tfor i := 0; i < 10; i++ {\n//\t\t\tsw.Send(i, nil)\n//\t\t}\n//\t}\ntype StreamWriter[T any] struct {\n\tstm *stream[T]\n}\n\n// Send sends a value to the stream.\n// e.g.\n//\n//\tclosed := sw.Send(i, nil)\n//\tif closed {\n//\t\t// the stream is closed\n//\t}\nfunc (sw *StreamWriter[T]) Send(chunk T, err error) (closed bool) {\n\treturn sw.stm.send(chunk, err)\n}\n\n// Close notify the receiver that the stream sender has finished.\n// The stream receiver will get an error of io.EOF from StreamReader.Recv().\n// Notice: always remember to call Close() after sending all data.\n// eg.\n//\n//\tdefer sw.Close()\n//\tfor i := 0; i < 10; i++ {\n//\t\tsw.Send(i, nil)\n//\t}\nfunc (sw *StreamWriter[T]) Close() {\n\tsw.stm.closeSend()\n}\n\n// StreamReader is the consumer side of an Eino stream.\n//\n// A StreamReader is read-once: only one goroutine should call Recv, and the\n// reader must be closed exactly once (whether the loop finishes normally or\n// exits early via break or return).\n//\n// Typical usage:\n//\n//\tdefer sr.Close() // always close, even after io.EOF\n//\tfor {\n//\t    chunk, err := sr.Recv()\n//\t    if errors.Is(err, io.EOF) {\n//\t        break\n//\t    }\n//\t    if err != nil {\n//\t        return err\n//\t    }\n//\t    process(chunk)\n//\t}\n//\n// To fan-out a single stream to N independent consumers, call [StreamReader.Copy]\n// before any Recv; the original reader becomes unusable after the call.\n//\n// StreamReaders are created by [Pipe], [StreamReaderFromArray],\n// [MergeStreamReaders], [MergeNamedStreamReaders], and [StreamReaderWithConvert].\ntype StreamReader[T any] struct {\n\ttyp readerType\n\n\tst *stream[T]\n\n\tar *arrayReader[T]\n\n\tmsr *multiStreamReader[T]\n\n\tsrw *streamReaderWithConvert[T]\n\n\tcsr *childStreamReader[T]\n}\n\n// Recv receives a value from the stream.\n// eg.\n//\n//\tfor chunk, err := sr.Recv() {\n//\t\tif errors.Is(err, io.EOF) {\n//\t\t\tbreak\n//\t\t}\n//\t\tif err != nil {\n//\t\tfmt.Println(chunk)\n//\t}\nfunc (sr *StreamReader[T]) Recv() (T, error) {\n\tswitch sr.typ {\n\tcase readerTypeStream:\n\t\treturn sr.st.recv()\n\tcase readerTypeArray:\n\t\treturn sr.ar.recv()\n\tcase readerTypeMultiStream:\n\t\treturn sr.msr.recv()\n\tcase readerTypeWithConvert:\n\t\treturn sr.srw.recv()\n\tcase readerTypeChild:\n\t\treturn sr.csr.recv()\n\tdefault:\n\t\tpanic(\"impossible\")\n\t}\n}\n\n// Close safely closes the StreamReader.\n// It should be called only once, as multiple calls may not work as expected.\n// Notice: always remember to call Close() after using Recv().\n// e.g.\n//\n//\tdefer sr.Close()\n//\n//\tfor chunk, err := sr.Recv() {\n//\t\tif errors.Is(err, io.EOF) {\n//\t\t\tbreak\n//\t\t}\n//\t\tfmt.Println(chunk)\n//\t}\nfunc (sr *StreamReader[T]) Close() {\n\tswitch sr.typ {\n\tcase readerTypeStream:\n\t\tsr.st.closeRecv()\n\tcase readerTypeArray:\n\n\tcase readerTypeMultiStream:\n\t\tsr.msr.close()\n\tcase readerTypeWithConvert:\n\t\tsr.srw.close()\n\tcase readerTypeChild:\n\t\tsr.csr.close()\n\tdefault:\n\t\tpanic(\"impossible\")\n\t}\n}\n\n// Copy creates n independent StreamReaders that each receive every element of\n// the original stream. The original StreamReader becomes unusable after Copy.\n//\n// Use Copy when two or more pipeline branches need the same stream —\n// for example, when a stream must be fed to both a callback handler and the\n// next node in a graph:\n//\n//\tcopies := sr.Copy(2)\n//\tsr1, sr2 := copies[0], copies[1]\n//\tdefer sr1.Close()\n//\tdefer sr2.Close()\n//\n//\t// sr1 and sr2 independently read the same elements\n//\n// n must be at least 1. If n < 2, the original reader is returned unchanged.\nfunc (sr *StreamReader[T]) Copy(n int) []*StreamReader[T] {\n\tif n < 2 {\n\t\treturn []*StreamReader[T]{sr}\n\t}\n\n\tif sr.typ == readerTypeArray {\n\t\tret := make([]*StreamReader[T], n)\n\t\tfor i, ar := range sr.ar.copy(n) {\n\t\t\tret[i] = &StreamReader[T]{typ: readerTypeArray, ar: ar}\n\t\t}\n\t\treturn ret\n\t}\n\n\treturn copyStreamReaders[T](sr, n)\n}\n\n// SetAutomaticClose sets the StreamReader to automatically close when it's no longer reachable and ready to be GCed.\n// NOT concurrency safe.\nfunc (sr *StreamReader[T]) SetAutomaticClose() {\n\tswitch sr.typ {\n\tcase readerTypeStream:\n\t\tif !sr.st.automaticClose {\n\t\t\tsr.st.automaticClose = true\n\t\t\tvar flag uint32\n\t\t\tsr.st.closedFlag = &flag\n\t\t\truntime.SetFinalizer(sr, func(s *StreamReader[T]) {\n\t\t\t\ts.Close()\n\t\t\t})\n\t\t}\n\tcase readerTypeMultiStream:\n\t\tfor _, s := range sr.msr.nonClosedStreams() {\n\t\t\tif !s.automaticClose {\n\t\t\t\ts.automaticClose = true\n\t\t\t\tvar flag uint32\n\t\t\t\ts.closedFlag = &flag\n\t\t\t\truntime.SetFinalizer(s, func(st *stream[T]) {\n\t\t\t\t\tst.closeRecv()\n\t\t\t\t})\n\t\t\t}\n\t\t}\n\tcase readerTypeChild:\n\t\tparent := sr.csr.parent.sr\n\t\tparent.SetAutomaticClose()\n\tcase readerTypeWithConvert:\n\t\tsr.srw.sr.SetAutomaticClose()\n\tcase readerTypeArray:\n\t\t// no need to clean up\n\tdefault:\n\t}\n}\n\nfunc (sr *StreamReader[T]) recvAny() (any, error) {\n\treturn sr.Recv()\n}\n\nfunc (sr *StreamReader[T]) copyAny(n int) []iStreamReader {\n\tret := make([]iStreamReader, n)\n\n\tsrs := sr.Copy(n)\n\n\tfor i := 0; i < n; i++ {\n\t\tret[i] = srs[i]\n\t}\n\n\treturn ret\n}\n\nfunc arrToStream[T any](arr []T) *stream[T] {\n\ts := newStream[T](len(arr))\n\tfor i := range arr {\n\t\ts.send(arr[i], nil)\n\t}\n\ts.closeSend()\n\n\treturn s\n}\n\nfunc (sr *StreamReader[T]) toStream() *stream[T] {\n\tswitch sr.typ {\n\tcase readerTypeStream:\n\t\treturn sr.st\n\tcase readerTypeArray:\n\t\treturn sr.ar.toStream()\n\tcase readerTypeMultiStream:\n\t\treturn sr.msr.toStream()\n\tcase readerTypeWithConvert:\n\t\treturn sr.srw.toStream()\n\tcase readerTypeChild:\n\t\treturn sr.csr.toStream()\n\tdefault:\n\t\tpanic(\"impossible\")\n\t}\n}\n\ntype readerType int\n\nconst (\n\treaderTypeStream readerType = iota\n\treaderTypeArray\n\treaderTypeMultiStream\n\treaderTypeWithConvert\n\treaderTypeChild\n)\n\ntype iStreamReader interface {\n\trecvAny() (any, error)\n\tcopyAny(int) []iStreamReader\n\tClose()\n\tSetAutomaticClose()\n}\n\n// stream is a channel-based stream with 1 sender and 1 receiver.\n// The sender calls closeSend() to notify the receiver that the stream sender has finished.\n// The receiver calls closeRecv() to notify the sender that the receiver stop receiving.\ntype stream[T any] struct {\n\titems chan streamItem[T]\n\n\tclosed chan struct{}\n\n\tautomaticClose bool\n\tclosedFlag     *uint32 // 0 = not closed, 1 = closed, only used when automaticClose is set\n}\n\ntype streamItem[T any] struct {\n\tchunk T\n\terr   error\n}\n\nfunc newStream[T any](cap int) *stream[T] {\n\treturn &stream[T]{\n\t\titems:  make(chan streamItem[T], cap),\n\t\tclosed: make(chan struct{}),\n\t}\n}\n\nfunc (s *stream[T]) asReader() *StreamReader[T] {\n\treturn &StreamReader[T]{typ: readerTypeStream, st: s}\n}\n\nfunc (s *stream[T]) recv() (chunk T, err error) {\n\titem, ok := <-s.items\n\n\tif !ok {\n\t\titem.err = io.EOF\n\t}\n\n\treturn item.chunk, item.err\n}\n\nfunc (s *stream[T]) send(chunk T, err error) (closed bool) {\n\t// if the stream is closed, return immediately\n\tselect {\n\tcase <-s.closed:\n\t\treturn true\n\tdefault:\n\t}\n\n\titem := streamItem[T]{chunk, err}\n\n\tselect {\n\tcase <-s.closed:\n\t\treturn true\n\tcase s.items <- item:\n\t\treturn false\n\t}\n}\n\nfunc (s *stream[T]) closeSend() {\n\tclose(s.items)\n}\n\nfunc (s *stream[T]) closeRecv() {\n\tif s.automaticClose {\n\t\tif atomic.CompareAndSwapUint32(s.closedFlag, 0, 1) {\n\t\t\tclose(s.closed)\n\t\t}\n\t\treturn\n\t}\n\n\tclose(s.closed)\n}\n\n// StreamReaderFromArray creates a StreamReader from a given slice of elements.\n// It takes an array of type T and returns a pointer to a StreamReader[T].\n// This allows for streaming the elements of the array in a controlled manner.\n// eg.\n//\n//\tsr := schema.StreamReaderFromArray([]int{1, 2, 3})\n//\tdefer sr.Close()\n//\n//\tfor chunk, err := sr.Recv() {\n//\t\tfmt.Println(chunk)\n//\t}\nfunc StreamReaderFromArray[T any](arr []T) *StreamReader[T] {\n\treturn &StreamReader[T]{ar: &arrayReader[T]{arr: arr}, typ: readerTypeArray}\n}\n\ntype arrayReader[T any] struct {\n\tarr   []T\n\tindex int\n}\n\nfunc (ar *arrayReader[T]) recv() (T, error) {\n\tif ar.index < len(ar.arr) {\n\t\tret := ar.arr[ar.index]\n\t\tar.index++\n\n\t\treturn ret, nil\n\t}\n\n\tvar t T\n\treturn t, io.EOF\n}\n\nfunc (ar *arrayReader[T]) copy(n int) []*arrayReader[T] {\n\tret := make([]*arrayReader[T], n)\n\n\tfor i := 0; i < n; i++ {\n\t\tret[i] = &arrayReader[T]{\n\t\t\tarr:   ar.arr,\n\t\t\tindex: ar.index,\n\t\t}\n\t}\n\n\treturn ret\n}\n\nfunc (ar *arrayReader[T]) toStream() *stream[T] {\n\treturn arrToStream(ar.arr[ar.index:])\n}\n\ntype multiArrayReader[T any] struct {\n\tars   []*arrayReader[T]\n\tindex int\n}\n\ntype multiStreamReader[T any] struct {\n\tsts []*stream[T]\n\n\titemsCases []reflect.SelectCase\n\n\tnonClosed []int\n\n\tsourceReaderNames []string\n}\n\nfunc newMultiStreamReader[T any](sts []*stream[T]) *multiStreamReader[T] {\n\tvar itemsCases []reflect.SelectCase\n\tif len(sts) > maxSelectNum {\n\t\titemsCases = make([]reflect.SelectCase, len(sts))\n\t\tfor i, st := range sts {\n\t\t\titemsCases[i] = reflect.SelectCase{\n\t\t\t\tDir:  reflect.SelectRecv,\n\t\t\t\tChan: reflect.ValueOf(st.items),\n\t\t\t}\n\t\t}\n\t}\n\n\tnonClosed := make([]int, len(sts))\n\tfor i := range sts {\n\t\tnonClosed[i] = i\n\t}\n\n\treturn &multiStreamReader[T]{\n\t\tsts:        sts,\n\t\titemsCases: itemsCases,\n\t\tnonClosed:  nonClosed,\n\t}\n}\n\nfunc (msr *multiStreamReader[T]) recv() (T, error) {\n\tfor len(msr.nonClosed) > 0 {\n\t\tvar chosen int\n\t\tvar ok bool\n\t\tif len(msr.nonClosed) > maxSelectNum {\n\t\t\tvar recv reflect.Value\n\t\t\tchosen, recv, ok = reflect.Select(msr.itemsCases)\n\t\t\tif ok {\n\t\t\t\titem := recv.Interface().(streamItem[T])\n\t\t\t\treturn item.chunk, item.err\n\t\t\t}\n\t\t\tmsr.itemsCases[chosen].Chan = reflect.Value{}\n\t\t} else {\n\t\t\tvar item *streamItem[T]\n\t\t\tchosen, item, ok = receiveN(msr.nonClosed, msr.sts)\n\t\t\tif ok {\n\t\t\t\treturn item.chunk, item.err\n\t\t\t}\n\t\t}\n\n\t\t// delete the closed stream\n\t\tfor i := range msr.nonClosed {\n\t\t\tif msr.nonClosed[i] == chosen {\n\t\t\t\tmsr.nonClosed = append(msr.nonClosed[:i], msr.nonClosed[i+1:]...)\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\tif len(msr.sourceReaderNames) > 0 {\n\t\t\tvar t T\n\t\t\treturn t, &SourceEOF{msr.sourceReaderNames[chosen]}\n\t\t}\n\t}\n\n\tvar t T\n\treturn t, io.EOF\n}\n\nfunc (msr *multiStreamReader[T]) nonClosedStreams() []*stream[T] {\n\tret := make([]*stream[T], len(msr.nonClosed))\n\n\tfor i, idx := range msr.nonClosed {\n\t\tret[i] = msr.sts[idx]\n\t}\n\n\treturn ret\n}\n\nfunc (msr *multiStreamReader[T]) close() {\n\tfor _, s := range msr.sts {\n\t\ts.closeRecv()\n\t}\n}\n\nfunc (msr *multiStreamReader[T]) toStream() *stream[T] {\n\treturn toStream[T, *multiStreamReader[T]](msr)\n}\n\ntype streamReaderWithConvert[T any] struct {\n\tsr iStreamReader\n\n\tconvert func(any) (T, error)\n\n\terrWrapper func(error) error\n}\n\nfunc newStreamReaderWithConvert[T any](origin iStreamReader, convert func(any) (T, error), opts ...ConvertOption) *StreamReader[T] {\n\topt := &convertOptions{}\n\tfor _, o := range opts {\n\t\to(opt)\n\t}\n\n\tsrw := &streamReaderWithConvert[T]{\n\t\tsr:         origin,\n\t\tconvert:    convert,\n\t\terrWrapper: opt.ErrWrapper,\n\t}\n\n\treturn &StreamReader[T]{\n\t\ttyp: readerTypeWithConvert,\n\t\tsrw: srw,\n\t}\n}\n\ntype convertOptions struct {\n\tErrWrapper func(error) error\n}\n\ntype ConvertOption func(*convertOptions)\n\n// WithErrWrapper wraps the first error encountered in a stream reader during conversion by StreamReaderWithConvert.\n// The error returned by the convert function will not be wrapped.\n// If the returned err is nil or is ErrNoValue, the stream chunk will be ignored\nfunc WithErrWrapper(wrapper func(error) error) ConvertOption {\n\treturn func(o *convertOptions) {\n\t\to.ErrWrapper = wrapper\n\t}\n}\n\n// StreamReaderWithConvert returns a new StreamReader[D] that wraps sr and\n// applies convert to every element. The original reader sr must not be used\n// after calling this function.\n//\n// Filtering: if convert returns [ErrNoValue], the element is silently dropped\n// and the next element is read. This lets you strip empty or irrelevant chunks\n// without surfacing an error to the caller.\n//\n// Error wrapping: use [WithErrWrapper] to wrap non-convert errors (e.g. those\n// arriving from an upstream source) before they reach the caller.\n//\n//\tintReader := schema.StreamReaderFromArray([]int{0, 1, 2, 3})\n//\tstrReader := schema.StreamReaderWithConvert(intReader,\n//\t    func(i int) (string, error) {\n//\t        if i == 0 {\n//\t            return \"\", schema.ErrNoValue // skip zero\n//\t        }\n//\t        return fmt.Sprintf(\"val_%d\", i), nil\n//\t    })\n//\tdefer strReader.Close()\n//\t// Recv yields \"val_1\", \"val_2\", \"val_3\"\nfunc StreamReaderWithConvert[T, D any](sr *StreamReader[T], convert func(T) (D, error), opts ...ConvertOption) *StreamReader[D] {\n\tc := func(a any) (D, error) {\n\t\treturn convert(a.(T))\n\t}\n\n\treturn newStreamReaderWithConvert(sr, c, opts...)\n}\n\nfunc (srw *streamReaderWithConvert[T]) recv() (T, error) {\n\tfor {\n\t\tout, err := srw.sr.recvAny()\n\n\t\tif err != nil {\n\t\t\tvar t T\n\t\t\tif err == io.EOF {\n\t\t\t\treturn t, err\n\t\t\t}\n\t\t\tif srw.errWrapper != nil {\n\t\t\t\terr = srw.errWrapper(err)\n\t\t\t\tif err != nil && !errors.Is(err, ErrNoValue) {\n\t\t\t\t\treturn t, err\n\t\t\t\t}\n\t\t\t}\n\n\t\t\treturn t, err\n\t\t}\n\n\t\tt, err := srw.convert(out)\n\t\tif err == nil {\n\t\t\treturn t, nil\n\t\t}\n\n\t\tif !errors.Is(err, ErrNoValue) {\n\t\t\treturn t, err\n\t\t}\n\t}\n}\n\nfunc (srw *streamReaderWithConvert[T]) close() {\n\tsrw.sr.Close()\n}\n\ntype reader[T any] interface {\n\trecv() (T, error)\n\tclose()\n}\n\nfunc toStream[T any, Reader reader[T]](r Reader) *stream[T] {\n\tret := newStream[T](5)\n\n\tgo func() {\n\t\tdefer func() {\n\t\t\tpanicErr := recover()\n\t\t\tif panicErr != nil {\n\t\t\t\te := safe.NewPanicErr(panicErr, debug.Stack())\n\n\t\t\t\tvar chunk T\n\t\t\t\t_ = ret.send(chunk, e)\n\t\t\t}\n\n\t\t\tret.closeSend()\n\t\t\tr.close()\n\t\t}()\n\n\t\tfor {\n\t\t\tout, err := r.recv()\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tclosed := ret.send(out, err)\n\t\t\tif closed {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}()\n\n\treturn ret\n}\n\nfunc (srw *streamReaderWithConvert[T]) toStream() *stream[T] {\n\treturn toStream[T, *streamReaderWithConvert[T]](srw)\n}\n\ntype cpStreamElement[T any] struct {\n\tonce sync.Once\n\tnext *cpStreamElement[T]\n\titem streamItem[T]\n}\n\n// copyStreamReaders creates multiple independent StreamReaders from a single StreamReader.\n// Each child StreamReader can read from the original stream independently.\nfunc copyStreamReaders[T any](sr *StreamReader[T], n int) []*StreamReader[T] {\n\tcpsr := &parentStreamReader[T]{\n\t\tsr:            sr,\n\t\tsubStreamList: make([]*cpStreamElement[T], n),\n\t\tclosedNum:     0,\n\t}\n\n\t// Initialize subStreamList with an empty element, which acts like a tail node.\n\t// A nil element (used for dereference) represents that the child has been closed.\n\t// It is challenging to link the previous and current elements when the length of the original channel is unknown.\n\t// Additionally, using a previous pointer complicates dereferencing elements, possibly requiring reference counting.\n\telem := &cpStreamElement[T]{}\n\n\tfor i := range cpsr.subStreamList {\n\t\tcpsr.subStreamList[i] = elem\n\t}\n\n\tret := make([]*StreamReader[T], n)\n\tfor i := range ret {\n\t\tret[i] = &StreamReader[T]{\n\t\t\tcsr: &childStreamReader[T]{\n\t\t\t\tparent: cpsr,\n\t\t\t\tindex:  i,\n\t\t\t},\n\t\t\ttyp: readerTypeChild,\n\t\t}\n\t}\n\n\treturn ret\n}\n\ntype parentStreamReader[T any] struct {\n\t// sr is the original StreamReader.\n\tsr *StreamReader[T]\n\n\t// subStreamList maps each child's index to its latest read chunk.\n\t// Each value comes from a hidden linked list of cpStreamElement.\n\tsubStreamList []*cpStreamElement[T]\n\n\t// closedNum is the count of closed children.\n\tclosedNum uint32\n}\n\n// peek is not safe for concurrent use with the same idx but is safe for different idx.\n// Ensure that each child StreamReader uses a for-loop in a single goroutine.\nfunc (p *parentStreamReader[T]) peek(idx int) (t T, err error) {\n\telem := p.subStreamList[idx]\n\tif elem == nil {\n\t\t// Unexpected call to receive after the child has been closed.\n\t\treturn t, ErrRecvAfterClosed\n\t}\n\n\t// The sync.Once here is used to:\n\t// 1. Write the content of this cpStreamElement.\n\t// 2. Initialize the 'next' field of this cpStreamElement with an empty cpStreamElement,\n\t//    similar to the initialization in copyStreamReaders.\n\telem.once.Do(func() {\n\t\tt, err = p.sr.Recv()\n\t\telem.item = streamItem[T]{chunk: t, err: err}\n\t\tif err != io.EOF {\n\t\t\telem.next = &cpStreamElement[T]{}\n\t\t\tp.subStreamList[idx] = elem.next\n\t\t}\n\t})\n\n\t// The element has been set and will not be modified again.\n\t// Therefore, children can read this element's content and 'next' pointer concurrently.\n\tt = elem.item.chunk\n\terr = elem.item.err\n\tif err != io.EOF {\n\t\tp.subStreamList[idx] = elem.next\n\t}\n\n\treturn t, err\n}\n\nfunc (p *parentStreamReader[T]) close(idx int) {\n\tif p.subStreamList[idx] == nil {\n\t\treturn // avoid close multiple times\n\t}\n\n\tp.subStreamList[idx] = nil\n\n\tcurClosedNum := atomic.AddUint32(&p.closedNum, 1)\n\n\tallClosed := int(curClosedNum) == len(p.subStreamList)\n\tif allClosed {\n\t\tp.sr.Close()\n\t}\n}\n\ntype childStreamReader[T any] struct {\n\tparent *parentStreamReader[T]\n\tindex  int\n}\n\nfunc (csr *childStreamReader[T]) recv() (T, error) {\n\treturn csr.parent.peek(csr.index)\n}\n\nfunc (csr *childStreamReader[T]) toStream() *stream[T] {\n\treturn toStream[T, *childStreamReader[T]](csr)\n}\n\nfunc (csr *childStreamReader[T]) close() {\n\tcsr.parent.close(csr.index)\n}\n\n// MergeStreamReaders fans in multiple StreamReaders into a single StreamReader.\n// Elements from all source streams are interleaved in arrival order (non-deterministic).\n// The merged reader reaches EOF only after every source stream has been exhausted.\n//\n// Callers must still close the merged reader; it propagates the close signal\n// to all underlying sources.\n//\n// Use [MergeNamedStreamReaders] instead when you need to know which source\n// stream ended first (it emits a [SourceEOF] per-source EOF rather than\n// silently discarding them).\n//\n// Returns nil if srs is empty.\nfunc MergeStreamReaders[T any](srs []*StreamReader[T]) *StreamReader[T] {\n\tif len(srs) < 1 {\n\t\treturn nil\n\t}\n\n\tif len(srs) < 2 {\n\t\treturn srs[0]\n\t}\n\n\tvar arr []T\n\tvar ss []*stream[T]\n\n\tfor _, sr := range srs {\n\t\tswitch sr.typ {\n\t\tcase readerTypeStream:\n\t\t\tss = append(ss, sr.st)\n\t\tcase readerTypeArray:\n\t\t\tarr = append(arr, sr.ar.arr[sr.ar.index:]...)\n\t\tcase readerTypeMultiStream:\n\t\t\tss = append(ss, sr.msr.nonClosedStreams()...)\n\t\tcase readerTypeWithConvert:\n\t\t\tss = append(ss, sr.srw.toStream())\n\t\tcase readerTypeChild:\n\t\t\tss = append(ss, sr.csr.toStream())\n\t\tdefault:\n\t\t\tpanic(\"impossible\")\n\t\t}\n\t}\n\n\tif len(ss) == 0 {\n\t\treturn &StreamReader[T]{\n\t\t\ttyp: readerTypeArray,\n\t\t\tar: &arrayReader[T]{\n\t\t\t\tarr:   arr,\n\t\t\t\tindex: 0,\n\t\t\t},\n\t\t}\n\t}\n\n\tif len(arr) != 0 {\n\t\ts := arrToStream(arr)\n\t\tss = append(ss, s)\n\t}\n\n\treturn &StreamReader[T]{\n\t\ttyp: readerTypeMultiStream,\n\t\tmsr: newMultiStreamReader(ss),\n\t}\n}\n\n// MergeNamedStreamReaders merges multiple named StreamReaders into one.\n// Unlike [MergeStreamReaders], when a source stream reaches EOF the merged\n// reader emits a [SourceEOF] error (instead of silently continuing) so you can\n// detect exactly which source finished. Use [GetSourceName] to retrieve the\n// name from a SourceEOF error. The merged reader itself signals io.EOF only\n// after all named sources are exhausted.\n//\n// This is useful when downstream logic must react differently to each source\n// completing — for example, draining one agent's output before proceeding:\n//\n//\tnamedStreams := map[string]*schema.StreamReader[string]{\n//\t    \"agent_a\": srA,\n//\t    \"agent_b\": srB,\n//\t}\n//\tmerged := schema.MergeNamedStreamReaders(namedStreams)\n//\tdefer merged.Close()\n//\tfor {\n//\t    chunk, err := merged.Recv()\n//\t    if errors.Is(err, io.EOF) { break }\n//\t    if name, ok := schema.GetSourceName(err); ok {\n//\t        fmt.Printf(\"%s finished\\n\", name)\n//\t        continue\n//\t    }\n//\t    if err != nil { return err }\n//\t    process(chunk)\n//\t}\n//\n// Returns nil if srs is empty.\nfunc MergeNamedStreamReaders[T any](srs map[string]*StreamReader[T]) *StreamReader[T] {\n\tif len(srs) < 1 {\n\t\treturn nil\n\t}\n\n\tss := make([]*StreamReader[T], len(srs))\n\tnames := make([]string, len(srs))\n\n\ti := 0\n\tfor name, sr := range srs {\n\t\tss[i] = sr\n\t\tnames[i] = name\n\t\ti++\n\t}\n\n\treturn InternalMergeNamedStreamReaders(ss, names)\n}\n\n// InternalMergeNamedStreamReaders merges multiple readers with their names\n// into a single multi-stream reader.\nfunc InternalMergeNamedStreamReaders[T any](srs []*StreamReader[T], names []string) *StreamReader[T] {\n\tss := make([]*stream[T], len(srs))\n\n\tfor i, sr := range srs {\n\t\tss[i] = sr.toStream()\n\t}\n\n\tmsr := newMultiStreamReader(ss)\n\tmsr.sourceReaderNames = names\n\n\treturn &StreamReader[T]{\n\t\ttyp: readerTypeMultiStream,\n\t\tmsr: msr,\n\t}\n}\n"
  },
  {
    "path": "schema/stream_copy_external_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"fmt\"\n\t\"io\"\n\t\"runtime\"\n\t\"sort\"\n\t\"sync\"\n\t\"sync/atomic\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestStream1(t *testing.T) {\n\truntime.GOMAXPROCS(1)\n\n\tsr, sw := Pipe[int](0)\n\tgo func() {\n\t\tfor i := 0; i < 100; i++ {\n\t\t\tsw.Send(i, nil)\n\t\t\ttime.Sleep(3 * time.Millisecond)\n\t\t}\n\t\tsw.Close()\n\t}()\n\tcopied := sr.Copy(2)\n\tvar (\n\t\tnow   = time.Now().UnixMilli()\n\t\tts    = []int64{now, now}\n\t\ttsOld = []int64{now, now}\n\t)\n\tvar count int32\n\twg := sync.WaitGroup{}\n\twg.Add(2)\n\tgo func() {\n\t\ti := 0\n\t\ts := copied[0]\n\t\tfor {\n\t\t\tn, e := s.Recv()\n\t\t\tif e != nil {\n\t\t\t\tif e == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t\ttsOld[0] = ts[0]\n\t\t\tts[0] = time.Now().UnixMilli()\n\t\t\tinterval := ts[0] - tsOld[0]\n\t\t\tif interval >= 6 {\n\t\t\t\tatomic.AddInt32(&count, 1)\n\t\t\t}\n\t\t\tassert.Equal(t, i, n)\n\t\t\ti++\n\t\t}\n\t\twg.Done()\n\t}()\n\tgo func() {\n\t\ti := 0\n\t\ts := copied[1]\n\t\tfor {\n\t\t\tn, e := s.Recv()\n\t\t\tif e != nil {\n\t\t\t\tif e == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t\ttsOld[1] = ts[1]\n\t\t\tts[1] = time.Now().UnixMilli()\n\t\t\tinterval := ts[1] - tsOld[1]\n\t\t\tif interval >= 6 {\n\t\t\t\tatomic.AddInt32(&count, 1)\n\t\t\t}\n\t\t\tassert.Equal(t, i, n)\n\t\t\ti++\n\t\t}\n\t\twg.Done()\n\t}()\n\twg.Wait()\n\tt.Logf(\"count= %d\", count)\n}\n\ntype info struct {\n\tidx     int\n\tts      int64\n\tafter   int64\n\tcontent string\n}\n\nfunc TestCopyDelay(t *testing.T) {\n\truntime.GOMAXPROCS(10)\n\tn := 3\n\t//m := 100\n\ts := newStream[string](0)\n\tscp := s.asReader().Copy(n)\n\tgo func() {\n\t\ts.send(\"1\", nil)\n\t\ts.send(\"2\", nil)\n\t\ttime.Sleep(time.Second)\n\t\ts.send(\"3\", nil)\n\t\ts.closeSend()\n\t}()\n\twg := sync.WaitGroup{}\n\twg.Add(n)\n\tinfoList := make([][]info, n)\n\tfor i := 0; i < n; i++ {\n\t\tj := i\n\t\tgo func() {\n\t\t\tdefer func() {\n\t\t\t\tscp[j].Close()\n\t\t\t\twg.Done()\n\t\t\t}()\n\t\t\tfor {\n\t\t\t\tlastTime := time.Now()\n\t\t\t\tstr, err := scp[j].Recv()\n\t\t\t\tif err == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\tnow := time.Now()\n\t\t\t\tinfoList[j] = append(infoList[j], info{\n\t\t\t\t\tidx:     j,\n\t\t\t\t\tts:      now.UnixMicro(),\n\t\t\t\t\tafter:   now.Sub(lastTime).Milliseconds(),\n\t\t\t\t\tcontent: str,\n\t\t\t\t})\n\t\t\t}\n\t\t}()\n\t}\n\twg.Wait()\n\tinfos := make([]info, 0)\n\tfor _, infoL := range infoList {\n\t\tinfos = append(infos, infoL...)\n\t}\n\tsort.Slice(infos, func(i, j int) bool {\n\t\treturn infos[i].ts < infos[j].ts\n\t})\n\tfor _, info := range infos {\n\t\tfmt.Printf(\"child[%d] ts[%d] after[%5dms] content[%s]\\n\", info.idx, info.ts, info.after, info.content)\n\t}\n}\n"
  },
  {
    "path": "schema/stream_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"math/rand\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestStream(t *testing.T) {\n\ts := newStream[int](0)\n\n\tvar wg sync.WaitGroup\n\twg.Add(1)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tfor i := 0; i < 10; i++ {\n\t\t\tclosed := s.send(i, nil)\n\t\t\tif closed {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\ts.closeSend()\n\t}()\n\n\ti := 0\n\tfor {\n\t\ti++\n\t\tif i > 5 {\n\t\t\ts.closeRecv()\n\t\t\tbreak\n\t\t}\n\t\tv, err := s.recv()\n\t\tif err != nil {\n\t\t\tassert.ErrorIs(t, err, io.EOF)\n\t\t\tbreak\n\t\t}\n\t\tt.Log(v)\n\t}\n\n\twg.Wait()\n}\n\nfunc TestStreamCopy(t *testing.T) {\n\ts := newStream[string](10)\n\tsrs := s.asReader().Copy(2)\n\n\ts.send(\"a\", nil)\n\ts.send(\"b\", nil)\n\ts.send(\"c\", nil)\n\ts.closeSend()\n\n\tdefer func() {\n\t\tfor _, sr := range srs {\n\t\t\tsr.Close()\n\t\t}\n\t}()\n\n\tfor {\n\t\tv, err := srs[0].Recv()\n\t\tif errors.Is(err, io.EOF) {\n\t\t\tbreak\n\t\t}\n\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tt.Log(\"copy 01 recv\", v)\n\t}\n\n\tfor {\n\t\tv, err := srs[1].Recv()\n\t\tif errors.Is(err, io.EOF) {\n\t\t\tbreak\n\t\t}\n\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tt.Log(\"copy 02 recv\", v)\n\t}\n\n\tfor {\n\t\tv, err := s.recv()\n\t\tif errors.Is(err, io.EOF) {\n\t\t\tbreak\n\t\t}\n\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\n\t\tt.Log(\"recv origin\", v)\n\t}\n\n\tt.Log(\"done\")\n}\n\nfunc TestNewStreamCopy(t *testing.T) {\n\tt.Run(\"test one index recv channel blocked while other indexes could recv\", func(t *testing.T) {\n\t\ts := newStream[string](1)\n\t\tscp := s.asReader().Copy(2)\n\n\t\tvar t1, t2 time.Time\n\n\t\tgo func() {\n\t\t\ts.send(\"a\", nil)\n\t\t\tt1 = time.Now()\n\t\t\ttime.Sleep(time.Millisecond * 200)\n\t\t\ts.send(\"a\", nil)\n\t\t\ts.closeSend()\n\t\t}()\n\t\twg := sync.WaitGroup{}\n\t\twg.Add(2)\n\n\t\tgo func() {\n\t\t\tdefer func() {\n\t\t\t\tscp[0].Close()\n\t\t\t\twg.Done()\n\t\t\t}()\n\n\t\t\tfor {\n\t\t\t\tstr, err := scp[0].Recv()\n\t\t\t\tif err == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tassert.NoError(t, err)\n\t\t\t\tassert.Equal(t, str, \"a\")\n\t\t\t}\n\t\t}()\n\n\t\tgo func() {\n\t\t\tdefer func() {\n\t\t\t\tscp[1].Close()\n\t\t\t\twg.Done()\n\t\t\t}()\n\n\t\t\ttime.Sleep(time.Millisecond * 100)\n\t\t\tfor {\n\t\t\t\tstr, err := scp[1].Recv()\n\t\t\t\tif err == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tif t2.IsZero() {\n\t\t\t\t\tt2 = time.Now()\n\t\t\t\t}\n\n\t\t\t\tassert.NoError(t, err)\n\t\t\t\tassert.Equal(t, str, \"a\")\n\t\t\t}\n\t\t}()\n\n\t\twg.Wait()\n\n\t\tassert.True(t, t2.Sub(t1) < time.Millisecond*200)\n\t})\n\n\tt.Run(\"test one index recv channel blocked and other index closed\", func(t *testing.T) {\n\t\ts := newStream[string](1)\n\t\tscp := s.asReader().Copy(2)\n\n\t\tgo func() {\n\t\t\ts.send(\"a\", nil)\n\t\t\ttime.Sleep(time.Millisecond * 200)\n\t\t\ts.send(\"a\", nil)\n\t\t\ts.closeSend()\n\t\t}()\n\n\t\twg := sync.WaitGroup{}\n\t\twg.Add(2)\n\n\t\t//buf := scp[0].csr.parent.mem.buf\n\t\tgo func() {\n\t\t\tdefer func() {\n\t\t\t\tscp[0].Close()\n\t\t\t\twg.Done()\n\t\t\t}()\n\n\t\t\tfor {\n\t\t\t\tstr, err := scp[0].Recv()\n\t\t\t\tif err == io.EOF {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tassert.NoError(t, err)\n\t\t\t\tassert.Equal(t, str, \"a\")\n\t\t\t}\n\t\t}()\n\n\t\tgo func() {\n\t\t\ttime.Sleep(time.Millisecond * 100)\n\t\t\tscp[1].Close()\n\t\t\tscp[1].Close() // try close multiple times\n\t\t\twg.Done()\n\t\t}()\n\n\t\twg.Wait()\n\n\t\t//assert.Equal(t, 0, buf.Len())\n\t})\n\n\tt.Run(\"test long time recv\", func(t *testing.T) {\n\t\ts := newStream[int](2)\n\t\tn := 1000\n\t\tgo func() {\n\t\t\tfor i := 0; i < n; i++ {\n\t\t\t\ts.send(i, nil)\n\t\t\t}\n\n\t\t\ts.closeSend()\n\t\t}()\n\n\t\tm := 100\n\t\twg := sync.WaitGroup{}\n\t\twg.Add(m)\n\t\tcopies := s.asReader().Copy(m)\n\t\tfor i := 0; i < m; i++ {\n\t\t\tidx := i\n\t\t\tgo func() {\n\t\t\t\tcp := copies[idx]\n\t\t\t\tl := 0\n\t\t\t\tdefer func() {\n\t\t\t\t\tassert.Equal(t, 1000, l)\n\t\t\t\t\tcp.Close()\n\t\t\t\t\twg.Done()\n\t\t\t\t}()\n\n\t\t\t\tfor {\n\t\t\t\t\texp, err := cp.Recv()\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\n\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\tassert.Equal(t, exp, l)\n\t\t\t\t\tl++\n\t\t\t\t}\n\t\t\t}()\n\t\t}\n\n\t\twg.Wait()\n\t\t//memo := copies[0].csr.parent.mem\n\t\t//assert.Equal(t, true, memo.hasFinished)\n\t\t//assert.Equal(t, 0, memo.buf.Len())\n\t})\n\n\tt.Run(\"test closes\", func(t *testing.T) {\n\t\ts := newStream[int](20)\n\t\tn := 1000\n\t\tgo func() {\n\t\t\tfor i := 0; i < n; i++ {\n\t\t\t\ts.send(i, nil)\n\t\t\t}\n\n\t\t\ts.closeSend()\n\t\t}()\n\n\t\tm := 100\n\t\twg := sync.WaitGroup{}\n\t\twg.Add(m)\n\n\t\twgEven := sync.WaitGroup{}\n\t\twgEven.Add(m / 2)\n\n\t\tsr := s.asReader()\n\t\tsr.SetAutomaticClose()\n\t\tcopies := sr.Copy(m)\n\t\tfor i := 0; i < m; i++ {\n\t\t\tidx := i\n\t\t\tgo func() {\n\t\t\t\tcp := copies[idx]\n\t\t\t\tl := 0\n\t\t\t\tdefer func() {\n\t\t\t\t\tcp.Close()\n\t\t\t\t\twg.Done()\n\t\t\t\t\tif idx%2 == 0 {\n\t\t\t\t\t\twgEven.Done()\n\t\t\t\t\t}\n\t\t\t\t}()\n\n\t\t\t\tfor {\n\t\t\t\t\tif idx%2 == 0 && l == idx {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\n\t\t\t\t\texp, err := cp.Recv()\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\n\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\tassert.Equal(t, exp, l)\n\t\t\t\t\tl++\n\t\t\t\t}\n\t\t\t}()\n\t\t}\n\n\t\twgEven.Wait()\n\t\twg.Wait()\n\t\tassert.Equal(t, m, int(copies[0].csr.parent.closedNum))\n\t})\n\n\tt.Run(\"test reader do no close\", func(t *testing.T) {\n\t\ts := newStream[int](20)\n\t\tn := 1000\n\t\tgo func() {\n\t\t\tfor i := 0; i < n; i++ {\n\t\t\t\ts.send(i, nil)\n\t\t\t}\n\n\t\t\ts.closeSend()\n\t\t}()\n\n\t\tm := 4\n\t\twg := sync.WaitGroup{}\n\t\twg.Add(m)\n\n\t\tcopies := s.asReader().Copy(m)\n\t\tfor i := 0; i < m; i++ {\n\t\t\tidx := i\n\t\t\tcp := copies[idx]\n\t\t\tcp.SetAutomaticClose()\n\t\t\tgo func() {\n\t\t\t\tl := 0\n\t\t\t\tdefer func() {\n\t\t\t\t\twg.Done()\n\t\t\t\t}()\n\n\t\t\t\tfor {\n\t\t\t\t\texp, err := cp.Recv()\n\t\t\t\t\tif err == io.EOF {\n\t\t\t\t\t\tbreak\n\t\t\t\t\t}\n\n\t\t\t\t\tassert.NoError(t, err)\n\t\t\t\t\tassert.Equal(t, exp, l)\n\t\t\t\t\tl++\n\t\t\t\t}\n\t\t\t}()\n\t\t}\n\n\t\twg.Wait()\n\t\tassert.Equal(t, 0, int(copies[0].csr.parent.closedNum)) // not closed\n\t})\n\n}\n\nfunc checkStream(s *StreamReader[int]) error {\n\tdefer s.Close()\n\n\tfor i := 0; i < 10; i++ {\n\t\tchunk, err := s.Recv()\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tif chunk != i {\n\t\t\treturn fmt.Errorf(\"receive err, expected:%d, actual: %d\", i, chunk)\n\t\t}\n\t}\n\t_, err := s.Recv()\n\tif err != io.EOF {\n\t\treturn fmt.Errorf(\"close chan fail\")\n\t}\n\treturn nil\n}\n\nfunc testStreamN(cap, n int) error {\n\ts := newStream[int](cap)\n\tgo func() {\n\t\tfor i := 0; i < 10; i++ {\n\t\t\ts.send(i, nil)\n\t\t}\n\t\ts.closeSend()\n\t}()\n\n\tvs := s.asReader().Copy(n)\n\terr := checkStream(vs[0])\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tvs = vs[1].Copy(n)\n\terr = checkStream(vs[0])\n\tif err != nil {\n\t\treturn err\n\t}\n\tvs = vs[1].Copy(n)\n\terr = checkStream(vs[0])\n\tif err != nil {\n\t\treturn err\n\t}\n\treturn nil\n}\n\nfunc TestCopy(t *testing.T) {\n\tfor i := 0; i < 10; i++ {\n\t\tfor j := 2; j < 10; j++ {\n\t\t\terr := testStreamN(i, j)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t}\n\t}\n}\n\nfunc TestCopy5(t *testing.T) {\n\ts := newStream[int](0)\n\tgo func() {\n\t\tfor i := 0; i < 10; i++ {\n\t\t\tclosed := s.send(i, nil)\n\t\t\tif closed {\n\t\t\t\tfmt.Printf(\"has closed\")\n\t\t\t}\n\t\t}\n\t\ts.closeSend()\n\t}()\n\tvs := s.asReader().Copy(5)\n\ttime.Sleep(time.Second)\n\tdefer func() {\n\t\tfor _, v := range vs {\n\t\t\tv.Close()\n\t\t}\n\t}()\n\tfor i := 0; i < 10; i++ {\n\t\tchunk, err := vs[0].Recv()\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tif chunk != i {\n\t\t\tt.Fatalf(\"receive err, expected:%d, actual: %d\", i, chunk)\n\t\t}\n\t}\n\t_, err := vs[0].Recv()\n\tif err != io.EOF {\n\t\tt.Fatalf(\"copied stream reader cannot return EOF\")\n\t}\n\t_, err = vs[0].Recv()\n\tif err != io.EOF {\n\t\tt.Fatalf(\"copied stream reader cannot return EOF repeatedly\")\n\t}\n}\n\nfunc TestStreamReaderWithConvert(t *testing.T) {\n\ts := newStream[int](2)\n\n\tvar cntA int\n\tvar e error\n\n\tconvA := func(src int) (int, error) {\n\t\tif src == 1 {\n\t\t\treturn 0, fmt.Errorf(\"mock err\")\n\t\t}\n\n\t\treturn src, nil\n\t}\n\n\tsta := StreamReaderWithConvert[int, int](s.asReader(), convA)\n\tsta.SetAutomaticClose()\n\n\ts.send(1, nil)\n\ts.send(2, nil)\n\ts.closeSend()\n\n\tfor {\n\t\titem, err := sta.Recv()\n\t\tif err != nil {\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\te = err\n\t\t\tcontinue\n\t\t}\n\n\t\tcntA += item\n\t}\n\n\tassert.NotNil(t, e)\n\tassert.Equal(t, cntA, 2)\n}\n\nfunc TestArrayStreamCombined(t *testing.T) {\n\tasr := &StreamReader[int]{\n\t\ttyp: readerTypeArray,\n\t\tar: &arrayReader[int]{\n\t\t\tarr:   []int{0, 1, 2},\n\t\t\tindex: 0,\n\t\t},\n\t}\n\n\ts := newStream[int](3)\n\tfor i := 3; i < 6; i++ {\n\t\ts.send(i, nil)\n\t}\n\ts.closeSend()\n\n\tnSR := MergeStreamReaders([]*StreamReader[int]{asr, s.asReader()})\n\tnSR.SetAutomaticClose()\n\n\trecord := make([]bool, 6)\n\tfor i := 0; i < 6; i++ {\n\t\tchunk, err := nSR.Recv()\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tif record[chunk] {\n\t\t\tt.Fatal(\"record duplicated\")\n\t\t}\n\t\trecord[chunk] = true\n\t}\n\n\t_, err := nSR.Recv()\n\tif err != io.EOF {\n\t\tt.Fatal(\"reader haven't finish correctly\")\n\t}\n\n\tfor i := range record {\n\t\tif !record[i] {\n\t\t\tt.Fatal(\"record missing\")\n\t\t}\n\t}\n}\n\nfunc TestMultiStream(t *testing.T) {\n\tvar sts []*stream[int]\n\tsum := 0\n\tfor i := 0; i < 10; i++ {\n\t\tsize := rand.Intn(10) + 1\n\t\tsum += size\n\t\tst := newStream[int](size)\n\t\tfor j := 1; j <= size; j++ {\n\t\t\tst.send(j&0xffff+i<<16, nil)\n\t\t}\n\t\tst.closeSend()\n\t\tsts = append(sts, st)\n\t}\n\tmst := newMultiStreamReader(sts)\n\treceiveList := make([]int, 10)\n\tfor i := 0; i < sum; i++ {\n\t\tchunk, err := mst.recv()\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t\tif receiveList[chunk>>16] >= chunk&0xffff {\n\t\t\tt.Fatal(\"out of order\")\n\t\t}\n\t\treceiveList[chunk>>16] = chunk & 0xffff\n\t}\n\t_, err := mst.recv()\n\tif err != io.EOF {\n\t\tt.Fatal(\"end stream haven't return EOF\")\n\t}\n}\n\n// TestMergeNamedStreamReaders tests the functionality of MergeNamedStreamReaders\n// with a focus on SourceEOF error handling.\nfunc TestMergeNamedStreamReaders(t *testing.T) {\n\tt.Run(\"BasicSourceEOF\", func(t *testing.T) {\n\t\t// Create two named streams\n\t\tsr1, sw1 := Pipe[string](2)\n\t\tsr2, sw2 := Pipe[string](2)\n\n\t\t// Merge the streams with names\n\t\tnamedStreams := map[string]*StreamReader[string]{\n\t\t\t\"stream1\": sr1,\n\t\t\t\"stream2\": sr2,\n\t\t}\n\t\tmergedSR := MergeNamedStreamReaders(namedStreams)\n\t\tmergedSR.SetAutomaticClose()\n\n\t\t// Send data to the first stream and close it immediately\n\t\tgo func() {\n\t\t\tdefer sw1.Close()\n\t\t\tsw1.Send(\"data1-1\", nil)\n\t\t\tsw1.Send(\"data1-2\", nil)\n\t\t\t// First stream ends\n\t\t}()\n\n\t\t// Send data to the second stream with a delay before closing\n\t\tgo func() {\n\t\t\tdefer sw2.Close()\n\t\t\tsw2.Send(\"data2-1\", nil)\n\t\t\tsw2.Send(\"data2-2\", nil)\n\t\t\tsw2.Send(\"data2-3\", nil)\n\t\t\t// Second stream ends\n\t\t}()\n\n\t\t// Track received data and EOF sources\n\t\treceivedData := make(map[string][]string)\n\t\teofSources := make([]string, 0, 2)\n\n\t\tfor {\n\t\t\tchunk, err := mergedSR.Recv()\n\t\t\tif err != nil {\n\t\t\t\t// Check if it's a SourceEOF error\n\t\t\t\tif sourceName, ok := GetSourceName(err); ok {\n\t\t\t\t\teofSources = append(eofSources, sourceName)\n\t\t\t\t\tt.Logf(\"Received EOF from source: %s\", sourceName)\n\t\t\t\t\tcontinue // Continue receiving from other streams\n\t\t\t\t}\n\n\t\t\t\t// If it's a regular EOF, all streams have ended\n\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\t// Handle other errors\n\t\t\t\tt.Errorf(\"Error receiving data: %v\", err)\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\t// Categorize data by prefix\n\t\t\tif len(chunk) >= 5 {\n\t\t\t\tprefix := chunk[:5]\n\t\t\t\tif prefix == \"data1\" {\n\t\t\t\t\treceivedData[\"stream1\"] = append(receivedData[\"stream1\"], chunk)\n\t\t\t\t} else if prefix == \"data2\" {\n\t\t\t\t\treceivedData[\"stream2\"] = append(receivedData[\"stream2\"], chunk)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\t// Verify we received both SourceEOF errors\n\t\tif len(eofSources) != 2 {\n\t\t\tt.Errorf(\"Expected 2 SourceEOF errors, got %d\", len(eofSources))\n\t\t}\n\n\t\t// Verify the source names are correct\n\t\texpectedSources := map[string]bool{\"stream1\": false, \"stream2\": false}\n\t\tfor _, source := range eofSources {\n\t\t\tif _, exists := expectedSources[source]; !exists {\n\t\t\t\tt.Errorf(\"Unexpected source name: %s\", source)\n\t\t\t} else {\n\t\t\t\texpectedSources[source] = true\n\t\t\t}\n\t\t}\n\n\t\t// Verify all expected sources were seen\n\t\tfor source, seen := range expectedSources {\n\t\t\tif !seen {\n\t\t\t\tt.Errorf(\"Did not receive SourceEOF for %s\", source)\n\t\t\t}\n\t\t}\n\n\t\t// Verify we received all expected data\n\t\tif len(receivedData[\"stream1\"]) != 2 {\n\t\t\tt.Errorf(\"Expected 2 items from stream1, got %d\", len(receivedData[\"stream1\"]))\n\t\t}\n\n\t\tif len(receivedData[\"stream2\"]) != 3 {\n\t\t\tt.Errorf(\"Expected 3 items from stream2, got %d\", len(receivedData[\"stream2\"]))\n\t\t}\n\t})\n\n\tt.Run(\"EmptyStream\", func(t *testing.T) {\n\t\t// Create two streams, one will be empty\n\t\tsr1, sw1 := Pipe[string](2)\n\t\tsr2, sw2 := Pipe[string](2)\n\n\t\t// Close the first stream immediately to make it empty\n\t\tsw1.Close()\n\n\t\t// Merge the streams with names\n\t\tnamedStreams := map[string]*StreamReader[string]{\n\t\t\t\"empty\": sr1,\n\t\t\t\"data\":  sr2,\n\t\t}\n\t\tmergedSR := MergeNamedStreamReaders(namedStreams)\n\t\tmergedSR.SetAutomaticClose()\n\n\t\t// Send data to the second stream\n\t\tgo func() {\n\t\t\tdefer sw2.Close()\n\t\t\tsw2.Send(\"test-data\", nil)\n\t\t}()\n\n\t\t// Track received EOFs and data\n\t\teofSources := make(map[string]bool, 2)\n\t\treceivedData := make([]string, 0, 1)\n\n\t\tfor {\n\t\t\tchunk, err := mergedSR.Recv()\n\t\t\tif err != nil {\n\t\t\t\tif sourceName, ok := GetSourceName(err); ok {\n\t\t\t\t\teofSources[sourceName] = true\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\n\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tt.Errorf(\"Error receiving data: %v\", err)\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\treceivedData = append(receivedData, chunk)\n\t\t}\n\n\t\t// Verify we received EOF from the empty stream\n\t\tif len(eofSources) != 2 {\n\t\t\tt.Errorf(\"Expected 2 SourceEOF errors, got %d\", len(eofSources))\n\t\t}\n\n\t\tif _, exist := eofSources[\"empty\"]; !exist {\n\t\t\tt.Errorf(\"Expected EOF from 'empty' stream, got '%v'\", eofSources)\n\t\t}\n\t\tif _, exist := eofSources[\"data\"]; !exist {\n\t\t\tt.Errorf(\"Expected EOF from 'data' stream, got '%v'\", eofSources)\n\t\t}\n\n\t\t// Verify we received the data from the non-empty stream\n\t\tif len(receivedData) != 1 || receivedData[0] != \"test-data\" {\n\t\t\tt.Errorf(\"Expected to receive 'test-data', got %v\", receivedData)\n\t\t}\n\t})\n\n\tt.Run(\"ArraySource\", func(t *testing.T) {\n\t\t// Create three named streams\n\t\tsr1, sw1 := Pipe[string](2)\n\t\tsr2, sw2 := Pipe[string](2)\n\t\tsr3 := StreamReaderFromArray([]string{\"data3-1\", \"data3-2\", \"data3-3\"})\n\n\t\t// Merge the streams with names\n\t\tnamedStreams := map[string]*StreamReader[string]{\n\t\t\t\"stream1\": sr1,\n\t\t\t\"stream2\": sr2,\n\t\t\t\"stream3\": sr3,\n\t\t}\n\t\tmergedSR := MergeNamedStreamReaders(namedStreams)\n\t\tmergedSR.SetAutomaticClose()\n\n\t\t// Send data and close streams in sequence\n\t\tgo func() {\n\t\t\t// First stream sends one item then closes\n\t\t\tsw1.Send(\"data1\", nil)\n\t\t\tsw1.Close()\n\n\t\t\t// Second stream sends two items then closes\n\t\t\tsw2.Send(\"data2-1\", nil)\n\t\t\tsw2.Send(\"data2-2\", nil)\n\t\t\tsw2.Close()\n\t\t}()\n\n\t\t// Track EOF order and data count\n\t\teofOrder := make([]string, 0, 3)\n\t\tdataCount := 0\n\n\t\tfor {\n\t\t\t_, err := mergedSR.Recv()\n\t\t\tif err != nil {\n\t\t\t\tif sourceName, ok := GetSourceName(err); ok {\n\t\t\t\t\teofOrder = append(eofOrder, sourceName)\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\n\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\tt.Errorf(\"Error receiving data: %v\", err)\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\tdataCount++\n\t\t}\n\n\t\t// Verify EOF count\n\t\tif len(eofOrder) != 3 {\n\t\t\tt.Errorf(\"Expected 3 SourceEOF errors, got %d\", len(eofOrder))\n\t\t}\n\n\t\t// Verify data count\n\t\tif dataCount != 6 {\n\t\t\tt.Errorf(\"Expected 6 data items, got %d\", dataCount)\n\t\t}\n\t})\n\n\tt.Run(\"ErrorPropagation\", func(t *testing.T) {\n\t\t// Create two streams\n\t\tsr1, sw1 := Pipe[string](2)\n\t\tsr2, sw2 := Pipe[string](2)\n\n\t\t// Merge the streams with names\n\t\tnamedStreams := map[string]*StreamReader[string]{\n\t\t\t\"normal\": sr1,\n\t\t\t\"error\":  sr2,\n\t\t}\n\t\tmergedSR := MergeNamedStreamReaders(namedStreams)\n\t\tdefer mergedSR.Close()\n\n\t\ttestError := errors.New(\"test error\")\n\n\t\t// Send normal data to first stream\n\t\tgo func() {\n\t\t\tdefer sw1.Close()\n\t\t\tsw1.Send(\"normal-data\", nil)\n\t\t}()\n\n\t\t// Send error to second stream\n\t\tgo func() {\n\t\t\tdefer sw2.Close()\n\t\t\tsw2.Send(\"\", testError)\n\t\t}()\n\n\t\t// Track received errors\n\t\tvar receivedError error\n\n\t\tfor {\n\t\t\t_, err := mergedSR.Recv()\n\t\t\tif err != nil {\n\t\t\t\t// Skip SourceEOF errors\n\t\t\t\tif _, ok := GetSourceName(err); ok {\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\n\t\t\t\tif errors.Is(err, io.EOF) {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\n\t\t\t\t// Store the first non-EOF error\n\t\t\t\treceivedError = err\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\n\t\t// Verify we received the test error\n\t\tif receivedError == nil || receivedError.Error() != testError.Error() {\n\t\t\tt.Errorf(\"Expected error '%v', got '%v'\", testError, receivedError)\n\t\t}\n\t})\n}\n"
  },
  {
    "path": "schema/tool.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"sort\"\n\n\t\"github.com/eino-contrib/jsonschema\"\n\torderedmap \"github.com/wk8/go-ordered-map/v2\"\n)\n\n// DataType is the type of the parameter.\n// It must be one of the following values: \"object\", \"number\", \"integer\", \"string\", \"array\", \"null\", \"boolean\", which is the same as the type of the parameter in JSONSchema.\ntype DataType string\n\n// Supported JSONSchema data types for tool parameters.\nconst (\n\tObject  DataType = \"object\"\n\tNumber  DataType = \"number\"\n\tInteger DataType = \"integer\"\n\tString  DataType = \"string\"\n\tArray   DataType = \"array\"\n\tNull    DataType = \"null\"\n\tBoolean DataType = \"boolean\"\n)\n\n// ToolChoice controls how the model uses the tools provided to it.\n// Pass as part of the model option via [model.WithToolChoice].\ntype ToolChoice string\n\nconst (\n\t// ToolChoiceForbidden instructs the model not to call any tools, even if\n\t// tools are bound. The model responds with a plain text message instead.\n\t// Corresponds to \"none\" in OpenAI Chat Completion.\n\tToolChoiceForbidden ToolChoice = \"forbidden\"\n\n\t// ToolChoiceAllowed lets the model decide: it may generate a plain message\n\t// or call one or more tools. This is the default when tools are provided.\n\t// Corresponds to \"auto\" in OpenAI Chat Completion.\n\tToolChoiceAllowed ToolChoice = \"allowed\"\n\n\t// ToolChoiceForced requires the model to call at least one tool. Use this\n\t// when you want to guarantee structured output via tool calling.\n\t// Corresponds to \"required\" in OpenAI Chat Completion.\n\tToolChoiceForced ToolChoice = \"forced\"\n)\n\n// ToolInfo describes a tool that can be passed to a ChatModel via\n// [ToolCallingChatModel.WithTools] or [ChatModel.BindTools].\n//\n// Name should be concise and unique within the tool set. Desc should explain\n// when and why to use the tool; few-shot examples in Desc significantly improve\n// model accuracy. ParamsOneOf may be nil for tools that take no arguments.\ntype ToolInfo struct {\n\t// The unique name of the tool that clearly communicates its purpose.\n\tName string\n\t// Used to tell the model how/when/why to use the tool.\n\t// You can provide few-shot examples as a part of the description.\n\tDesc string\n\t// Extra is the extra information for the tool.\n\tExtra map[string]any\n\n\t// The parameters the functions accepts (different models may require different parameter types).\n\t// can be described in two ways:\n\t//  - use params: schema.NewParamsOneOfByParams(params)\n\t//  - use jsonschema: schema.NewParamsOneOfByJSONSchema(jsonschema)\n\t// If is nil, signals that the tool does not need any input parameter\n\t*ParamsOneOf\n}\n\n// ParameterInfo is the information of a parameter.\n// It is used to describe the parameters of a tool.\ntype ParameterInfo struct {\n\t// The type of the parameter.\n\tType DataType\n\t// The element type of the parameter, only for array.\n\tElemInfo *ParameterInfo\n\t// The sub parameters of the parameter, only for object.\n\tSubParams map[string]*ParameterInfo\n\t// The description of the parameter.\n\tDesc string\n\t// The enum values of the parameter, only for string.\n\tEnum []string\n\t// Whether the parameter is required.\n\tRequired bool\n}\n\n// ParamsOneOf holds a tool's parameter schema using exactly one of two\n// representations. Choose the one that best fits your needs:\n//\n//  1. [NewParamsOneOfByParams] — lightweight: describe parameters as a\n//     map[string]*[ParameterInfo]. Covers the most common cases (scalars,\n//     arrays, nested objects, enums, required flags).\n//\n//  2. [NewParamsOneOfByJSONSchema] — powerful: supply a full\n//     *jsonschema.Schema (JSON Schema 2020-12). Required when you need\n//     features not expressible via ParameterInfo, such as anyOf, oneOf, or\n//     $defs references. [utils.InferTool] generates this form automatically\n//     from Go struct tags.\n//\n// You must use exactly one constructor — setting both fields is invalid.\n// If ParamsOneOf is nil, the tool takes no input parameters.\ntype ParamsOneOf struct {\n\t// use NewParamsOneOfByParams to set this field\n\tparams map[string]*ParameterInfo\n\n\tjsonschema *jsonschema.Schema\n}\n\n// NewParamsOneOfByParams creates a ParamsOneOf with map[string]*ParameterInfo.\nfunc NewParamsOneOfByParams(params map[string]*ParameterInfo) *ParamsOneOf {\n\treturn &ParamsOneOf{\n\t\tparams: params,\n\t}\n}\n\n// NewParamsOneOfByJSONSchema creates a ParamsOneOf with *jsonschema.Schema.\nfunc NewParamsOneOfByJSONSchema(s *jsonschema.Schema) *ParamsOneOf {\n\treturn &ParamsOneOf{\n\t\tjsonschema: s,\n\t}\n}\n\n// ToJSONSchema parses ParamsOneOf, converts the parameter description that user actually provides, into the format ready to be passed to Model.\nfunc (p *ParamsOneOf) ToJSONSchema() (*jsonschema.Schema, error) {\n\tif p == nil {\n\t\treturn nil, nil\n\t}\n\n\tif p.params != nil {\n\t\tsc := &jsonschema.Schema{\n\t\t\tProperties: orderedmap.New[string, *jsonschema.Schema](),\n\t\t\tType:       string(Object),\n\t\t\tRequired:   make([]string, 0, len(p.params)),\n\t\t}\n\n\t\tkeys := make([]string, 0, len(p.params))\n\t\tfor k := range p.params {\n\t\t\tkeys = append(keys, k)\n\t\t}\n\t\tsort.Strings(keys)\n\n\t\tfor _, k := range keys {\n\t\t\tv := p.params[k]\n\t\t\tsc.Properties.Set(k, paramInfoToJSONSchema(v))\n\t\t\tif v.Required {\n\t\t\t\tsc.Required = append(sc.Required, k)\n\t\t\t}\n\t\t}\n\n\t\treturn sc, nil\n\t}\n\n\treturn p.jsonschema, nil\n}\n\nfunc paramInfoToJSONSchema(paramInfo *ParameterInfo) *jsonschema.Schema {\n\tjs := &jsonschema.Schema{\n\t\tType:        string(paramInfo.Type),\n\t\tDescription: paramInfo.Desc,\n\t}\n\n\tif len(paramInfo.Enum) > 0 {\n\t\tjs.Enum = make([]any, len(paramInfo.Enum))\n\t\tfor i, enum := range paramInfo.Enum {\n\t\t\tjs.Enum[i] = enum\n\t\t}\n\t}\n\n\tif paramInfo.ElemInfo != nil {\n\t\tjs.Items = paramInfoToJSONSchema(paramInfo.ElemInfo)\n\t}\n\n\tif len(paramInfo.SubParams) > 0 {\n\t\trequired := make([]string, 0, len(paramInfo.SubParams))\n\t\tjs.Properties = orderedmap.New[string, *jsonschema.Schema]()\n\t\tkeys := make([]string, 0, len(paramInfo.SubParams))\n\t\tfor k := range paramInfo.SubParams {\n\t\t\tkeys = append(keys, k)\n\t\t}\n\t\tsort.Strings(keys)\n\n\t\tfor _, k := range keys {\n\t\t\tv := paramInfo.SubParams[k]\n\t\t\titem := paramInfoToJSONSchema(v)\n\t\t\tjs.Properties.Set(k, item)\n\t\t\tif v.Required {\n\t\t\t\trequired = append(required, k)\n\t\t\t}\n\t\t}\n\n\t\tjs.Required = required\n\t}\n\n\treturn js\n}\n"
  },
  {
    "path": "schema/tool_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage schema\n\nimport (\n\t\"encoding/json\"\n\t\"testing\"\n\n\t\"github.com/eino-contrib/jsonschema\"\n\t\"github.com/smartystreets/goconvey/convey\"\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestParamsOneOfToJSONSchema(t *testing.T) {\n\tconvey.Convey(\"ParamsOneOfToJSONSchema\", t, func() {\n\t\tvar (\n\t\t\toneOf     ParamsOneOf\n\t\t\tconverted any\n\t\t\terr       error\n\t\t)\n\n\t\tconvey.Convey(\"user provides JSON schema directly, use what the user provides\", func() {\n\t\t\toneOf.jsonschema = &jsonschema.Schema{\n\t\t\t\tType:        \"string\",\n\t\t\t\tDescription: \"this is the only argument\",\n\t\t\t}\n\t\t\tconverted, err = oneOf.ToJSONSchema()\n\t\t\tconvey.So(err, convey.ShouldBeNil)\n\t\t\tconvey.So(converted, convey.ShouldResemble, oneOf.jsonschema)\n\t\t})\n\n\t\tconvey.Convey(\"user provides map[string]ParameterInfo, converts to json schema\", func() {\n\t\t\toneOf.params = map[string]*ParameterInfo{\n\t\t\t\t\"arg1\": {\n\t\t\t\t\tType:     String,\n\t\t\t\t\tDesc:     \"this is the first argument\",\n\t\t\t\t\tRequired: true,\n\t\t\t\t\tEnum:     []string{\"1\", \"2\"},\n\t\t\t\t},\n\t\t\t\t\"arg2\": {\n\t\t\t\t\tType: Object,\n\t\t\t\t\tDesc: \"this is the second argument\",\n\t\t\t\t\tSubParams: map[string]*ParameterInfo{\n\t\t\t\t\t\t\"sub_arg1\": {\n\t\t\t\t\t\t\tType:     String,\n\t\t\t\t\t\t\tDesc:     \"this is the sub argument\",\n\t\t\t\t\t\t\tRequired: true,\n\t\t\t\t\t\t\tEnum:     []string{\"1\", \"2\"},\n\t\t\t\t\t\t},\n\t\t\t\t\t\t\"sub_arg2\": {\n\t\t\t\t\t\t\tType: String,\n\t\t\t\t\t\t\tDesc: \"this is the sub argument 2\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t\tRequired: true,\n\t\t\t\t},\n\t\t\t\t\"arg3\": {\n\t\t\t\t\tType: Array,\n\t\t\t\t\tDesc: \"this is the third argument\",\n\t\t\t\t\tElemInfo: &ParameterInfo{\n\t\t\t\t\t\tType:     String,\n\t\t\t\t\t\tDesc:     \"this is the element of the third argument\",\n\t\t\t\t\t\tRequired: true,\n\t\t\t\t\t\tEnum:     []string{\"1\", \"2\"},\n\t\t\t\t\t},\n\t\t\t\t\tRequired: true,\n\t\t\t\t},\n\t\t\t}\n\t\t\tconverted, err = oneOf.ToJSONSchema()\n\t\t\tconvey.So(err, convey.ShouldBeNil)\n\t\t})\n\n\t\tconvey.Convey(\"user provides map[string]ParameterInfo, converts to json schema in order\", func() {\n\t\t\tparams := &ParamsOneOf{\n\t\t\t\tparams: map[string]*ParameterInfo{\n\t\t\t\t\t\"c\": {\n\t\t\t\t\t\tType: \"string\",\n\t\t\t\t\t},\n\t\t\t\t\t\"a\": {\n\t\t\t\t\t\tType: \"object\",\n\t\t\t\t\t\tSubParams: map[string]*ParameterInfo{\n\t\t\t\t\t\t\t\"z\": {\n\t\t\t\t\t\t\t\tType: \"number\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\"y\": {\n\t\t\t\t\t\t\t\tType: \"string\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t\t\"b\": {\n\t\t\t\t\t\tType: \"array\",\n\t\t\t\t\t\tElemInfo: &ParameterInfo{\n\t\t\t\t\t\t\tType: \"object\",\n\t\t\t\t\t\t\tSubParams: map[string]*ParameterInfo{\n\t\t\t\t\t\t\t\t\"p\": {\n\t\t\t\t\t\t\t\t\tType: \"integer\",\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t\"o\": {\n\t\t\t\t\t\t\t\t\tType: \"boolean\",\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}\n\n\t\t\tschema1, err := params.ToJSONSchema()\n\t\t\tassert.NoError(t, err)\n\t\t\tjson1, err := json.Marshal(schema1)\n\t\t\tassert.NoError(t, err)\n\n\t\t\tschema2, err := params.ToJSONSchema()\n\t\t\tassert.NoError(t, err)\n\t\t\tjson2, err := json.Marshal(schema2)\n\t\t\tassert.NoError(t, err)\n\n\t\t\tassert.Equal(t, string(json1), string(json2))\n\t\t})\n\n\t})\n}\n"
  },
  {
    "path": "scripts/dev_setup.sh",
    "content": "#!/usr/bin/env bash\n# dev_setup.sh — set up a local multi-module workspace for eino development.\n#\n# BACKGROUND\n#   eino, eino-ext, and eino-examples live in separate GitHub repositories to\n#   keep their Go modules, versioning, and maintenance independent. However,\n#   working across them is inconvenient: editors and AI coding tools lack\n#   cross-repo type information and can't navigate between them.\n#\n#   This script brings all three repos together locally so that a single\n#   go.work file provides full cross-module LSP (go-to-definition, type\n#   inference, autocomplete) across all ~83 modules — without touching any\n#   remote repository.\n#\n# WHAT IT DOES\n#   1. Clones eino-ext  → ext/\n#   2. Clones eino-examples → examples/\n#   3. Registers ext/ and examples/ in .git/info/exclude so eino's git\n#      never sees them (local-only, never committed)\n#   4. Creates go.work at the repo root covering eino + all modules in\n#      ext/ and examples/ (go.work is already in .gitignore)\n#\n# RESULTING LAYOUT\n#   eino/               ← you are here (github.com/cloudwego/eino)\n#   eino/ext/           ← github.com/cloudwego/eino-ext  (full git repo)\n#   eino/examples/      ← github.com/cloudwego/eino-examples  (full git repo)\n#   eino/go.work        ← wires all modules together (gitignored)\n#\n# WORKING ACROSS REPOS\n#   Each subdirectory is a full independent git repo tracking its own remote.\n#   To contribute to eino-ext or eino-examples, work inside that directory:\n#\n#     cd ext\n#     git checkout -b feat/my-feature\n#     # make changes — editor has full cross-repo type info via go.work\n#     git commit -m \"feat: ...\"\n#     git push origin feat/my-feature   # pushes to cloudwego/eino-ext\n#\n# KEEPING REPOS UP TO DATE\n#     git -C ext pull\n#     git -C examples pull\n#\n# USAGE\n#   bash scripts/dev_setup.sh           # first-time setup\n#   bash scripts/dev_setup.sh --reset   # re-clone everything from scratch\n\nset -euo pipefail\n\nREPO_ROOT=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")/..\" && pwd)\"\ncd \"$REPO_ROOT\"\n\nEXT_DIR=\"ext\"\nEXAMPLES_DIR=\"examples\"\nEINO_EXT_REPO=\"https://github.com/cloudwego/eino-ext\"\nEINO_EXAMPLES_REPO=\"https://github.com/cloudwego/eino-examples\"\n\n# Parse flags\nRESET=false\nfor arg in \"$@\"; do\n  case $arg in\n    --reset) RESET=true ;;\n  esac\ndone\n\necho \"==> Setting up eino dev workspace in: $REPO_ROOT\"\n\n# --reset: remove existing dirs\nif [ \"$RESET\" = true ]; then\n  echo \"==> --reset: removing $EXT_DIR/ and $EXAMPLES_DIR/\"\n  rm -rf \"$EXT_DIR\" \"$EXAMPLES_DIR\" go.work go.work.sum\nfi\n\n# Clone repos if not already present\nif [ ! -d \"$EXT_DIR/.git\" ]; then\n  echo \"==> Cloning eino-ext into $EXT_DIR/\"\n  git clone \"$EINO_EXT_REPO\" \"$EXT_DIR\"\nelse\n  echo \"==> $EXT_DIR/ already exists, skipping clone\"\nfi\n\nif [ ! -d \"$EXAMPLES_DIR/.git\" ]; then\n  echo \"==> Cloning eino-examples into $EXAMPLES_DIR/\"\n  git clone \"$EINO_EXAMPLES_REPO\" \"$EXAMPLES_DIR\"\nelse\n  echo \"==> $EXAMPLES_DIR/ already exists, skipping clone\"\nfi\n\n# Exclude dirs from eino's git tracking (local only, not committed)\nEXCLUDE_FILE=\".git/info/exclude\"\nadd_exclude() {\n  local entry=\"$1\"\n  if ! grep -qxF \"$entry\" \"$EXCLUDE_FILE\" 2>/dev/null; then\n    echo \"$entry\" >> \"$EXCLUDE_FILE\"\n    echo \"==> Added '$entry' to $EXCLUDE_FILE\"\n  fi\n}\nadd_exclude \"$EXT_DIR/\"\nadd_exclude \"$EXAMPLES_DIR/\"\n\n# Build go.work covering eino root + every go.mod found in ext/ and examples/\nif [ ! -f \"go.work\" ]; then\n  echo \"==> Creating go.work\"\n  go work init .\n\n  # Collect all module directories (directories containing a go.mod)\n  while IFS= read -r modfile; do\n    dir=\"$(dirname \"$modfile\")\"\n    go work use \"$dir\"\n  done < <(find \"$EXT_DIR\" \"$EXAMPLES_DIR\" -name \"go.mod\" | sort)\n\n  echo \"==> go.work created with $(grep -c '^\\s\\+\\.' go.work || true) module(s)\"\nelse\n  echo \"==> go.work already exists, skipping (use --reset to recreate)\"\nfi\n\necho \"\"\necho \"Done. Your workspace includes:\"\necho \"  .             — github.com/cloudwego/eino\"\necho \"  $EXT_DIR/          — github.com/cloudwego/eino-ext ($(find \"$EXT_DIR\" -name \"go.mod\" | wc -l | tr -d ' ') modules)\"\necho \"  $EXAMPLES_DIR/  — github.com/cloudwego/eino-examples ($(find \"$EXAMPLES_DIR\" -name \"go.mod\" | wc -l | tr -d ' ') modules)\"\necho \"\"\necho \"Run 'go build ./...' or open this directory in your editor.\"\n"
  },
  {
    "path": "scripts/eino_setup.sh",
    "content": "#!/usr/bin/env bash\n# eino_setup.sh — fetch eino framework source into your project for AI-assisted development.\n#\n# BACKGROUND\n#   When building applications with eino, your AI coding assistant (Claude Code,\n#   Cursor, Copilot, etc.) only sees your code. It cannot navigate into eino's\n#   source to understand how components work, what patterns are idiomatic, or\n#   how to wire things together correctly.\n#\n#   This script clones eino, eino-ext, and eino-examples into a _eino/ directory\n#   inside your project. Your AI assistant can then browse the actual source,\n#   examples, and extensions — giving it full context to help you build correctly.\n#\n# WHAT IT DOES\n#   1. Clones eino        → _eino/eino/\n#   2. Clones eino-ext    → _eino/eino-ext/\n#   3. Clones eino-examples → _eino/eino-examples/\n#   4. Adds _eino/ to .gitignore (read-only reference, never committed)\n#   5. Writes a _eino/README.md explaining the directory to future readers\n#\n# RESULTING LAYOUT\n#   your-project/\n#   ├── _eino/\n#   │   ├── eino/          ← github.com/cloudwego/eino (core framework)\n#   │   ├── eino-ext/      ← github.com/cloudwego/eino-ext (components & integrations)\n#   │   └── eino-examples/ ← github.com/cloudwego/eino-examples (patterns & recipes)\n#   └── ... your code\n#\n# NOTE: _eino/ is read-only reference material. Do not edit files inside it.\n#       Your go.mod is unchanged — eino remains a normal dependency.\n#\n# KEEPING UP TO DATE\n#   bash eino_setup.sh --update    # pull latest on all three repos\n#\n# USAGE\n#   bash eino_setup.sh             # first-time setup\n#   bash eino_setup.sh --reset     # re-clone everything from scratch\n#   bash eino_setup.sh --update    # pull latest without re-cloning\n#\n# SYSTEM PROMPT\n#   After running this script, add the following to your AI assistant's project\n#   instructions (CLAUDE.md, .cursorrules, .github/copilot-instructions.md, etc.):\n#\n#   ---\n#   ## eino Framework Reference\n#\n#   This project uses the eino framework (github.com/cloudwego/eino).\n#   The full framework source is available locally in `_eino/`:\n#\n#   - `_eino/eino/`          — core framework (components, graph, compose, callbacks)\n#   - `_eino/eino-ext/`      — official components and integrations (models, tools, retrievers, etc.)\n#   - `_eino/eino-examples/` — working examples and patterns\n#\n#   When answering questions about eino APIs, component wiring, graph construction,\n#   callbacks, or any eino-specific patterns: explore `_eino/` first.\n#   Prefer examples from `_eino/eino-examples/` as the canonical reference for\n#   idiomatic usage.\n#   ---\n\nset -euo pipefail\n\nPROJECT_ROOT=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\ncd \"$PROJECT_ROOT\"\n\nEINO_DIR=\"_eino\"\nEINO_REPO=\"https://github.com/cloudwego/eino\"\nEINO_EXT_REPO=\"https://github.com/cloudwego/eino-ext\"\nEINO_EXAMPLES_REPO=\"https://github.com/cloudwego/eino-examples\"\n\n# Parse flags\nRESET=false\nUPDATE=false\nfor arg in \"$@\"; do\n  case $arg in\n    --reset)  RESET=true ;;\n    --update) UPDATE=true ;;\n  esac\ndone\n\necho \"==> eino setup in: $PROJECT_ROOT\"\n\n# --reset: remove and re-clone\nif [ \"$RESET\" = true ]; then\n  echo \"==> --reset: removing $EINO_DIR/\"\n  rm -rf \"$EINO_DIR\"\nfi\n\n# --update: pull latest on existing clones\nif [ \"$UPDATE\" = true ]; then\n  for repo in eino eino-ext eino-examples; do\n    dir=\"$EINO_DIR/$repo\"\n    if [ -d \"$dir/.git\" ]; then\n      echo \"==> Updating $dir/\"\n      git -C \"$dir\" pull --ff-only\n    else\n      echo \"==> $dir/ not found, skipping update (run without --update to clone)\"\n    fi\n  done\n  echo \"\"\n  echo \"Done. Run 'bash eino_setup.sh' to clone any missing repos.\"\n  exit 0\nfi\n\nmkdir -p \"$EINO_DIR\"\n\n# Clone repos (shallow — we only need source to read, not full history)\nclone_if_missing() {\n  local repo_url=\"$1\"\n  local dest=\"$2\"\n  if [ ! -d \"$dest/.git\" ]; then\n    echo \"==> Cloning $(basename \"$dest\")/\"\n    git clone --depth=1 \"$repo_url\" \"$dest\"\n  else\n    echo \"==> $dest/ already exists, skipping clone\"\n  fi\n}\n\nclone_if_missing \"$EINO_REPO\"          \"$EINO_DIR/eino\"\nclone_if_missing \"$EINO_EXT_REPO\"      \"$EINO_DIR/eino-ext\"\nclone_if_missing \"$EINO_EXAMPLES_REPO\" \"$EINO_DIR/eino-examples\"\n\n# Add _eino/ to .gitignore\nGITIGNORE=\".gitignore\"\nif ! grep -qxF \"$EINO_DIR/\" \"$GITIGNORE\" 2>/dev/null; then\n  echo \"\" >> \"$GITIGNORE\"\n  echo \"# eino framework source (AI coding reference — see eino_setup.sh)\" >> \"$GITIGNORE\"\n  echo \"$EINO_DIR/\" >> \"$GITIGNORE\"\n  echo \"==> Added '$EINO_DIR/' to $GITIGNORE\"\nfi\n\n# Write a README so the directory is self-explanatory\ncat > \"$EINO_DIR/README.md\" <<'EOF'\n# _eino — eino framework source reference\n\nThis directory contains read-only clones of the eino framework repositories,\nchecked out for use by AI coding assistants (Claude Code, Cursor, Copilot, etc.).\n\n| Directory      | Repository                              | Purpose                        |\n|----------------|-----------------------------------------|--------------------------------|\n| `eino/`        | github.com/cloudwego/eino               | Core framework source          |\n| `eino-ext/`    | github.com/cloudwego/eino-ext           | Components and integrations    |\n| `eino-examples/` | github.com/cloudwego/eino-examples    | Patterns, recipes, and samples |\n\n**Do not edit files here.** This directory is in `.gitignore` and is never committed.\n\nTo update to the latest:\n\n    bash eino_setup.sh --update\n\nTo re-clone from scratch:\n\n    bash eino_setup.sh --reset\nEOF\n\necho \"\"\necho \"Done. Your AI assistant now has full eino context in $EINO_DIR/:\"\necho \"  $EINO_DIR/eino/          — core framework ($(find \"$EINO_DIR/eino\" -name \"*.go\" | wc -l | tr -d ' ') .go files)\"\necho \"  $EINO_DIR/eino-ext/      — components & integrations ($(find \"$EINO_DIR/eino-ext\" -name \"*.go\" | wc -l | tr -d ' ') .go files)\"\necho \"  $EINO_DIR/eino-examples/ — patterns & recipes ($(find \"$EINO_DIR/eino-examples\" -name \"*.go\" | wc -l | tr -d ' ') .go files)\"\necho \"\"\necho \"Add the following to your AI assistant's system prompt or project instructions\"\necho \"(e.g. CLAUDE.md, .cursorrules, .github/copilot-instructions.md):\"\necho \"\"\necho \"---\"\ncat <<'PROMPT'\n## eino Framework Reference\n\nThis project uses the eino framework (github.com/cloudwego/eino).\nThe full framework source is available locally in `_eino/`:\n\n- `_eino/eino/`          — core framework (components, graph, compose, callbacks)\n- `_eino/eino-ext/`      — official components and integrations (models, tools, retrievers, etc.)\n- `_eino/eino-examples/` — working examples and patterns\n\nWhen answering questions about eino APIs, component wiring, graph construction,\ncallbacks, or any eino-specific patterns: explore `_eino/` first.\nPrefer examples from `_eino/eino-examples/` as the canonical reference for\nidiomatic usage.\nPROMPT\necho \"---\"\n"
  },
  {
    "path": "utils/callbacks/template.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n// Package callbacks provides ready-to-use callback handler templates for components.\npackage callbacks\n\nimport (\n\t\"context\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\n// NewHandlerHelper creates a new component template handler builder.\n// This builder can be used to configure and build a component template handler,\n// which can handle callback events for different components with its own struct definition,\n// and fallbackTemplate can be used to handle scenarios where none of the cases are hit as a fallback.\nfunc NewHandlerHelper() *HandlerHelper {\n\treturn &HandlerHelper{\n\t\tcomposeTemplates: map[components.Component]callbacks.Handler{},\n\t}\n}\n\n// HandlerHelper is a builder for creating a callbacks.Handler with specific handlers for different component types.\n// create a handler with callbacks.NewHandlerHelper().\n// eg.\n//\n//\thelper := template.NewHandlerHelper().\n//\t\tChatModel(&model.IndexerCallbackHandler{}).\n//\t\tPrompt(&prompt.IndexerCallbackHandler{}).\n//\t\tHandler()\n//\n// then use the handler with runnable.Invoke(ctx, input, compose.WithCallbacks(handler))\ntype HandlerHelper struct {\n\tpromptHandler      *PromptCallbackHandler\n\tchatModelHandler   *ModelCallbackHandler\n\tembeddingHandler   *EmbeddingCallbackHandler\n\tindexerHandler     *IndexerCallbackHandler\n\tretrieverHandler   *RetrieverCallbackHandler\n\tloaderHandler      *LoaderCallbackHandler\n\ttransformerHandler *TransformerCallbackHandler\n\ttoolHandler        *ToolCallbackHandler\n\ttoolsNodeHandler   *ToolsNodeCallbackHandlers\n\tagentHandler       *AgentCallbackHandler\n\tcomposeTemplates   map[components.Component]callbacks.Handler\n}\n\n// Handler returns the callbacks.Handler created by HandlerHelper.\nfunc (c *HandlerHelper) Handler() callbacks.Handler {\n\treturn &handlerTemplate{c}\n}\n\n// Prompt sets the prompt handler for the handler helper, which will be called when the prompt component is executed.\nfunc (c *HandlerHelper) Prompt(handler *PromptCallbackHandler) *HandlerHelper {\n\tc.promptHandler = handler\n\treturn c\n}\n\n// ChatModel sets the chat model handler for the handler helper, which will be called when the chat model component is executed.\nfunc (c *HandlerHelper) ChatModel(handler *ModelCallbackHandler) *HandlerHelper {\n\tc.chatModelHandler = handler\n\treturn c\n}\n\n// Embedding sets the embedding handler for the handler helper, which will be called when the embedding component is executed.\nfunc (c *HandlerHelper) Embedding(handler *EmbeddingCallbackHandler) *HandlerHelper {\n\tc.embeddingHandler = handler\n\treturn c\n}\n\n// Indexer sets the indexer handler for the handler helper, which will be called when the indexer component is executed.\nfunc (c *HandlerHelper) Indexer(handler *IndexerCallbackHandler) *HandlerHelper {\n\tc.indexerHandler = handler\n\treturn c\n}\n\n// Retriever sets the retriever handler for the handler helper, which will be called when the retriever component is executed.\nfunc (c *HandlerHelper) Retriever(handler *RetrieverCallbackHandler) *HandlerHelper {\n\tc.retrieverHandler = handler\n\treturn c\n}\n\n// Loader sets the loader handler for the handler helper, which will be called when the loader component is executed.\nfunc (c *HandlerHelper) Loader(handler *LoaderCallbackHandler) *HandlerHelper {\n\tc.loaderHandler = handler\n\treturn c\n}\n\n// Transformer sets the transformer handler for the handler helper, which will be called when the transformer component is executed.\nfunc (c *HandlerHelper) Transformer(handler *TransformerCallbackHandler) *HandlerHelper {\n\tc.transformerHandler = handler\n\treturn c\n}\n\n// Tool sets the tool handler for the handler helper, which will be called when the tool component is executed.\nfunc (c *HandlerHelper) Tool(handler *ToolCallbackHandler) *HandlerHelper {\n\tc.toolHandler = handler\n\treturn c\n}\n\n// ToolsNode sets the tools node handler for the handler helper, which will be called when the tools node is executed.\nfunc (c *HandlerHelper) ToolsNode(handler *ToolsNodeCallbackHandlers) *HandlerHelper {\n\tc.toolsNodeHandler = handler\n\treturn c\n}\n\n// Agent sets the agent handler for the handler helper, which will be called when the agent is executed.\nfunc (c *HandlerHelper) Agent(handler *AgentCallbackHandler) *HandlerHelper {\n\tc.agentHandler = handler\n\treturn c\n}\n\n// Graph sets the graph handler for the handler helper, which will be called when the graph is executed.\nfunc (c *HandlerHelper) Graph(handler callbacks.Handler) *HandlerHelper {\n\tc.composeTemplates[compose.ComponentOfGraph] = handler\n\treturn c\n}\n\n// Chain sets the chain handler for the handler helper, which will be called when the chain is executed.\nfunc (c *HandlerHelper) Chain(handler callbacks.Handler) *HandlerHelper {\n\tc.composeTemplates[compose.ComponentOfChain] = handler\n\treturn c\n}\n\n// Lambda sets the lambda handler for the handler helper, which will be called when the lambda is executed.\nfunc (c *HandlerHelper) Lambda(handler callbacks.Handler) *HandlerHelper {\n\tc.composeTemplates[compose.ComponentOfLambda] = handler\n\treturn c\n}\n\ntype handlerTemplate struct {\n\t*HandlerHelper\n}\n\n// OnStart is the callback function for the start event of a component.\n// implement the callbacks Handler interface.\nfunc (c *handlerTemplate) OnStart(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\tswitch info.Component {\n\tcase components.ComponentOfPrompt:\n\t\treturn c.promptHandler.OnStart(ctx, info, prompt.ConvCallbackInput(input))\n\tcase components.ComponentOfChatModel:\n\t\treturn c.chatModelHandler.OnStart(ctx, info, model.ConvCallbackInput(input))\n\tcase components.ComponentOfEmbedding:\n\t\treturn c.embeddingHandler.OnStart(ctx, info, embedding.ConvCallbackInput(input))\n\tcase components.ComponentOfIndexer:\n\t\treturn c.indexerHandler.OnStart(ctx, info, indexer.ConvCallbackInput(input))\n\tcase components.ComponentOfRetriever:\n\t\treturn c.retrieverHandler.OnStart(ctx, info, retriever.ConvCallbackInput(input))\n\tcase components.ComponentOfLoader:\n\t\treturn c.loaderHandler.OnStart(ctx, info, document.ConvLoaderCallbackInput(input))\n\tcase components.ComponentOfTransformer:\n\t\treturn c.transformerHandler.OnStart(ctx, info, document.ConvTransformerCallbackInput(input))\n\tcase components.ComponentOfTool:\n\t\treturn c.toolHandler.OnStart(ctx, info, tool.ConvCallbackInput(input))\n\tcase compose.ComponentOfToolsNode:\n\t\treturn c.toolsNodeHandler.OnStart(ctx, info, convToolsNodeCallbackInput(input))\n\tcase adk.ComponentOfAgent:\n\t\treturn c.agentHandler.OnStart(ctx, info, adk.ConvAgentCallbackInput(input))\n\tcase compose.ComponentOfGraph,\n\t\tcompose.ComponentOfChain,\n\t\tcompose.ComponentOfLambda:\n\t\treturn c.composeTemplates[info.Component].OnStart(ctx, info, input)\n\tdefault:\n\t\treturn ctx\n\t}\n}\n\n// OnEnd is the callback function for the end event of a component.\n// implement the callbacks Handler interface.\nfunc (c *handlerTemplate) OnEnd(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\tswitch info.Component {\n\tcase components.ComponentOfPrompt:\n\t\treturn c.promptHandler.OnEnd(ctx, info, prompt.ConvCallbackOutput(output))\n\tcase components.ComponentOfChatModel:\n\t\treturn c.chatModelHandler.OnEnd(ctx, info, model.ConvCallbackOutput(output))\n\tcase components.ComponentOfEmbedding:\n\t\treturn c.embeddingHandler.OnEnd(ctx, info, embedding.ConvCallbackOutput(output))\n\tcase components.ComponentOfIndexer:\n\t\treturn c.indexerHandler.OnEnd(ctx, info, indexer.ConvCallbackOutput(output))\n\tcase components.ComponentOfRetriever:\n\t\treturn c.retrieverHandler.OnEnd(ctx, info, retriever.ConvCallbackOutput(output))\n\tcase components.ComponentOfLoader:\n\t\treturn c.loaderHandler.OnEnd(ctx, info, document.ConvLoaderCallbackOutput(output))\n\tcase components.ComponentOfTransformer:\n\t\treturn c.transformerHandler.OnEnd(ctx, info, document.ConvTransformerCallbackOutput(output))\n\tcase components.ComponentOfTool:\n\t\treturn c.toolHandler.OnEnd(ctx, info, tool.ConvCallbackOutput(output))\n\tcase compose.ComponentOfToolsNode:\n\t\treturn c.toolsNodeHandler.OnEnd(ctx, info, convToolsNodeCallbackOutput(output))\n\tcase adk.ComponentOfAgent:\n\t\treturn c.agentHandler.OnEnd(ctx, info, adk.ConvAgentCallbackOutput(output))\n\tcase compose.ComponentOfGraph,\n\t\tcompose.ComponentOfChain,\n\t\tcompose.ComponentOfLambda:\n\t\treturn c.composeTemplates[info.Component].OnEnd(ctx, info, output)\n\tdefault:\n\t\treturn ctx\n\t}\n}\n\n// OnError is the callback function for the error event of a component.\n// implement the callbacks Handler interface.\nfunc (c *handlerTemplate) OnError(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\tswitch info.Component {\n\tcase components.ComponentOfPrompt:\n\t\treturn c.promptHandler.OnError(ctx, info, err)\n\tcase components.ComponentOfChatModel:\n\t\treturn c.chatModelHandler.OnError(ctx, info, err)\n\tcase components.ComponentOfEmbedding:\n\t\treturn c.embeddingHandler.OnError(ctx, info, err)\n\tcase components.ComponentOfIndexer:\n\t\treturn c.indexerHandler.OnError(ctx, info, err)\n\tcase components.ComponentOfRetriever:\n\t\treturn c.retrieverHandler.OnError(ctx, info, err)\n\tcase components.ComponentOfLoader:\n\t\treturn c.loaderHandler.OnError(ctx, info, err)\n\tcase components.ComponentOfTransformer:\n\t\treturn c.transformerHandler.OnError(ctx, info, err)\n\tcase components.ComponentOfTool:\n\t\treturn c.toolHandler.OnError(ctx, info, err)\n\tcase compose.ComponentOfToolsNode:\n\t\treturn c.toolsNodeHandler.OnError(ctx, info, err)\n\tcase compose.ComponentOfGraph,\n\t\tcompose.ComponentOfChain,\n\t\tcompose.ComponentOfLambda:\n\t\treturn c.composeTemplates[info.Component].OnError(ctx, info, err)\n\tdefault:\n\t\treturn ctx\n\t}\n}\n\n// OnStartWithStreamInput is the callback function for the start event of a component with stream input.\n// implement the callbacks Handler interface.\nfunc (c *handlerTemplate) OnStartWithStreamInput(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {\n\tswitch info.Component {\n\t// currently no components.Component receive stream as input\n\tcase compose.ComponentOfGraph,\n\t\tcompose.ComponentOfChain,\n\t\tcompose.ComponentOfLambda:\n\t\treturn c.composeTemplates[info.Component].OnStartWithStreamInput(ctx, info, input)\n\tdefault:\n\t\treturn ctx\n\t}\n}\n\n// OnEndWithStreamOutput is the callback function for the end event of a component with stream output.\n// implement the callbacks Handler interface.\nfunc (c *handlerTemplate) OnEndWithStreamOutput(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context {\n\tswitch info.Component {\n\tcase components.ComponentOfChatModel:\n\t\treturn c.chatModelHandler.OnEndWithStreamOutput(ctx, info,\n\t\t\tschema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) (*model.CallbackOutput, error) {\n\t\t\t\treturn model.ConvCallbackOutput(item), nil\n\t\t\t}))\n\tcase components.ComponentOfTool:\n\t\treturn c.toolHandler.OnEndWithStreamOutput(ctx, info,\n\t\t\tschema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) (*tool.CallbackOutput, error) {\n\t\t\t\treturn tool.ConvCallbackOutput(item), nil\n\t\t\t}))\n\tcase compose.ComponentOfToolsNode:\n\t\treturn c.toolsNodeHandler.OnEndWithStreamOutput(ctx, info,\n\t\t\tschema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) ([]*schema.Message, error) {\n\t\t\t\treturn convToolsNodeCallbackOutput(item), nil\n\t\t\t}))\n\tcase compose.ComponentOfGraph,\n\t\tcompose.ComponentOfChain,\n\t\tcompose.ComponentOfLambda:\n\t\treturn c.composeTemplates[info.Component].OnEndWithStreamOutput(ctx, info, output)\n\tdefault:\n\t\treturn ctx\n\t}\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (c *handlerTemplate) Needed(ctx context.Context, info *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tif info == nil {\n\t\treturn false\n\t}\n\n\tswitch info.Component {\n\tcase components.ComponentOfChatModel:\n\t\tif c.chatModelHandler != nil && c.chatModelHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase components.ComponentOfEmbedding:\n\t\tif c.embeddingHandler != nil && c.embeddingHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase components.ComponentOfIndexer:\n\t\tif c.indexerHandler != nil && c.indexerHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase components.ComponentOfLoader:\n\t\tif c.loaderHandler != nil && c.loaderHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase components.ComponentOfPrompt:\n\t\tif c.promptHandler != nil && c.promptHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase components.ComponentOfRetriever:\n\t\tif c.retrieverHandler != nil && c.retrieverHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase components.ComponentOfTool:\n\t\tif c.toolHandler != nil && c.toolHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase components.ComponentOfTransformer:\n\t\tif c.transformerHandler != nil && c.transformerHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase compose.ComponentOfToolsNode:\n\t\tif c.toolsNodeHandler != nil && c.toolsNodeHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase adk.ComponentOfAgent:\n\t\tif c.agentHandler != nil && c.agentHandler.Needed(ctx, info, timing) {\n\t\t\treturn true\n\t\t}\n\tcase compose.ComponentOfGraph,\n\t\tcompose.ComponentOfChain,\n\t\tcompose.ComponentOfLambda:\n\t\thandler := c.composeTemplates[info.Component]\n\t\tif handler != nil {\n\t\t\tchecker, ok := handler.(callbacks.TimingChecker)\n\t\t\tif !ok || checker.Needed(ctx, info, timing) {\n\t\t\t\treturn true\n\t\t\t}\n\t\t}\n\tdefault:\n\t\treturn false\n\t}\n\n\treturn false\n}\n\n// LoaderCallbackHandler is the handler for the loader callback.\ntype LoaderCallbackHandler struct {\n\tOnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.LoaderCallbackInput) context.Context\n\tOnEnd   func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.LoaderCallbackOutput) context.Context\n\tOnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (ch *LoaderCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// TransformerCallbackHandler is the handler for the transformer callback.\ntype TransformerCallbackHandler struct {\n\tOnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context\n\tOnEnd   func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.TransformerCallbackOutput) context.Context\n\tOnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (ch *TransformerCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// EmbeddingCallbackHandler is the handler for the embedding callback.\ntype EmbeddingCallbackHandler struct {\n\tOnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *embedding.CallbackInput) context.Context\n\tOnEnd   func(ctx context.Context, runInfo *callbacks.RunInfo, output *embedding.CallbackOutput) context.Context\n\tOnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (ch *EmbeddingCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// IndexerCallbackHandler is the handler for the indexer callback.\ntype IndexerCallbackHandler struct {\n\tOnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *indexer.CallbackInput) context.Context\n\tOnEnd   func(ctx context.Context, runInfo *callbacks.RunInfo, output *indexer.CallbackOutput) context.Context\n\tOnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (ch *IndexerCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// ModelCallbackHandler is the handler for the model callback.\ntype ModelCallbackHandler struct {\n\tOnStart               func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context\n\tOnEnd                 func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.CallbackOutput) context.Context\n\tOnEndWithStreamOutput func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context\n\tOnError               func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (ch *ModelCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tcase callbacks.TimingOnEndWithStreamOutput:\n\t\treturn ch.OnEndWithStreamOutput != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// PromptCallbackHandler is the handler for the callback.\ntype PromptCallbackHandler struct {\n\t// OnStart is the callback function for the start of the callback.\n\tOnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context\n\t// OnEnd is the callback function for the end of the callback.\n\tOnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context\n\t// OnError is the callback function for the error of the callback.\n\tOnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (ch *PromptCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// RetrieverCallbackHandler is the handler for the retriever callback.\ntype RetrieverCallbackHandler struct {\n\t// OnStart is the callback function for the start of the retriever.\n\tOnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *retriever.CallbackInput) context.Context\n\t// OnEnd is the callback function for the end of the retriever.\n\tOnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *retriever.CallbackOutput) context.Context\n\t// OnError is the callback function for the error of the retriever.\n\tOnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (ch *RetrieverCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// ToolCallbackHandler is the handler for the tool callback.\ntype ToolCallbackHandler struct {\n\tOnStart               func(ctx context.Context, info *callbacks.RunInfo, input *tool.CallbackInput) context.Context\n\tOnEnd                 func(ctx context.Context, info *callbacks.RunInfo, output *tool.CallbackOutput) context.Context\n\tOnEndWithStreamOutput func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[*tool.CallbackOutput]) context.Context\n\tOnError               func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed checks if the callback handler is needed for the given timing.\nfunc (ch *ToolCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnEndWithStreamOutput:\n\t\treturn ch.OnEndWithStreamOutput != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\n// ToolsNodeCallbackHandlers defines optional callbacks for the Tools node\n// lifecycle events.\ntype ToolsNodeCallbackHandlers struct {\n\tOnStart               func(ctx context.Context, info *callbacks.RunInfo, input *schema.Message) context.Context\n\tOnEnd                 func(ctx context.Context, info *callbacks.RunInfo, input []*schema.Message) context.Context\n\tOnEndWithStreamOutput func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.Message]) context.Context\n\tOnError               func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context\n}\n\n// Needed reports whether a handler is registered for the given timing.\nfunc (ch *ToolsNodeCallbackHandlers) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tcase callbacks.TimingOnEndWithStreamOutput:\n\t\treturn ch.OnEndWithStreamOutput != nil\n\tcase callbacks.TimingOnError:\n\t\treturn ch.OnError != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n\nfunc convToolsNodeCallbackInput(src callbacks.CallbackInput) *schema.Message {\n\tswitch t := src.(type) {\n\tcase *schema.Message:\n\t\treturn t\n\tdefault:\n\t\treturn nil\n\t}\n}\n\nfunc convToolsNodeCallbackOutput(src callbacks.CallbackInput) []*schema.Message {\n\tswitch t := src.(type) {\n\tcase []*schema.Message:\n\t\treturn t\n\tdefault:\n\t\treturn nil\n\t}\n}\n\ntype AgentCallbackHandler struct {\n\tOnStart func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context\n\tOnEnd   func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context\n}\n\nfunc (ch *AgentCallbackHandler) Needed(ctx context.Context, info *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {\n\tswitch timing {\n\tcase callbacks.TimingOnStart:\n\t\treturn ch.OnStart != nil\n\tcase callbacks.TimingOnEnd:\n\t\treturn ch.OnEnd != nil\n\tdefault:\n\t\treturn false\n\t}\n}\n"
  },
  {
    "path": "utils/callbacks/template_test.go",
    "content": "/*\n * Copyright 2024 CloudWeGo Authors\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\npackage callbacks\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\n\t\"github.com/cloudwego/eino/adk\"\n\t\"github.com/cloudwego/eino/callbacks\"\n\t\"github.com/cloudwego/eino/components\"\n\t\"github.com/cloudwego/eino/components/document\"\n\t\"github.com/cloudwego/eino/components/embedding\"\n\t\"github.com/cloudwego/eino/components/indexer\"\n\t\"github.com/cloudwego/eino/components/model\"\n\t\"github.com/cloudwego/eino/components/prompt\"\n\t\"github.com/cloudwego/eino/components/retriever\"\n\t\"github.com/cloudwego/eino/components/tool\"\n\t\"github.com/cloudwego/eino/compose\"\n\t\"github.com/cloudwego/eino/schema\"\n)\n\nfunc TestNewComponentTemplate(t *testing.T) {\n\tt.Run(\"TestNewComponentTemplate\", func(t *testing.T) {\n\t\tcnt := 0\n\t\ttpl := NewHandlerHelper()\n\t\ttpl.ChatModel(&ModelCallbackHandler{\n\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.CallbackOutput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context {\n\t\t\t\toutput.Close()\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t}}).\n\t\t\tEmbedding(&EmbeddingCallbackHandler{\n\t\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *embedding.CallbackInput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *embedding.CallbackOutput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t}).\n\t\t\tPrompt(&PromptCallbackHandler{\n\t\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t}).\n\t\t\tRetriever(&RetrieverCallbackHandler{\n\t\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *retriever.CallbackInput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *retriever.CallbackOutput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t}).\n\t\t\tTool(&ToolCallbackHandler{\n\t\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *tool.CallbackInput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *tool.CallbackOutput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*tool.CallbackOutput]) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t\tOnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t},\n\t\t\t}).\n\t\t\tLambda(callbacks.NewHandlerBuilder().\n\t\t\t\tOnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t}).\n\t\t\t\tOnStartWithStreamInputFn(func(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {\n\t\t\t\t\tinput.Close()\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t}).\n\t\t\t\tOnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t}).\n\t\t\t\tOnEndWithStreamOutputFn(func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context {\n\t\t\t\t\toutput.Close()\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t}).\n\t\t\t\tOnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\t\tcnt++\n\t\t\t\t\treturn ctx\n\t\t\t\t}).Build()).\n\t\t\tHandler()\n\n\t\ttypes := []components.Component{\n\t\t\tcomponents.ComponentOfPrompt,\n\t\t\tcomponents.ComponentOfChatModel,\n\t\t\tcomponents.ComponentOfEmbedding,\n\t\t\tcomponents.ComponentOfRetriever,\n\t\t\tcomponents.ComponentOfTool,\n\t\t\tcompose.ComponentOfLambda,\n\t\t}\n\n\t\thandler := tpl.Handler()\n\t\tctx := context.Background()\n\t\tfor _, typ := range types {\n\t\t\thandler.OnStart(ctx, &callbacks.RunInfo{Component: typ}, nil)\n\t\t\thandler.OnEnd(ctx, &callbacks.RunInfo{Component: typ}, nil)\n\t\t\thandler.OnError(ctx, &callbacks.RunInfo{Component: typ}, fmt.Errorf(\"mock err\"))\n\n\t\t\tsir, siw := schema.Pipe[callbacks.CallbackInput](1)\n\t\t\tsiw.Close()\n\t\t\thandler.OnStartWithStreamInput(ctx, &callbacks.RunInfo{Component: typ}, sir)\n\n\t\t\tsor, sow := schema.Pipe[callbacks.CallbackOutput](1)\n\t\t\tsow.Close()\n\t\t\thandler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: typ}, sor)\n\t\t}\n\n\t\tassert.Equal(t, 22, cnt)\n\n\t\tctx = context.Background()\n\t\tctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, handler)\n\t\tcallbacks.OnStart[any](ctx, nil)\n\t\tassert.Equal(t, 22, cnt)\n\n\t\tctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfPrompt})\n\t\tctx = callbacks.OnStart[any](ctx, nil)\n\t\tassert.Equal(t, 23, cnt)\n\n\t\tctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer})\n\t\tcallbacks.OnEnd[any](ctx, nil)\n\t\tassert.Equal(t, 23, cnt)\n\n\t\tctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding})\n\t\tcallbacks.OnError(ctx, nil)\n\t\tassert.Equal(t, 24, cnt)\n\n\t\tctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader})\n\t\tcallbacks.OnStart[any](ctx, nil)\n\t\tassert.Equal(t, 24, cnt)\n\n\t\ttpl.Transformer(&TransformerCallbackHandler{\n\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.TransformerCallbackOutput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t}).Indexer(&IndexerCallbackHandler{\n\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *indexer.CallbackInput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *indexer.CallbackOutput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t}).Loader(&LoaderCallbackHandler{\n\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.LoaderCallbackInput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.LoaderCallbackOutput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t}).ToolsNode(&ToolsNodeCallbackHandlers{\n\t\t\tOnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *schema.Message) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[[]*schema.Message]) context.Context {\n\t\t\t\tcnt++\n\n\t\t\t\tif output == nil {\n\t\t\t\t\treturn ctx\n\t\t\t\t}\n\n\t\t\t\tfor {\n\t\t\t\t\t_, err := output.Recv()\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn ctx\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t},\n\t\t})\n\n\t\thandler = tpl.Handler()\n\t\tctx = context.Background()\n\t\tctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, handler)\n\n\t\tctx = callbacks.OnStart[any](ctx, nil)\n\t\tassert.Equal(t, 25, cnt)\n\n\t\tcallbacks.OnEnd[any](ctx, nil)\n\t\tassert.Equal(t, 26, cnt)\n\n\t\tctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader})\n\t\tcallbacks.OnEnd[any](ctx, nil)\n\t\tassert.Equal(t, 27, cnt)\n\n\t\tctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode})\n\t\tcallbacks.OnStart[any](ctx, nil)\n\t\tassert.Equal(t, 28, cnt)\n\n\t\tsr, sw := schema.Pipe[any](0)\n\t\tsw.Close()\n\t\tcallbacks.OnEndWithStreamOutput[any](ctx, sr)\n\t\tassert.Equal(t, 29, cnt)\n\n\t\tsr1, sw1 := schema.Pipe[[]*schema.Message](1)\n\t\tsw1.Send([]*schema.Message{{}}, nil)\n\t\tsw1.Close()\n\t\tcallbacks.OnEndWithStreamOutput[[]*schema.Message](ctx, sr1)\n\t\tassert.Equal(t, 30, cnt)\n\n\t\tcallbacks.OnError(ctx, nil)\n\t\tassert.Equal(t, 30, cnt)\n\n\t\tctx = callbacks.ReuseHandlers(ctx, nil)\n\t\tcallbacks.OnStart[any](ctx, nil)\n\t\tassert.Equal(t, 30, cnt)\n\t})\n}\n\nfunc TestAgentCallbackHandler(t *testing.T) {\n\tt.Run(\"Needed returns correct values\", func(t *testing.T) {\n\t\thandler := &AgentCallbackHandler{\n\t\t\tOnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context {\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t}\n\n\t\tctx := context.Background()\n\t\tinfo := &callbacks.RunInfo{Component: adk.ComponentOfAgent}\n\n\t\tassert.True(t, handler.Needed(ctx, info, callbacks.TimingOnStart))\n\t\tassert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))\n\t})\n\n\tt.Run(\"Needed with OnEnd set\", func(t *testing.T) {\n\t\thandler := &AgentCallbackHandler{\n\t\t\tOnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context {\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t}\n\n\t\tctx := context.Background()\n\t\tinfo := &callbacks.RunInfo{Component: adk.ComponentOfAgent}\n\n\t\tassert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart))\n\t\tassert.True(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))\n\t})\n\n\tt.Run(\"Needed with nil handlers\", func(t *testing.T) {\n\t\thandler := &AgentCallbackHandler{}\n\n\t\tctx := context.Background()\n\t\tinfo := &callbacks.RunInfo{Component: adk.ComponentOfAgent}\n\n\t\tassert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart))\n\t\tassert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))\n\t})\n}\n\nfunc TestHandlerHelperWithAgent(t *testing.T) {\n\tt.Run(\"Agent method sets handler correctly\", func(t *testing.T) {\n\t\tcnt := 0\n\t\ttpl := NewHandlerHelper()\n\t\ttpl.Agent(&AgentCallbackHandler{\n\t\t\tOnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t\tOnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context {\n\t\t\t\tcnt++\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t})\n\n\t\thandler := tpl.Handler()\n\t\tctx := context.Background()\n\t\tctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: adk.ComponentOfAgent}, handler)\n\n\t\tctx = callbacks.OnStart[any](ctx, nil)\n\t\tassert.Equal(t, 1, cnt)\n\n\t\tcallbacks.OnEnd[any](ctx, nil)\n\t\tassert.Equal(t, 2, cnt)\n\t})\n}\n\nfunc TestHandlerTemplateWithAgentComponent(t *testing.T) {\n\tt.Run(\"OnStart routes to agent handler\", func(t *testing.T) {\n\t\tcalled := false\n\t\ttpl := NewHandlerHelper()\n\t\ttpl.Agent(&AgentCallbackHandler{\n\t\t\tOnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context {\n\t\t\t\tcalled = true\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t})\n\n\t\thandler := tpl.Handler()\n\t\tctx := context.Background()\n\t\tinfo := &callbacks.RunInfo{Component: adk.ComponentOfAgent, Name: \"TestAgent\"}\n\n\t\thandler.OnStart(ctx, info, &adk.AgentCallbackInput{})\n\t\tassert.True(t, called)\n\t})\n\n\tt.Run(\"OnEnd routes to agent handler\", func(t *testing.T) {\n\t\tcalled := false\n\t\ttpl := NewHandlerHelper()\n\t\ttpl.Agent(&AgentCallbackHandler{\n\t\t\tOnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context {\n\t\t\t\tcalled = true\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t})\n\n\t\thandler := tpl.Handler()\n\t\tctx := context.Background()\n\t\tinfo := &callbacks.RunInfo{Component: adk.ComponentOfAgent, Name: \"TestAgent\"}\n\n\t\thandler.OnEnd(ctx, info, &adk.AgentCallbackOutput{})\n\t\tassert.True(t, called)\n\t})\n\n\tt.Run(\"Needed returns true for agent component\", func(t *testing.T) {\n\t\ttpl := NewHandlerHelper()\n\t\ttpl.Agent(&AgentCallbackHandler{\n\t\t\tOnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context {\n\t\t\t\treturn ctx\n\t\t\t},\n\t\t})\n\n\t\thandler := tpl.Handler()\n\t\tctx := context.Background()\n\t\tinfo := &callbacks.RunInfo{Component: adk.ComponentOfAgent}\n\n\t\tchecker, ok := handler.(callbacks.TimingChecker)\n\t\tassert.True(t, ok, \"handler should implement TimingChecker\")\n\t\tassert.True(t, checker.Needed(ctx, info, callbacks.TimingOnStart))\n\t})\n}\n"
  }
]