chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:49:10 +08:00
commit 3165c4acf8
382 changed files with 154008 additions and 0 deletions
+31
View File
@@ -0,0 +1,31 @@
coverage:
status:
project: #add everything under here, more options at https://docs.codecov.com/docs/commit-status
default: # default is the status check's name, not default settings
target: auto #default
threshold: 1% #allow coverage to drop by 1%
base: auto
if_ci_failed: error #success, failure, error, ignore
patch:
default:
target: 82% #default
threshold: 1% #allow coverage to drop by 1%
base: auto
if_ci_failed: error #success, failure, error, ignore
comment: #this is a top-level key
layout: " diff, flags, files"
behavior: default
require_changes: false # if true: only post the comment if coverage changes
require_base: false # [true :: must have a base report to post]
require_head: true # [true :: must have a head report to post]
hide_project_coverage: false # [true :: only show coverage on the git diff aka patch coverage]
# sample regex patterns
ignore:
- "tests"
- "examples/"
- "mock/"
- "callbacks/interface.go"
- "utils/safe"
- "components/tool/utils/create_options.go"
+41
View File
@@ -0,0 +1,41 @@
{
"allowedTypes": [
"feat",
"fix",
"docs",
"style",
"refactor",
"perf",
"test",
"build",
"ci",
"chore",
"revert"
],
"allowedScopes": [
"adk",
"adk/filesystem",
"callbacks",
"components",
"compose",
"deep",
"dynamictool",
"filesystem",
"flow",
"internal",
"middlewares",
"planexecute",
"plantask",
"prebuilt",
"reduction",
"schema",
"skill",
"summarization",
"supervisor",
"toolsearch",
"utils",
"docs",
"ci",
"serialization"
]
}
+40
View File
@@ -0,0 +1,40 @@
---
name: Bug report
about: Create a report to help us improve
title: ''
labels: ''
assignees: ''
---
**Describe the bug**
A clear and concise description of what the bug is.
**To Reproduce**
Steps to reproduce the behavior:
1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error
**Expected behavior**
A clear and concise description of what you expected to happen.
**Screenshots**
If applicable, add screenshots to help explain your problem.
**Version:**
Please provide the version of {project_name} you are using.
**Environment:**
The output of `go env`.
**Additional context**
Add any other context about the problem here.
+24
View File
@@ -0,0 +1,24 @@
---
name: Feature request
about: Suggest an idea for this project
title: ''
labels: ''
assignees: ''
---
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
**Describe the solution you'd like**
A clear and concise description of what you want to happen.
**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.
**Additional context**
Add any other context or screenshots about the feature request here.
+49
View File
@@ -0,0 +1,49 @@
#### What type of PR is this?
<!--
Add one of the following kinds:
build: Changes that affect the build system or external dependencies (example scopes: gulp, broccoli, npm)
ci: Changes to our CI configuration files and scripts (example scopes: Travis, Circle, BrowserStack, SauceLabs)
docs: Documentation only changes
feat: A new feature
optimize: A new optimization
fix: A bug fix
perf: A code change that improves performance
refactor: A code change that neither fixes a bug nor adds a feature
style: Changes that do not affect the meaning of the code (white space, formatting, missing semi-colons, etc)
test: Adding missing tests or correcting existing tests
chore: Changes to the build process or auxiliary tools and libraries such as documentation generation
-->
#### Check the PR title.
<!--
The description of the title will be attached in Release Notes,
so please describe it from user-oriented, what this PR does / why we need it.
Please check your PR title with the below requirements:
-->
- [ ] This PR title match the format: \<type\>(optional scope): \<description\>
- [ ] The description of this PR title is user-oriented and clear enough for others to understand.
- [ ] 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)
#### (Optional) Translate the PR title into Chinese.
#### (Optional) More detailed description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review(e.g., it's recommended to provide perf data if this is a perf type PR).
-->
en:
zh(optional):
#### (Optional) Which issue(s) this PR fixes:
<!--
Automatically closes linked issue when PR is merged.
Eg: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
-->
#### (optional) The PR that updates user documentation:
<!--
If 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)
-->
Binary file not shown.

After

Width:  |  Height:  |  Size: 85 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 304 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 242 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 784 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 446 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 186 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 39 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 256 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 150 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 311 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 171 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 171 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 82 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 61 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 226 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 482 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 188 KiB

+156
View File
@@ -0,0 +1,156 @@
name: Pull Request Check
on: [ pull_request ]
jobs:
compliant:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Check License Header
uses: apache/skywalking-eyes/header@v0.4.0
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- name: Check Spell
uses: crate-ci/typos@v1.42.3
golangci-lint:
runs-on: ubuntu-latest
permissions:
contents: write
pull-requests: write
repository-projects: write
steps:
- uses: actions/checkout@v4
- name: Set up Go
uses: actions/setup-go@v5
with:
go-version: 1.18
# for self-hosted, the cache path is shared across projects,
# and it works well without the cache of GitHub actions
# Enable it if we're going to use GitHub only
cache: true
- name: Golang CI Lint
# https://golangci-lint.run/
uses: golangci/golangci-lint-action@v9.2.0
with:
version: v2.8.0
args: --timeout 5m
commit-msg-check:
name: Commit Message Check
runs-on: ubuntu-latest
permissions:
contents: read
pull-requests: read
steps:
- uses: actions/checkout@v4
- name: Validate commit messages format and scope
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const fs = require('fs');
const pr = context.payload.pull_request;
if (!pr) {
core.setFailed('This workflow must run on pull_request events.');
return;
}
let allowedTypes = [];
let allowedScopes = [];
try {
const raw = fs.readFileSync('.github/.commit-rules.json', 'utf8');
const cfg = JSON.parse(raw);
allowedTypes = Array.isArray(cfg.allowedTypes) ? cfg.allowedTypes : [];
allowedScopes = Array.isArray(cfg.allowedScopes) ? cfg.allowedScopes : [];
} catch (e) {
core.setFailed('Cannot read .github/.commit-rules.json: ' + e.message);
return;
}
if (!allowedTypes.length) {
core.setFailed('allowedTypes is empty in .github/.commit-rules.json');
return;
}
const { owner, repo } = context.repo;
const pull_number = pr.number;
const commits = await github.paginate(
github.rest.pulls.listCommits,
{ owner, repo, pull_number, per_page: 100 }
);
let errors = [];
for (const c of commits) {
const sha = c.sha.slice(0, 7);
const subject = (c.commit.message || '').split('\n')[0];
const m = subject.match(/^([a-z]+)(\(([a-z0-9\-\/]+)\))?:\s(.+)$/);
if (!m) {
errors.push(`(${sha}) invalid format: "${subject}"`);
continue;
}
const type = m[1];
const scope = m[3]; // may be undefined
const desc = m[4];
if (!allowedTypes.includes(type)) {
errors.push(`(${sha}) invalid type "${type}"`);
}
if (!desc || !desc.trim()) {
errors.push(`(${sha}) description must be non-empty`);
}
if (scope) {
const topScope = scope.split('/')[0];
if (allowedScopes.length && !allowedScopes.includes(topScope)) {
errors.push(`(${sha}) invalid scope "${scope}"`);
}
}
}
if (errors.length) {
core.setFailed('Commit message check failed:\n' + errors.join('\n'));
} else {
core.info('All commit messages conform to "<type>(optional scope): <description>" and scope rules.');
}
pr-title-check:
name: PR Title Check
runs-on: ubuntu-latest
permissions:
contents: read
pull-requests: read
steps:
- uses: actions/checkout@v4
- name: Read commit rules
id: rules
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const fs = require('fs');
let cfg;
try {
const raw = fs.readFileSync('.github/.commit-rules.json', 'utf8');
cfg = JSON.parse(raw);
} catch (e) {
core.setFailed('Cannot read .github/.commit-rules.json: ' + e.message);
return;
}
const toMultiline = (list) => Array.isArray(list) ? list.join('\n') : '';
core.setOutput('types', toMultiline(cfg.allowedTypes));
core.setOutput('scopes', toMultiline(cfg.allowedScopes));
- name: Validate PR title
uses: amannn/action-semantic-pull-request@v6.1.1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
types: ${{ steps.rules.outputs.types }}
scopes: ${{ steps.rules.outputs.scopes }}
requireScope: false
+120
View File
@@ -0,0 +1,120 @@
name: Tag Notification
on:
push:
tags:
- 'v*'
jobs:
notify:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
with:
fetch-depth: 0
fetch-tags: true
- name: Fetch tag info
run: |
git fetch --tags -f
- name: Get tag info and send notification
run: |
# Get the tag name
TAG_NAME="${{ github.ref_name }}"
echo "Processing tag: $TAG_NAME"
# Get tag message
echo "Getting tag message..."
TAG_MESSAGE=$(git tag -l --format='%(contents)' "$TAG_NAME")
echo "Tag message:"
echo "$TAG_MESSAGE"
echo "---"
# Create base content parts
HEADER="### 🏷️ Eino New Tag Created: \`$TAG_NAME\`"
VERSION_INFO="📦 Version: \`$TAG_NAME\`"
# Prepare the message parts for jq
if [ ! -z "$TAG_MESSAGE" ]; then
# Pass all parts to jq and let it handle the formatting
jq -n \
--arg header "$HEADER" \
--arg version "$VERSION_INFO" \
--arg notes "$TAG_MESSAGE" \
--arg repo_url "https://github.com/${{ github.repository }}/releases/tag/$TAG_NAME" \
'{
"msg_type": "interactive",
"card": {
"elements": [
{
"tag": "markdown",
"content": ($header + "\n\n" + $version + "\n\n### 📝 Release Notes:\n" + $notes)
},
{
"tag": "action",
"actions": [
{
"tag": "button",
"text": {
"tag": "plain_text",
"content": "🔗 View Tag"
},
"url": $repo_url,
"type": "default"
}
]
}
],
"header": {
"title": {
"tag": "plain_text",
"content": "🏷️ Eino New Tag Created"
}
}
}
}' > webhook_payload.json
else
# Without release notes
jq -n \
--arg header "$HEADER" \
--arg version "$VERSION_INFO" \
--arg repo_url "https://github.com/${{ github.repository }}/releases/tag/$TAG_NAME" \
'{
"msg_type": "interactive",
"card": {
"elements": [
{
"tag": "markdown",
"content": ($header + "\n\n" + $version)
},
{
"tag": "action",
"actions": [
{
"tag": "button",
"text": {
"tag": "plain_text",
"content": "🔗 View Tag"
},
"url": $repo_url,
"type": "default"
}
]
}
],
"header": {
"title": {
"tag": "plain_text",
"content": "🏷️ Eino New Tag Created"
}
}
}
}' > webhook_payload.json
fi
# Send webhook
curl -X POST \
-H "Content-Type: application/json" \
-d @webhook_payload.json \
"${{ secrets.FEISHU_WEBHOOK_URL }}"
+262
View File
@@ -0,0 +1,262 @@
name: Eino Tests
on:
pull_request:
push:
branches:
- main
env:
DEFAULT_GO_VERSION: "1.18"
jobs:
unit-test:
name: eino-unit-test
runs-on: ubuntu-latest
permissions:
contents: write
pull-requests: write
repository-projects: write
env:
COVERAGE_FILE: coverage.out
BREAKDOWN_FILE: main.breakdown
steps:
- uses: actions/checkout@v4
- name: Set up Go
uses: actions/setup-go@v5
with:
go-version: ${{ env.DEFAULT_GO_VERSION }}
- name: Exec Go Test
run: |
modules=`find . -name "go.mod" -exec dirname {} \;`
echo $modules
list=""
coverpkg=""
if [[ ! -f "go.work" ]];then go work init;fi
for module in $modules; do go work use $module; list=$module"/... "$list; coverpkg=$module"/...,"$coverpkg; done
go work sync
go test -race -v -coverprofile=${{ env.COVERAGE_FILE }} -gcflags="all=-l -N" -coverpkg=$coverpkg $list
- name: Upload coverage to Codecov
uses: codecov/codecov-action@v5
with:
name: eino-unit-test
env_vars: GOLANG,EINO
files: ${{ env.COVERAGE_FILE }}
token: ${{ secrets.CODECOV_TOKEN }}
codecov_yml_path: ./github/.codecov.yml
benchmark-test:
runs-on: ubuntu-latest
permissions:
contents: write
pull-requests: write
repository-projects: write
steps:
- uses: actions/checkout@v4
- name: Set up Go
uses: actions/setup-go@v5
with:
go-version: ${{ env.DEFAULT_GO_VERSION }}
- name: Run Benchmark Tests
run: go test -bench=. -benchmem -run=none ./...
compatibility-test:
strategy:
matrix:
go: [ "1.19", "1.20", "1.21", "1.22", "1.23", "1.24" ]
runs-on: ubuntu-latest
permissions:
contents: write
pull-requests: write
repository-projects: write
steps:
- uses: actions/checkout@v4
- name: Set up Go
uses: actions/setup-go@v5
with:
go-version: ${{ matrix.go }}
cache: true
- name: Compatibility Test
run: |
# just basic unit test, no coverage report
go test -race ./...
api-compatibility:
name: api-compatibility-check
runs-on: ubuntu-latest
permissions:
contents: write
pull-requests: write
repository-projects: write
if: github.event_name == 'pull_request'
steps:
- uses: actions/checkout@v4
with:
fetch-depth: 0
- name: Set up Go
uses: actions/setup-go@v5
with:
go-version: "1.22"
- name: Install go-apidiff
run: go install github.com/joelanford/go-apidiff@v0.8.2
- name: Check API compatibility
id: apidiff
run: |
BASE_SHA=${{ github.event.pull_request.base.sha }}
HEAD_SHA=${{ github.event.pull_request.head.sha }}
echo "Checking API compatibility between $BASE_SHA and $HEAD_SHA"
go mod tidy
if ! DIFF_OUTPUT=$(go-apidiff $BASE_SHA $HEAD_SHA 2>&1); then
echo "go-apidiff output: $DIFF_OUTPUT"
fi
echo "diff_output<<EOF" >> $GITHUB_ENV
echo "$DIFF_OUTPUT" >> $GITHUB_ENV
echo "EOF" >> $GITHUB_ENV
if echo "$DIFF_OUTPUT" | grep -q "Incompatible changes:"; then
echo "has_breaking_changes=true" >> $GITHUB_OUTPUT
else
echo "has_breaking_changes=false" >> $GITHUB_OUTPUT
fi
- name: Create Review Thread
if: steps.apidiff.outputs.has_breaking_changes == 'true'
continue-on-error: true
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const reviewComments = await github.rest.pulls.listReviewComments({
owner: context.repo.owner,
repo: context.repo.repo,
pull_number: context.issue.number
});
const existingPackageComments = new Map();
for (const comment of reviewComments.data) {
if (comment.body.includes('Breaking API Changes Detected')) {
const packageMatch = comment.body.match(/Package: `([^`]+)`/);
if (packageMatch) {
const pkg = packageMatch[1];
if (!existingPackageComments.has(pkg)) {
existingPackageComments.set(pkg, new Set());
}
existingPackageComments.get(pkg).add(comment.path);
}
}
}
const files = await github.rest.pulls.listFiles({
owner: context.repo.owner,
repo: context.repo.repo,
pull_number: context.issue.number
});
const diffOutput = process.env.diff_output || '';
const breakingChanges = new Map();
let currentPackage = '';
let isInIncompatibleSection = false;
const lines = diffOutput.split('\n');
for (let i = 0; i < lines.length; i++) {
const line = lines[i].trim();
if (line.startsWith('github.com/')) {
currentPackage = line;
if (!breakingChanges.has(currentPackage)) {
breakingChanges.set(currentPackage, []);
}
continue;
}
if (line === 'Incompatible changes:') {
isInIncompatibleSection = true;
continue;
}
if (line === '') {
isInIncompatibleSection = false;
continue;
}
if (isInIncompatibleSection && line.startsWith('- ')) {
const change = line.substring(2);
if (currentPackage) {
breakingChanges.get(currentPackage).push(change);
}
}
}
const changedFiles = files.data;
for (const [pkg, changes] of breakingChanges) {
if (changes.length === 0) continue;
const pkgPath = pkg.split('/').slice(3).join('/');
const matchingFile = changedFiles.find(file =>
file.filename.includes(pkgPath)
) || changedFiles[0];
const hasCommentForPackage = existingPackageComments.has(pkg) &&
existingPackageComments.get(pkg).has(matchingFile.filename);
if (matchingFile && !hasCommentForPackage) {
const changesList = changes.map(change => {
const [name, desc] = change.split(':').map(s => s.trim());
return `- **${name}:** ${desc}`;
}).join('\n');
const commentBody = [
'🚨 **Breaking API Changes Detected**',
'',
`Package: \`${pkg}\``,
'',
'Incompatible changes:',
changesList,
'',
'<details>',
'<summary>Review Guidelines</summary>',
'',
'Please ensure that:',
'- The changes are absolutely necessary',
'- They are properly documented',
'- Migration guides are provided if needed',
'</details>',
'',
'⚠️ Please resolve this thread after reviewing the breaking changes.'
].join('\n');
await github.rest.pulls.createReview({
owner: context.repo.owner,
repo: context.repo.repo,
pull_number: context.issue.number,
event: 'COMMENT',
comments: [{
path: matchingFile.filename,
position: matchingFile.patch ? matchingFile.patch.split('\n').findIndex(line => line.startsWith('+')) + 1 : 1,
body: commentBody
}]
});
if (!existingPackageComments.has(pkg)) {
existingPackageComments.set(pkg, new Set());
}
existingPackageComments.get(pkg).add(matchingFile.filename);
}
}
+66
View File
@@ -0,0 +1,66 @@
# Binaries for programs and plugins
*.exe
*.exe~
*.dll
*.so
*.dylib
# Test binary, built with `go test -c`
*.test
# Output of the go coverage tool, specifically when used with LiteIDE
*.out
# Dependency directories (remove the comment below to include it)
# vendor/
# Go workspace file
go.work
go.work.sum
# env file
.env
# the result of the go build
output*
output/*
# Files generated by IDEs
.idea/
*.iml
# Vim swap files
*.swp
# Vscode files
.vscode
/patches
/vendor
# Trae files
.trae
# Specs files (internal documentation)
**/specs/
# Reports (generated analysis files)
reports/
/todos
.DS_Store
*.log*
.claude
CLAUDE.md
*.jsonl
*.txt
.agents/
# Specs directories
*/specs
/todos
/.claude/
# Internal dev setup (not for public repo)
/scripts/dev_setup_internal.sh
+148
View File
@@ -0,0 +1,148 @@
# output configuration options
version: "2"
# All available settings of specific linters.
# Refer to https://golangci-lint.run/usage/linters
linters:
default: standard
enable:
- revive
- godoclint
- funlen
- cyclop
disable:
- errcheck
- staticcheck
- unused
- ineffassign
exclusions:
generated: lax
paths:
- ".*_test.go"
- ".*_mock.go"
rules:
- path: "^internal/.*"
linters:
- revive
- text: "var-naming: don't use underscores in Go names"
linters:
- revive
- path: "/utils/"
text: "var-naming: avoid meaningless package names"
linters:
- revive
- text: "exported: type name will be used as agent.AgentOption by other packages"
linters:
- revive
- path: "adk/prebuilt/deep/task_tool.go"
text: "argument-limit: maximum number of arguments per function exceeded"
linters:
- revive
- path: "compose/component_to_graph_node.go"
text: "argument-limit: maximum number of arguments per function exceeded"
linters:
- revive
- path: "compose/graph_run.go"
text: "argument-limit: maximum number of arguments per function exceeded"
linters:
- revive
- path: "adk/workflow.go"
text: "argument-limit: maximum number of arguments per function exceeded"
linters:
- revive
- path: "compose/graph.go"
linters:
- cyclop
text: "calculated cyclomatic complexity for function compile"
- path: "schema/message.go"
linters:
- cyclop
text: "calculated cyclomatic complexity for function ConcatMessages"
- path: "compose/graph_run.go"
linters:
- cyclop
text: "calculated cyclomatic complexity for function run"
- path: "compose/graph.go"
linters:
- funlen
text: "Function 'compile' is too long"
- path: "compose/graph_run.go"
linters:
- funlen
text: "Function 'run' is too long"
settings:
govet:
enable-all: true
# Disable analyzers by name.
# Run `go tool vet help` to see all analyzers.
disable:
- fieldalignment
revive:
# Sets the default failure confidence.
# This means that linting errors with less than 0.8 confidence will be ignored.
# Default: 0.8
confidence: 0.8
rules:
# Exported function and methods should have comments.
- name: exported
severity: error
exclude:
- "^internal/.*"
arguments:
- "disable-checks-on-constants"
- "disable-checks-on-variables"
- "disable-checks-on-types"
- "disable-checks-on-methods"
- name: package-comments
disabled: false
- name: var-naming
disabled: false
arguments:
# AllowList
- [ "utils", "s_", "err_", "err__", "plan_", "userInput_", "executedSteps_", "executedStep_", "iterator_", "in_", "out_" ]
# DenyList
- [ ]
- - extra-bad-package-names:
- helpers
- models
- name: argument-limit
arguments: [ 6 ]
- name: function-length
arguments: [ 120, 0 ]
godoclint:
check-exported: true
require-package-documentation: true
funlen:
lines: 200
statements: 120
cyclop:
max-complexity: 40
package-average: 20
formatters:
enable:
- gci
- gofmt
settings:
gofmt:
# Simplify code: gofmt with `-s` option.
# Default: true
simplify: true
# Apply the rewrite rules to the source before reformatting.
# https://pkg.go.dev/cmd/gofmt
# Default: []
rewrite-rules:
- pattern: 'interface{}'
replacement: 'any'
- pattern: 'a[b:len(a)]'
replacement: 'a[b:]'
gci:
# Section configuration to compare against.
# Section names are case-insensitive and may contain parameters in ().
# The default order of sections is `standard > default > custom > blank > dot > alias > localmodule`.
# If `custom-order` is `true`, it follows the order of `sections` option.
# Default: ["standard", "default"]
sections:
- standard
- default
- localmodule
custom-order: true
+27
View File
@@ -0,0 +1,27 @@
header:
license:
spdx-id: Apache-2.0
copyright-owner: CloudWeGo Authors
template: |
/*
* Copyright {{ .Year }} CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
paths:
- '**/*.go'
- '**/*.s'
comment: on-failure
+128
View File
@@ -0,0 +1,128 @@
# Contributor Covenant Code of Conduct
## Our Pledge
We as members, contributors, and leaders pledge to make participation in our
community a harassment-free experience for everyone, regardless of age, body
size, visible or invisible disability, ethnicity, sex characteristics, gender
identity and expression, level of experience, education, socio-economic status,
nationality, personal appearance, race, religion, or sexual identity
and orientation.
We pledge to act and interact in ways that contribute to an open, welcoming,
diverse, inclusive, and healthy community.
## Our Standards
Examples of behavior that contributes to a positive environment for our
community include:
* Demonstrating empathy and kindness toward other people
* Being respectful of differing opinions, viewpoints, and experiences
* Giving and gracefully accepting constructive feedback
* Accepting responsibility and apologizing to those affected by our mistakes,
and learning from the experience
* Focusing on what is best not just for us as individuals, but for the
overall community
Examples of unacceptable behavior include:
* The use of sexualized language or imagery, and sexual attention or
advances of any kind
* Trolling, insulting or derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or email
address, without their explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting
## Enforcement Responsibilities
Community leaders are responsible for clarifying and enforcing our standards of
acceptable behavior and will take appropriate and fair corrective action in
response to any behavior that they deem inappropriate, threatening, offensive,
or harmful.
Community leaders have the right and responsibility to remove, edit, or reject
comments, commits, code, wiki edits, issues, and other contributions that are
not aligned to this Code of Conduct, and will communicate reasons for moderation
decisions when appropriate.
## Scope
This Code of Conduct applies within all community spaces, and also applies when
an individual is officially representing the community in public spaces.
Examples of representing our community include using an official e-mail address,
posting via an official social media account, or acting as an appointed
representative at an online or offline event.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported to the community leaders responsible for enforcement at
conduct@cloudwego.io.
All complaints will be reviewed and investigated promptly and fairly.
All community leaders are obligated to respect the privacy and security of the
reporter of any incident.
## Enforcement Guidelines
Community leaders will follow these Community Impact Guidelines in determining
the consequences for any action they deem in violation of this Code of Conduct:
### 1. Correction
**Community Impact**: Use of inappropriate language or other behavior deemed
unprofessional or unwelcome in the community.
**Consequence**: A private, written warning from community leaders, providing
clarity around the nature of the violation and an explanation of why the
behavior was inappropriate. A public apology may be requested.
### 2. Warning
**Community Impact**: A violation through a single incident or series
of actions.
**Consequence**: A warning with consequences for continued behavior. No
interaction with the people involved, including unsolicited interaction with
those enforcing the Code of Conduct, for a specified period of time. This
includes avoiding interactions in community spaces as well as external channels
like social media. Violating these terms may lead to a temporary or
permanent ban.
### 3. Temporary Ban
**Community Impact**: A serious violation of community standards, including
sustained inappropriate behavior.
**Consequence**: A temporary ban from any sort of interaction or public
communication with the community for a specified period of time. No public or
private interaction with the people involved, including unsolicited interaction
with those enforcing the Code of Conduct, is allowed during this period.
Violating these terms may lead to a permanent ban.
### 4. Permanent Ban
**Community Impact**: Demonstrating a pattern of violation of community
standards, including sustained inappropriate behavior, harassment of an
individual, or aggression toward or disparagement of classes of individuals.
**Consequence**: A permanent ban from any sort of public interaction within
the community.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage],
version 2.0, available at
https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
Community Impact Guidelines were inspired by [Mozilla's code of conduct
enforcement ladder](https://github.com/mozilla/diversity).
[homepage]: https://www.contributor-covenant.org
For answers to common questions about this code of conduct, see the FAQ at
https://www.contributor-covenant.org/faq. Translations are available at
https://www.contributor-covenant.org/translations.
+51
View File
@@ -0,0 +1,51 @@
# How to Contribute
## Your First Pull Request
We 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).
## Branch Organization
We 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)
## Bugs
### 1. How to Find Known Issues
We 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 doesnt already exist.
### 2. Reporting New Issues
Providing a reduced test code is a recommended way for reporting issues. Then can place in:
- Just in issues
- [Golang Playground](https://play.golang.org/)
### 3. Security Bugs
Please do not report the safe disclosure of bugs to public issues. Contact us by [Support Email](mailto:conduct@cloudwego.io)
## How to Get in Touch
- [Email](mailto:conduct@cloudwego.io)
## Submit a Pull Request
Before you submit your Pull Request (PR) consider the following guidelines:
1. 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.
2. 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.
3. [Fork](https://docs.github.com/en/github/getting-started-with-github/fork-a-repo) the cloudwego {project_name} repo.
4. In your forked repository, make your changes in a new git branch:
```
git checkout -b my-fix-branch develop
```
5. Create your patch, including appropriate test cases.
6. Follow our [Style Guides](#code-style-guides).
7. Commit your changes using a descriptive commit message that follows [AngularJS Git Commit Message Conventions](https://docs.google.com/document/d/1QrDFcIiPjSLDn3EL15IJygNPiHORgU1_OOAqWjiDU5Y/edit).
Adherence to these conventions is necessary because release notes are automatically generated from these messages.
8. Push your branch to GitHub:
```
git push origin my-fix-branch
```
9. In GitHub, send a pull request to `{project_name}:develop`
## Contribution Prerequisites
- Our development environment keeps up with [Go Official](https://golang.org/project/).
- 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)
- You are familiar with [GitHub](https://github.com)
- Maybe you need familiar with [Actions](https://github.com/features/actions)(our default workflow tool).
## Code Style Guides
- [Effective Go](https://golang.org/doc/effective_go)
- [Go Code Review Comments](https://github.com/golang/go/wiki/CodeReviewComments)
+177
View File
@@ -0,0 +1,177 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
+209
View File
@@ -0,0 +1,209 @@
# Eino
![coverage](https://raw.githubusercontent.com/cloudwego/eino/badges/.badges/main/coverage.svg)
[![Release](https://img.shields.io/github/v/release/cloudwego/eino)](https://github.com/cloudwego/eino/releases)
[![WebSite](https://img.shields.io/website?up_message=cloudwego&url=https%3A%2F%2Fwww.cloudwego.io%2F)](https://www.cloudwego.io/)
[![License](https://img.shields.io/github/license/cloudwego/eino)](https://github.com/cloudwego/eino/blob/main/LICENSE)
[![Go Report Card](https://goreportcard.com/badge/github.com/cloudwego/eino)](https://goreportcard.com/report/github.com/cloudwego/eino)
[![OpenIssue](https://img.shields.io/github/issues/cloudwego/eino)](https://github.com/cloudwego/kitex/eino)
[![ClosedIssue](https://img.shields.io/github/issues-closed/cloudwego/eino)](https://github.com/cloudwego/eino/issues?q=is%3Aissue+is%3Aclosed)
![Stars](https://img.shields.io/github/stars/cloudwego/eino)
![Forks](https://img.shields.io/github/forks/cloudwego/eino)
English | [中文](README.zh_CN.md)
# Overview
**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.
Eino provides:
- **[Components](https://github.com/cloudwego/eino-ext)**: reusable building blocks like `ChatModel`, `Tool`, `Retriever`, and `ChatTemplate`, with official implementations for OpenAI, Ollama, and more.
- **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.
- **Composition**: connect components into graphs and workflows that can run standalone or be exposed as tools for agents.
- **[Examples](https://github.com/cloudwego/eino-examples)**: working code for common patterns and real-world use cases.
![](.github/static/img/eino/eino_concept.jpeg)
# Quick Start
## ChatModelAgent
Configure a ChatModel, optionally add tools, and you have a working agent:
```Go
chatModel, _ := openai.NewChatModel(ctx, &openai.ChatModelConfig{
Model: "gpt-4o",
APIKey: os.Getenv("OPENAI_API_KEY"),
})
agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
Model: chatModel,
})
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent})
iter := runner.Query(ctx, "Hello, who are you?")
for {
event, ok := iter.Next()
if !ok {
break
}
fmt.Println(event.Message.Content)
}
```
Add tools to give the agent capabilities:
```Go
agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
Model: chatModel,
ToolsConfig: adk.ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{weatherTool, calculatorTool},
},
},
})
```
The agent handles the ReAct loop internally — it decides when to call tools and when to respond.
→ [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/)
## DeepAgent
For complex tasks, use DeepAgent. It breaks down problems into steps, delegates to sub-agents, and tracks progress:
```Go
deepAgent, _ := deep.New(ctx, &deep.Config{
ChatModel: chatModel,
SubAgents: []adk.Agent{researchAgent, codeAgent},
ToolsConfig: adk.ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{shellTool, pythonTool, webSearchTool},
},
},
})
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: deepAgent})
iter := runner.Query(ctx, "Analyze the sales data in report.csv and generate a summary chart")
```
DeepAgent can be configured to coordinate multiple specialized agents, run shell commands, execute Python code, and search the web.
→ [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/)
## Composition
When you need precise control over execution flow, use `compose` to build graphs and workflows:
```Go
graph := compose.NewGraph[*Input, *Output]()
graph.AddLambdaNode("validate", validateFn)
graph.AddChatModelNode("generate", chatModel)
graph.AddLambdaNode("format", formatFn)
graph.AddEdge(compose.START, "validate")
graph.AddEdge("validate", "generate")
graph.AddEdge("generate", "format")
graph.AddEdge("format", compose.END)
runnable, _ := graph.Compile(ctx)
result, _ := runnable.Invoke(ctx, input)
```
Compositions can be exposed as tools for agents, bridging deterministic workflows with autonomous behavior:
```Go
tool, _ := graphtool.NewInvokableGraphTool(graph, "data_pipeline", "Process and validate data")
agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
Model: chatModel,
ToolsConfig: adk.ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{tool},
},
},
})
```
This lets you build domain-specific pipelines with exact control, then let agents decide when to use them.
→ [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/)
# Key Features
## Component Ecosystem
Eino defines component abstractions (ChatModel, Tool, Retriever, Embedding, etc.) with official implementations for OpenAI, Claude, Gemini, Ark, Ollama, Elasticsearch, and more.
→ [eino-ext](https://github.com/cloudwego/eino-ext)
## Stream Processing
Eino 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.
→ [docs](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/)
## Callback Aspects
Inject logging, tracing, and metrics at fixed points (OnStart, OnEnd, OnError, OnStartWithStreamInput, OnEndWithStreamOutput) across components, graphs, and agents.
→ [docs](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/callback_manual/)
## Interrupt/Resume
Any agent or tool can pause execution for human input and resume from checkpoint. The framework handles state persistence and routing.
→ [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)
# Framework Structure
![](.github/static/img/eino/eino_framework.jpeg)
The Eino framework consists of:
- Eino (this repo): Type definitions, streaming mechanism, component abstractions, orchestration, agent implementations, aspect mechanisms
- [EinoExt](https://github.com/cloudwego/eino-ext): Component implementations, callback handlers, usage examples, evaluators, prompt optimizers
- [Eino Devops](https://github.com/cloudwego/eino-ext/tree/main/devops): Visualized development and debugging
- [EinoExamples](https://github.com/cloudwego/eino-examples): Example applications and best practices
## Documentation
- [Eino User Manual](https://www.cloudwego.io/zh/docs/eino/)
- [Eino: Quick Start](https://www.cloudwego.io/zh/docs/eino/quick_start/)
## Dependencies
- Go 1.18 and above.
## Code Style
This repo uses `golangci-lint`. Check locally with:
```bash
golangci-lint run ./...
```
Rules enforced:
- Exported functions, interfaces, packages, etc. should have GoDoc comments
- Code should be formatted with `gofmt -s`
- Import order should follow `goimports` (std -> third party -> local)
## Security
If you discover a potential security issue in this project, or think you may
have discovered a security issue, we ask that you notify Bytedance Security via our [security center](https://security.bytedance.com/src) or [vulnerability reporting email](mailto:sec@bytedance.com?subject=Feedback%20On%20Eino).
Do **not** create a public GitHub issue.
## Contact
- Membership: [COMMUNITY MEMBERSHIP](https://github.com/cloudwego/community/blob/main/COMMUNITY_MEMBERSHIP.md)
- Issues: [Issues](https://github.com/cloudwego/eino/issues)
- Lark: Scan the QR code below with [Feishu](https://www.feishu.cn/en/) to join the CloudWeGo/eino user group.
&ensp;&ensp;&ensp; <img src=".github/static/img/eino/lark_group_zh.png" alt="LarkGroup" width="200"/>
## License
This project is licensed under the [Apache-2.0 License](LICENSE-APACHE).
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目:`cloudwego/eino`
- 原始仓库:https://github.com/cloudwego/eino
- 导入方式:上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源,不代表 WeHub 是原项目作者
+210
View File
@@ -0,0 +1,210 @@
# Eino
![coverage](https://raw.githubusercontent.com/cloudwego/eino/badges/.badges/main/coverage.svg)
[![Release](https://img.shields.io/github/v/release/cloudwego/eino)](https://github.com/cloudwego/eino/releases)
[![WebSite](https://img.shields.io/website?up_message=cloudwego&url=https%3A%2F%2Fwww.cloudwego.io%2F)](https://www.cloudwego.io/)
[![License](https://img.shields.io/github/license/cloudwego/eino)](https://github.com/cloudwego/eino/blob/main/LICENSE)
[![Go Report Card](https://goreportcard.com/badge/github.com/cloudwego/eino)](https://goreportcard.com/report/github.com/cloudwego/eino)
[![OpenIssue](https://img.shields.io/github/issues/cloudwego/eino)](https://github.com/cloudwego/kitex/eino)
[![ClosedIssue](https://img.shields.io/github/issues-closed/cloudwego/eino)](https://github.com/cloudwego/eino/issues?q=is%3Aissue+is%3Aclosed)
![Stars](https://img.shields.io/github/stars/cloudwego/eino)
![Forks](https://img.shields.io/github/forks/cloudwego/eino)
[English](README.md) | 中文
# 简介
**Eino['aino]** 是一个 Go 语言的 LLM 应用开发框架,借鉴了 LangChain、Google ADK 等开源项目,按照 Go 的惯例设计。
Eino 提供:
- **[组件](https://github.com/cloudwego/eino-ext)**`ChatModel``Tool``Retriever``ChatTemplate` 等可复用模块,官方实现覆盖 OpenAI、Ollama 等
- **智能体开发套件(ADK)**:支持工具调用、多智能体协同、上下文管理、中断/恢复等人机交互,以及开箱即用的智能体模式
- **编排**:把组件组装成图或工作流,既能独立运行,也能作为工具给智能体调用
- **[示例](https://github.com/cloudwego/eino-examples)**:常见模式和实际场景的可运行代码
![](.github/static/img/eino/eino_concept.jpeg)
# 快速上手
## ChatModelAgent
配置好 ChatModel,加上工具(可选),就能跑起来:
```Go
chatModel, _ := openai.NewChatModel(ctx, &openai.ChatModelConfig{
Model: "gpt-4o",
APIKey: os.Getenv("OPENAI_API_KEY"),
})
agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
Model: chatModel,
})
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent})
iter := runner.Query(ctx, "Hello, who are you?")
for {
event, ok := iter.Next()
if !ok {
break
}
fmt.Println(event.Message.Content)
}
```
加工具让智能体有更多能力:
```Go
agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
Model: chatModel,
ToolsConfig: adk.ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{weatherTool, calculatorTool},
},
},
})
```
智能体内部自动处理 ReAct 循环,自己判断什么时候调工具、什么时候回复。
→ [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/)
## DeepAgent
复杂任务用 DeepAgent,它会把问题拆成步骤,分派给子智能体,并追踪进度:
```Go
deepAgent, _ := deep.New(ctx, &deep.Config{
ChatModel: chatModel,
SubAgents: []adk.Agent{researchAgent, codeAgent},
ToolsConfig: adk.ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{shellTool, pythonTool, webSearchTool},
},
},
})
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: deepAgent})
iter := runner.Query(ctx, "Analyze the sales data in report.csv and generate a summary chart")
```
DeepAgent 可以配置成:协调多个专业智能体、跑 shell 命令、执行 Python、搜索网络。
→ [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/)
## 编排
需要精确控制执行流程时,用 `compose` 搭图或工作流:
```Go
graph := compose.NewGraph[*Input, *Output]()
graph.AddLambdaNode("validate", validateFn)
graph.AddChatModelNode("generate", chatModel)
graph.AddLambdaNode("format", formatFn)
graph.AddEdge(compose.START, "validate")
graph.AddEdge("validate", "generate")
graph.AddEdge("generate", "format")
graph.AddEdge("format", compose.END)
runnable, _ := graph.Compile(ctx)
result, _ := runnable.Invoke(ctx, input)
```
编排出来的流程可以包装成工具给智能体用,把确定性流程和自主决策结合起来:
```Go
tool, _ := graphtool.NewInvokableGraphTool(graph, "data_pipeline", "Process and validate data")
agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
Model: chatModel,
ToolsConfig: adk.ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{tool},
},
},
})
```
这样你可以写出精确可控的业务流程,再让智能体决定什么时候调用。
→ [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/)
# 主要特性
## 组件生态
Eino 定义了组件抽象(ChatModel、Tool、Retriever、Embedding 等),官方实现覆盖 OpenAI、Claude、Gemini、Ark、Ollama、Elasticsearch 等。
→ [eino-ext](https://github.com/cloudwego/eino-ext)
## 流式处理
Eino 在编排中自动处理流式:拼接、装箱、合并、复制。组件只需实现有业务意义的流式范式,框架处理剩下的。
→ [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/)
## 回调切面
在固定切点(OnStart、OnEnd、OnError、OnStartWithStreamInput、OnEndWithStreamOutput)注入日志、追踪、指标,适用于组件、图、智能体。
→ [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/chain_and_graph_orchestration/callback_manual/)
## 中断/恢复
任何智能体或工具都能暂停等待人工输入,从检查点恢复。框架处理状态持久化和路由。
→ [文档](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)
# 框架结构
![](.github/static/img/eino/eino_framework.jpeg)
Eino 框架包含:
- Eino(本仓库):类型定义、流处理机制、组件抽象、编排、智能体实现、切面机制
- [EinoExt](https://github.com/cloudwego/eino-ext):组件实现、回调处理器、使用示例、评估器、提示优化器
- [Eino Devops](https://github.com/cloudwego/eino-ext/tree/main/devops):可视化开发和调试
- [EinoExamples](https://github.com/cloudwego/eino-examples):示例应用和最佳实践
## 文档
- [Eino 用户手册](https://www.cloudwego.io/zh/docs/eino/)
- [Eino: 快速开始](https://www.cloudwego.io/zh/docs/eino/quick_start/)
## 依赖
- Go 1.18 及以上
## 代码规范
本仓库使用 `golangci-lint`,本地检查:
```bash
golangci-lint run ./...
```
规则:
- 导出的函数、接口、package 等需要 GoDoc 注释
- 代码格式符合 `gofmt -s`
- import 顺序符合 `goimports`std -> third party -> local
## 安全
如果你在该项目中发现潜在的安全问题,或你认为可能发现了安全问题,请通过我们的[安全中心](https://security.bytedance.com/src)
或[漏洞报告邮箱](mailto:sec@bytedance.com?subject=关于Eino的反馈)通知字节跳动安全团队。
请**不要**创建公开的 GitHub Issue。
## 联系我们
- 成为 member[COMMUNITY MEMBERSHIP](https://github.com/cloudwego/community/blob/main/COMMUNITY_MEMBERSHIP.md)
- Issues[Issues](https://github.com/cloudwego/eino/issues)
- 飞书:扫码加入 CloudWeGo/eino 用户群
&ensp;&ensp;&ensp; <img src=".github/static/img/eino/lark_group_zh.png" alt="LarkGroup" width="200"/>
## 开源许可证
本项目基于 [Apache-2.0 许可证](LICENSE-APACHE) 开源。
+14
View File
@@ -0,0 +1,14 @@
# Typo check: https://github.com/crate-ci/typos
[default]
[default.extend-words]
Invokable = "Invokable"
invokable = "invokable"
InvokableLambda = "InvokableLambda"
InvokableRun = "InvokableRun"
typ = "typ"
byted = "byted"
cpy = "cpy"
mak = "mak"
[files]
extend-exclude = ["go.mod", "go.sum", "check_branch_name.sh"]
+417
View File
@@ -0,0 +1,417 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Package adk provides core agent development kit utilities and types.
package adk
import (
"context"
"errors"
"fmt"
"github.com/bytedance/sonic"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/compose"
"github.com/cloudwego/eino/schema"
)
var (
defaultAgentToolParam = schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
"request": {
Desc: "request to be processed",
Required: true,
Type: schema.String,
},
})
)
type AgentToolOptions struct {
fullChatHistoryAsInput bool
agentInputSchema *schema.ParamsOneOf
}
type AgentToolOption func(*AgentToolOptions)
// WithFullChatHistoryAsInput enables using the full chat history as input.
func WithFullChatHistoryAsInput() AgentToolOption {
return func(options *AgentToolOptions) {
options.fullChatHistoryAsInput = true
}
}
// WithAgentInputSchema sets a custom input schema for the agent tool.
func WithAgentInputSchema(schema *schema.ParamsOneOf) AgentToolOption {
return func(options *AgentToolOptions) {
options.agentInputSchema = schema
}
}
func withAgentToolEnableStreaming(enabled bool) tool.Option {
return tool.WrapImplSpecificOptFn(func(opt *agentToolOptions) {
opt.enableStreaming = enabled
})
}
// NewAgentTool creates a tool that wraps an agent for invocation.
//
// The agent must have a non-empty Name and Description, as they are used as
// the tool's name and description respectively. This is validated when Info()
// is called during tool setup.
//
// Event Streaming:
// When EmitInternalEvents is enabled in ToolsConfig, the agent tool will emit AgentEvent
// from the inner agent to the parent agent's AsyncGenerator, allowing real-time streaming
// of the inner agent's output to the end-user via Runner.
//
// Note that these forwarded events are NOT recorded in the parent agent's runSession.
// They are only emitted to the end-user and have no effect on the parent agent's state
// or checkpoint. The only exception is Interrupted action, which is propagated via
// CompositeInterrupt to enable proper interrupt/resume across agent boundaries.
//
// Action Scoping:
// Actions emitted by the inner agent are scoped to the agent tool boundary:
// - Interrupted: Propagated via CompositeInterrupt to allow proper interrupt/resume across boundaries
// - Exit, TransferToAgent, BreakLoop: Ignored outside the agent tool; these actions only affect
// the inner agent's execution and do not propagate to the parent agent
//
// This scoping ensures that nested agents cannot unexpectedly terminate or transfer control
// of their parent agent's execution flow.
func NewAgentTool(_ context.Context, agent Agent, options ...AgentToolOption) tool.BaseTool {
opts := &AgentToolOptions{}
for _, opt := range options {
opt(opts)
}
return &agentTool{
agent: agent,
fullChatHistoryAsInput: opts.fullChatHistoryAsInput,
inputSchema: opts.agentInputSchema,
}
}
// NewTypedAgentTool creates a new agent tool that wraps a TypedAgent as a tool.BaseTool.
func NewTypedAgentTool[M MessageType](_ context.Context, agent TypedAgent[M], options ...AgentToolOption) tool.BaseTool {
opts := &AgentToolOptions{}
for _, opt := range options {
opt(opts)
}
return &typedAgentTool[M]{
agent: agent,
fullChatHistoryAsInput: opts.fullChatHistoryAsInput,
inputSchema: opts.agentInputSchema,
}
}
type typedAgentTool[M MessageType] struct {
agent TypedAgent[M]
fullChatHistoryAsInput bool
inputSchema *schema.ParamsOneOf
}
type agentTool = typedAgentTool[*schema.Message]
type agentToolRequest struct {
Request string `json:"request"`
}
func (at *typedAgentTool[M]) Info(ctx context.Context) (*schema.ToolInfo, error) {
name := at.agent.Name(ctx)
if name == "" {
return nil, errors.New("agent tool requires a non-empty Name")
}
desc := at.agent.Description(ctx)
if desc == "" {
return nil, errors.New("agent tool requires a non-empty Description")
}
param := at.inputSchema
if param == nil {
param = defaultAgentToolParam
}
return &schema.ToolInfo{
Name: name,
Desc: desc,
ParamsOneOf: param,
}, nil
}
func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
if cancelCtx := getCancelContext(ctx); cancelCtx != nil {
cancelCtx.markAgentToolDescendant()
}
gen, enableStreaming := getEmitGeneratorAndEnableStreaming[M](opts)
var ms *bridgeStore
var iter *AsyncIterator[*TypedAgentEvent[M]]
var err error
wasInterrupted, hasState, state := tool.GetInterruptState[[]byte](ctx)
if !wasInterrupted {
ms = newBridgeStore()
var input []M
if at.fullChatHistoryAsInput {
var zero M
if _, ok := any(zero).(*schema.Message); !ok {
// fullChatHistoryAsInput is only supported for *schema.Message agents and will not
// be extended to *schema.AgenticMessage. The chat history format and role semantics
// differ fundamentally between Message and AgenticMessage, and the history rewriting
// logic (role attribution, system message filtering, transfer messages) is specific
// to the Message model.
return "", fmt.Errorf("fullChatHistoryAsInput is only supported for *schema.Message agents")
}
msgInput, histErr := getReactChatHistory(ctx, at.agent.Name(ctx))
if histErr != nil {
return "", histErr
}
input = any(msgInput).([]M)
} else {
if at.inputSchema == nil {
req := &agentToolRequest{}
err = sonic.UnmarshalString(argumentsInJSON, req)
if err != nil {
return "", err
}
argumentsInJSON = req.Request
}
input = newTypedUserMessages[M](argumentsInJSON)
}
runner := newTypedInvokableAgentToolRunner(at.agent, ms, enableStreaming)
iter = runner.Run(ctx, input,
append(extractAndDeriveAgentToolCancelCtx(ctx, at.agent.Name(ctx), opts), WithCheckPointID(bridgeCheckpointID), withSharedParentSession())...)
} else {
if !hasState {
return "", fmt.Errorf("agent tool '%s' interrupt has happened, but cannot find interrupt state", at.agent.Name(ctx))
}
ms = newResumeBridgeStore(bridgeCheckpointID, state)
agentOpts := extractAndDeriveAgentToolCancelCtx(ctx, at.agent.Name(ctx), opts)
agentOpts = append(agentOpts, withSharedParentSession())
runner := newTypedInvokableAgentToolRunner(at.agent, ms, enableStreaming)
iter, err = runner.Resume(ctx, bridgeCheckpointID, agentOpts...)
if err != nil {
return "", err
}
}
var lastEvent *TypedAgentEvent[M]
for {
event, ok := iter.Next()
if !ok {
break
}
if lastEvent != nil &&
lastEvent.Output != nil &&
lastEvent.Output.MessageOutput != nil &&
lastEvent.Output.MessageOutput.MessageStream != nil {
lastEvent.Output.MessageOutput.MessageStream.Close()
}
if event.Err != nil {
return "", event.Err
}
if gen != nil {
if event.Action == nil || event.Action.Interrupted == nil {
if parentRunCtx := getRunCtx(ctx); parentRunCtx != nil && len(parentRunCtx.RunPath) > 0 {
rp := make([]RunStep, 0, len(parentRunCtx.RunPath)+len(event.RunPath))
rp = append(rp, parentRunCtx.RunPath...)
rp = append(rp, event.RunPath...)
event.RunPath = rp
}
tmp := copyTypedAgentEvent(event)
gen.Send(event)
event = tmp
}
}
lastEvent = event
}
if lastEvent != nil && lastEvent.Action != nil && lastEvent.Action.Interrupted != nil {
data, existed, err_ := ms.Get(ctx, bridgeCheckpointID)
if err_ != nil {
return "", fmt.Errorf("failed to get interrupt info: %w", err_)
}
if !existed {
return "", fmt.Errorf("interrupt has happened, but cannot find interrupt info")
}
return "", tool.CompositeInterrupt(ctx, "agent tool interrupt", data,
lastEvent.Action.internalInterrupted)
}
if lastEvent == nil {
return "", errors.New("no event returned")
}
var ret string
if lastEvent.Output != nil {
if output := lastEvent.Output.MessageOutput; output != nil {
msg, err := output.GetMessage()
if err != nil {
return "", err
}
ret = extractTextContent(msg)
}
}
return ret, nil
}
// agentToolOptions is a wrapper structure used to convert AgentRunOption slices to tool.Option.
// It stores the agent name and corresponding run options for tool-specific processing.
type agentToolOptions struct {
agentName string
opts []AgentRunOption
enableStreaming bool
}
// typedAgentToolEventOptions carries the parent runner's event generator for a
// specific message type. This keeps forwarded internal events type-compatible
// with the parent event stream.
type typedAgentToolEventOptions[M MessageType] struct {
generator *AsyncGenerator[*TypedAgentEvent[M]]
}
func withAgentToolOptions(agentName string, opts []AgentRunOption) tool.Option {
return tool.WrapImplSpecificOptFn(func(opt *agentToolOptions) {
opt.agentName = agentName
opt.opts = opts
})
}
func withAgentToolEventGenerator(gen *AsyncGenerator[*AgentEvent]) tool.Option {
return withTypedAgentToolEventGenerator(gen)
}
func withTypedAgentToolEventGenerator[M MessageType](gen *AsyncGenerator[*TypedAgentEvent[M]]) tool.Option {
return tool.WrapImplSpecificOptFn(func(o *typedAgentToolEventOptions[M]) {
o.generator = gen
})
}
func getOptionsByAgentName(agentName string, opts []tool.Option) []AgentRunOption {
var ret []AgentRunOption
for _, opt := range opts {
o := tool.GetImplSpecificOptions[agentToolOptions](nil, opt)
if o != nil && o.agentName == agentName {
ret = append(ret, o.opts...)
}
}
return ret
}
func extractAndDeriveAgentToolCancelCtx(ctx context.Context, agentName string, opts []tool.Option) []AgentRunOption {
agentOpts := getOptionsByAgentName(agentName, opts)
baseOpts := getCommonOptions(nil, agentOpts...)
parentCtx := baseOpts.cancelCtx
if parentCtx == nil {
parentCtx = getCancelContext(ctx)
}
if parentCtx != nil {
parentCtx.markAgentToolDescendant()
childCtx := parentCtx.deriveAgentToolCancelContext(ctx)
agentOpts = append(agentOpts, WrapImplSpecificOptFn(func(o *options) {
o.cancelCtx = childCtx
}))
}
return agentOpts
}
func getEmitGeneratorAndEnableStreaming[M MessageType](opts []tool.Option) (*AsyncGenerator[*TypedAgentEvent[M]], bool) {
o := tool.GetImplSpecificOptions[agentToolOptions](nil, opts...)
eventOptions := tool.GetImplSpecificOptions[typedAgentToolEventOptions[M]](nil, opts...)
if o == nil && eventOptions == nil {
return nil, false
}
var gen *AsyncGenerator[*TypedAgentEvent[M]]
if eventOptions != nil {
gen = eventOptions.generator
}
var enableStreaming bool
if o != nil {
enableStreaming = o.enableStreaming
}
return gen, enableStreaming
}
func getReactChatHistory(ctx context.Context, destAgentName string) ([]Message, error) {
var messages []Message
err := compose.ProcessState(ctx, func(ctx context.Context, st *State) error {
if len(st.Messages) == 0 {
return nil
}
messages = make([]Message, len(st.Messages)-1)
copy(messages, st.Messages[:len(st.Messages)-1])
return nil
})
if err != nil {
return nil, fmt.Errorf("failed to get chat history from state: %w", err)
}
var agentName string
if runCtx := getRunCtx(ctx); runCtx != nil && len(runCtx.RunPath) > 0 {
agentName = runCtx.RunPath[len(runCtx.RunPath)-1].agentName
}
a, t := GenTransferMessages(ctx, destAgentName)
messages = append(messages, a, t)
history := make([]Message, 0, len(messages))
for _, msg := range messages {
if msg.Role == schema.System {
continue
}
if msg.Role == schema.Assistant || msg.Role == schema.Tool {
msg = rewriteMessage(msg, agentName)
}
history = append(history, msg)
}
return history, nil
}
func newTypedUserMessages[M MessageType](text string) []M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any([]Message{schema.UserMessage(text)}).([]M)
case *schema.AgenticMessage:
return any([]*schema.AgenticMessage{schema.UserAgenticMessage(text)}).([]M)
default:
return nil
}
}
func newTypedInvokableAgentToolRunner[M MessageType](agent TypedAgent[M], store compose.CheckPointStore, enableStreaming bool) *TypedRunner[M] {
return &TypedRunner[M]{
a: agent,
enableStreaming: enableStreaming,
store: store,
}
}
File diff suppressed because it is too large Load Diff
+268
View File
@@ -0,0 +1,268 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/cloudwego/eino/callbacks"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/schema"
)
type agenticCallbackRecorder struct {
mu sync.Mutex
onStartCalled bool
onEndCalled bool
runInfo *callbacks.RunInfo
inputReceived *TypedAgentCallbackInput[*schema.AgenticMessage]
eventsReceived []*TypedAgentEvent[*schema.AgenticMessage]
eventsDone chan struct{}
closeOnce sync.Once
}
func (r *agenticCallbackRecorder) getOnStartCalled() bool {
r.mu.Lock()
defer r.mu.Unlock()
return r.onStartCalled
}
func (r *agenticCallbackRecorder) getOnEndCalled() bool {
r.mu.Lock()
defer r.mu.Unlock()
return r.onEndCalled
}
func (r *agenticCallbackRecorder) getEventsReceived() []*TypedAgentEvent[*schema.AgenticMessage] {
r.mu.Lock()
defer r.mu.Unlock()
result := make([]*TypedAgentEvent[*schema.AgenticMessage], len(r.eventsReceived))
copy(result, r.eventsReceived)
return result
}
func newAgenticRecordingHandler(recorder *agenticCallbackRecorder) callbacks.Handler {
recorder.eventsDone = make(chan struct{})
return callbacks.NewHandlerBuilder().
OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
if info.Component != ComponentOfAgenticAgent {
return ctx
}
recorder.mu.Lock()
defer recorder.mu.Unlock()
recorder.onStartCalled = true
recorder.runInfo = info
if agentInput := ConvTypedCallbackInput[*schema.AgenticMessage](input); agentInput != nil {
recorder.inputReceived = agentInput
}
return ctx
}).
OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
if info.Component != ComponentOfAgenticAgent {
return ctx
}
recorder.mu.Lock()
recorder.onEndCalled = true
recorder.runInfo = info
recorder.mu.Unlock()
if agentOutput := ConvTypedCallbackOutput[*schema.AgenticMessage](output); agentOutput != nil {
if agentOutput.Events != nil {
go func() {
defer recorder.closeOnce.Do(func() { close(recorder.eventsDone) })
for {
event, ok := agentOutput.Events.Next()
if !ok {
break
}
recorder.mu.Lock()
recorder.eventsReceived = append(recorder.eventsReceived, event)
recorder.mu.Unlock()
}
}()
return ctx
}
}
recorder.closeOnce.Do(func() { close(recorder.eventsDone) })
return ctx
}).
Build()
}
func TestAgenticCallback(t *testing.T) {
ctx := context.Background()
expectedContent := "This is the test response content"
m := &mockAgenticModel{
generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
return agenticMsg(expectedContent), nil
},
}
agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
Name: "TestChatAgent",
Description: "Test chat agent",
Instruction: "You are a test agent",
Model: m,
})
require.NoError(t, err)
recorder := &agenticCallbackRecorder{}
handler := newAgenticRecordingHandler(recorder)
var agentEvents []*TypedAgentEvent[*schema.AgenticMessage]
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent})
iter := runner.Query(ctx, "hello", WithCallbacks(handler))
for {
event, ok := iter.Next()
if !ok {
break
}
agentEvents = append(agentEvents, event)
}
<-recorder.eventsDone
assertAgenticEventRoleFields(t, agentEvents)
t.Run("OnStart_Invocation", func(t *testing.T) {
assert.True(t, recorder.getOnStartCalled(), "OnStart should be called")
require.NotNil(t, recorder.inputReceived, "Input should be received")
require.NotNil(t, recorder.inputReceived.Input, "AgentInput should be set")
assert.Len(t, recorder.inputReceived.Input.Messages, 1)
})
t.Run("OnEnd_Invocation", func(t *testing.T) {
assert.True(t, recorder.getOnEndCalled(), "OnEnd should be called")
assert.Len(t, recorder.getEventsReceived(), 1)
})
t.Run("RunInfo_Fields", func(t *testing.T) {
require.NotNil(t, recorder.runInfo)
assert.Equal(t, "TestChatAgent", recorder.runInfo.Name)
assert.Equal(t, ComponentOfAgenticAgent, recorder.runInfo.Component)
})
t.Run("Events_MatchAgentOutput", func(t *testing.T) {
require.NotEmpty(t, agentEvents, "Agent should emit events")
received := recorder.getEventsReceived()
require.NotEmpty(t, received, "Callback should receive events")
require.Len(t, received, 1, "Callback should receive exactly 1 event")
require.NotNil(t, received[0].Output)
require.NotNil(t, received[0].Output.MessageOutput)
require.NotNil(t, received[0].Output.MessageOutput.Message)
assert.Equal(t, expectedContent, agenticTextContent(received[0].Output.MessageOutput.Message))
})
}
func TestAgenticCallbackMultipleHandlers(t *testing.T) {
ctx := context.Background()
m := &mockAgenticModel{
generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
return agenticMsg("test response"), nil
},
}
agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
Name: "TestAgent",
Description: "Test agent",
Instruction: "You are a test agent",
Model: m,
})
require.NoError(t, err)
recorder1 := &agenticCallbackRecorder{}
recorder2 := &agenticCallbackRecorder{}
handler1 := newAgenticRecordingHandler(recorder1)
handler2 := newAgenticRecordingHandler(recorder2)
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent})
iter := runner.Query(ctx, "hello", WithCallbacks(handler1, handler2))
for {
_, ok := iter.Next()
if !ok {
break
}
}
<-recorder1.eventsDone
<-recorder2.eventsDone
assert.True(t, recorder1.getOnStartCalled(), "Handler1 OnStart should be called")
assert.True(t, recorder2.getOnStartCalled(), "Handler2 OnStart should be called")
assert.True(t, recorder1.getOnEndCalled(), "Handler1 OnEnd should be called")
assert.True(t, recorder2.getOnEndCalled(), "Handler2 OnEnd should be called")
assert.NotEmpty(t, recorder1.getEventsReceived(), "Handler1 should receive events")
assert.NotEmpty(t, recorder2.getEventsReceived(), "Handler2 should receive events")
}
func TestCoverage_WrapAgenticIterWithOnEnd(t *testing.T) {
ctx := context.Background()
var onEndCalled bool
handler := callbacks.NewHandlerBuilder().
OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
return ctx
}).
OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
if info.Component == ComponentOfAgenticAgent {
onEndCalled = true
}
return ctx
}).
Build()
ctx = initAgenticCallbacks(ctx, "test-agent", "ChatModel",
WithCallbacks(handler))
cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{
Input: &TypedAgentInput[*schema.AgenticMessage]{
Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")},
},
}
ctx = callbacks.OnStart(ctx, cbInput)
origIter, origGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer origGen.Close()
origGen.Send(&TypedAgentEvent[*schema.AgenticMessage]{
Output: &TypedAgentOutput[*schema.AgenticMessage]{
MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
Message: agenticMsg("done"),
},
},
})
}()
wrappedIter := wrapAgenticIterWithOnEnd(ctx, origIter)
for {
_, ok := wrappedIter.Next()
if !ok {
break
}
}
assert.True(t, onEndCalled, "OnEnd callback should have been called")
}
+665
View File
@@ -0,0 +1,665 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"encoding/json"
"sync/atomic"
"testing"
"time"
"github.com/eino-contrib/jsonschema"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/compose"
"github.com/cloudwego/eino/schema"
)
func agenticMsg(text string) *schema.AgenticMessage {
return &schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.AssistantGenText{Text: text}),
},
}
}
func agenticTextContent(msg *schema.AgenticMessage) string {
for _, b := range msg.ContentBlocks {
if b.AssistantGenText != nil {
return b.AssistantGenText.Text
}
}
return ""
}
func TestAgenticIntegration_ChatModelSingleShot(t *testing.T) {
ctx := context.Background()
m := &mockAgenticModel{
generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
return agenticMsg("Handled internally with tool result: 42"), nil
},
}
dummyTool := newSlowTool("calculator", 0, "42")
agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
Name: "ToolCallAgent",
Description: "Agent with tools for agentic model",
Instruction: "You are a calculator.",
Model: m,
ToolsConfig: ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{dummyTool},
},
},
})
require.NoError(t, err)
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: agent,
})
iter := runner.Query(ctx, "What is 6*7?")
var events []*TypedAgentEvent[*schema.AgenticMessage]
for {
event, ok := iter.Next()
if !ok {
break
}
events = append(events, event)
}
require.Len(t, events, 1)
assertAgenticEventRoleFields(t, events)
lastEvent := events[len(events)-1]
require.Nil(t, lastEvent.Err)
require.NotNil(t, lastEvent.Output)
require.NotNil(t, lastEvent.Output.MessageOutput)
assert.Equal(t, "Handled internally with tool result: 42",
agenticTextContent(lastEvent.Output.MessageOutput.Message))
}
func TestAgenticIntegration_ChatModelToolsPassedViaOptions(t *testing.T) {
ctx := context.Background()
var receivedTools []*schema.ToolInfo
m := &mockAgenticModel{
generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
o := model.GetCommonOptions(&model.Options{}, opts...)
receivedTools = o.Tools
return agenticMsg("done"), nil
},
}
dummyTool := newSlowTool("my_tool", 0, "result")
agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
Name: "ToolOptAgent",
Description: "Agent verifying tools are passed via options",
Model: m,
ToolsConfig: ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{dummyTool},
},
},
})
require.NoError(t, err)
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: agent,
})
iter := runner.Query(ctx, "test tools")
for {
_, ok := iter.Next()
if !ok {
break
}
}
require.NotNil(t, receivedTools, "tools should be passed via model.Options")
require.Len(t, receivedTools, 1)
assert.Equal(t, "my_tool", receivedTools[0].Name)
}
func TestAgenticIntegration_StreamingWithRunner(t *testing.T) {
ctx := context.Background()
chunk1 := &schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello "}),
},
}
chunk2 := &schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}),
},
}
m := &mockAgenticModel{
streamFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) {
r, w := schema.Pipe[*schema.AgenticMessage](2)
go func() {
defer w.Close()
w.Send(chunk1, nil)
w.Send(chunk2, nil)
}()
return r, nil
},
}
agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
Name: "StreamRunner",
Description: "Streaming runner agent",
Model: m,
})
require.NoError(t, err)
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: agent,
EnableStreaming: true,
})
iter := runner.Query(ctx, "stream me")
event, ok := iter.Next()
require.True(t, ok)
assert.Nil(t, event.Err)
require.NotNil(t, event.Output)
require.NotNil(t, event.Output.MessageOutput)
if event.Output.MessageOutput.IsStreaming {
require.NotNil(t, event.Output.MessageOutput.MessageStream)
var chunks []*schema.AgenticMessage
for {
chunk, err := event.Output.MessageOutput.MessageStream.Recv()
if err != nil {
break
}
chunks = append(chunks, chunk)
}
assert.Equal(t, 2, len(chunks))
} else {
assert.NotNil(t, event.Output.MessageOutput.Message)
}
_, ok = iter.Next()
assert.False(t, ok)
}
func TestAgenticIntegration_CancelDuringExecution(t *testing.T) {
ctx := context.Background()
modelStarted := make(chan struct{}, 1)
modelBlocked := make(chan struct{})
m := &mockAgenticModel{
generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
select {
case modelStarted <- struct{}{}:
default:
}
select {
case <-modelBlocked:
return agenticMsg("should not reach"), nil
case <-ctx.Done():
return nil, ctx.Err()
}
},
}
agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
Name: "CancelAgent",
Description: "cancel test",
Model: m,
})
require.NoError(t, err)
cancelCtx, cancel := context.WithCancel(ctx)
defer cancel()
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: agent,
})
iter := runner.Run(cancelCtx, []*schema.AgenticMessage{
schema.UserAgenticMessage("Hi"),
})
<-modelStarted
cancel()
var capturedErr error
for {
event, ok := iter.Next()
if !ok {
break
}
if event.Err != nil {
capturedErr = event.Err
}
}
require.Error(t, capturedErr, "should propagate cancel error")
assert.ErrorIs(t, capturedErr, context.Canceled)
}
func TestAgenticIntegration_CancelWithTimeout(t *testing.T) {
ctx := context.Background()
sa := &myAgenticAgent{
name: "slow-agent",
runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer generator.Close()
select {
case <-time.After(10 * time.Second):
generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{
Output: &TypedAgentOutput[*schema.AgenticMessage]{
MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
Message: agenticMsg("slow response"),
},
},
})
case <-ctx.Done():
generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{
Err: ctx.Err(),
})
}
}()
return iter
},
}
timeoutCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
defer cancel()
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: sa,
})
iter := runner.Run(timeoutCtx, []*schema.AgenticMessage{
schema.UserAgenticMessage("slow request"),
})
var capturedErr error
for {
event, ok := iter.Next()
if !ok {
break
}
if event.Err != nil {
capturedErr = event.Err
}
}
require.Error(t, capturedErr, "should get timeout/cancel error")
assert.ErrorIs(t, capturedErr, context.DeadlineExceeded)
}
func TestAgenticIntegration_AgentTool(t *testing.T) {
ctx := context.Background()
innerModel := &mockAgenticModel{
generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
return agenticMsg("inner tool result"), nil
},
}
innerAgent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
Name: "InnerAgent",
Description: "An agent used as a tool",
Model: innerModel,
})
require.NoError(t, err)
agentTool := NewTypedAgentTool(ctx, TypedAgent[*schema.AgenticMessage](innerAgent))
require.NotNil(t, agentTool)
info, err := agentTool.Info(ctx)
require.NoError(t, err)
assert.Equal(t, "InnerAgent", info.Name)
assert.Equal(t, "An agent used as a tool", info.Desc)
outerModel := &mockAgenticModel{
generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
return agenticMsg("outer response after inner tool"), nil
},
}
outerAgent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
Name: "OuterAgent",
Description: "Outer agent with agent tool",
Model: outerModel,
ToolsConfig: ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: []tool.BaseTool{agentTool},
},
},
})
require.NoError(t, err)
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: outerAgent,
})
iter := runner.Query(ctx, "delegate to inner")
var events []*TypedAgentEvent[*schema.AgenticMessage]
for {
event, ok := iter.Next()
if !ok {
break
}
events = append(events, event)
}
require.NotEmpty(t, events)
assertAgenticEventRoleFields(t, events)
lastEvent := events[len(events)-1]
assert.Nil(t, lastEvent.Err)
assert.NotNil(t, lastEvent.Output)
}
func TestAgenticIntegration_InterruptEventFormation(t *testing.T) {
ctx := context.Background()
t.Run("simple interrupt", func(t *testing.T) {
agent := &myAgenticAgent{
name: "int-agent",
runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer generator.Close()
intEvent := TypedInterrupt[*schema.AgenticMessage](ctx, "approval needed")
intEvent.Action.Interrupted.Data = "approval data"
generator.Send(intEvent)
}()
return iter
},
}
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: agent,
})
iter := runner.Query(ctx, "interrupt test")
var interruptEvent *TypedAgentEvent[*schema.AgenticMessage]
for {
event, ok := iter.Next()
if !ok {
break
}
if event.Action != nil && event.Action.Interrupted != nil {
interruptEvent = event
}
}
require.NotNil(t, interruptEvent)
assert.Equal(t, "approval data", interruptEvent.Action.Interrupted.Data)
require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)
assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID)
assert.Equal(t, "approval needed", interruptEvent.Action.Interrupted.InterruptContexts[0].Info)
assert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause)
})
t.Run("stateful interrupt", func(t *testing.T) {
agent := &myAgenticAgent{
name: "st-agent",
runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer generator.Close()
intEvent := TypedStatefulInterrupt[*schema.AgenticMessage](ctx, "state interrupt", "my-state")
intEvent.Action.Interrupted.Data = "stateful data"
generator.Send(intEvent)
}()
return iter
},
}
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: agent,
})
iter := runner.Query(ctx, "stateful test")
var interruptEvent *TypedAgentEvent[*schema.AgenticMessage]
for {
event, ok := iter.Next()
if !ok {
break
}
if event.Action != nil && event.Action.Interrupted != nil {
interruptEvent = event
}
}
require.NotNil(t, interruptEvent)
assert.Equal(t, "stateful data", interruptEvent.Action.Interrupted.Data)
require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)
assert.Equal(t, "state interrupt", interruptEvent.Action.Interrupted.InterruptContexts[0].Info)
})
}
func TestAgenticIntegration_CheckpointInterruptResume(t *testing.T) {
ctx := context.Background()
var resumeCalled int32
agent := &myAgenticAgent{
name: "ckpt-agent",
runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer generator.Close()
generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{
AgentName: "ckpt-agent",
Output: &TypedAgentOutput[*schema.AgenticMessage]{
MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
Message: agenticMsg("before interrupt"),
},
},
})
intEvent := TypedInterrupt[*schema.AgenticMessage](ctx, "need approval")
intEvent.Action.Interrupted.Data = "approval data"
generator.Send(intEvent)
}()
return iter
},
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
atomic.StoreInt32(&resumeCalled, 1)
iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer generator.Close()
generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{
AgentName: "ckpt-agent",
Output: &TypedAgentOutput[*schema.AgenticMessage]{
MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
Message: agenticMsg("after resume"),
},
},
})
}()
return iter
},
}
store := newMyStore()
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: agent,
CheckPointStore: store,
})
iter := runner.Query(ctx, "checkpoint test", WithCheckPointID("ckpt-1"))
var interruptEvent *TypedAgentEvent[*schema.AgenticMessage]
var preInterruptOutputs []string
for {
event, ok := iter.Next()
if !ok {
break
}
require.Nil(t, event.Err)
if event.Action != nil && event.Action.Interrupted != nil {
interruptEvent = event
}
if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil {
preInterruptOutputs = append(preInterruptOutputs, agenticTextContent(event.Output.MessageOutput.Message))
}
}
require.NotNil(t, interruptEvent, "should receive interrupt event")
assert.Contains(t, preInterruptOutputs, "before interrupt")
require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)
interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID
require.NotEmpty(t, interruptID)
resumeIter, err := runner.ResumeWithParams(ctx, "ckpt-1", &ResumeParams{
Targets: map[string]any{
interruptID: nil,
},
})
require.NoError(t, err)
var postResumeOutputs []string
for {
event, ok := resumeIter.Next()
if !ok {
break
}
if event.Err != nil {
t.Fatalf("unexpected error during resume: %v", event.Err)
}
if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil {
postResumeOutputs = append(postResumeOutputs, agenticTextContent(event.Output.MessageOutput.Message))
}
}
assert.Equal(t, int32(1), atomic.LoadInt32(&resumeCalled), "resume function should have been called")
assert.Contains(t, postResumeOutputs, "after resume")
}
func TestAgenticIntegration_CheckpointWithMCPListToolsResult(t *testing.T) {
ctx := context.Background()
inputSchemaJSON := `{
"type": "object",
"properties": {
"query": {"type": "string", "description": "search query"},
"limit": {"type": "integer", "description": "max results"}
},
"required": ["query"]
}`
var inputSchema jsonschema.Schema
require.NoError(t, json.Unmarshal([]byte(inputSchemaJSON), &inputSchema))
mcpMsg := &schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
{
Type: schema.ContentBlockTypeMCPListToolsResult,
MCPListToolsResult: &schema.MCPListToolsResult{
ServerLabel: "test-server",
Tools: []*schema.MCPListToolsItem{
{
Name: "search",
Description: "search the web",
InputSchema: &inputSchema,
},
},
},
},
schema.NewContentBlock(&schema.AssistantGenText{Text: "here are tools"}),
},
}
var resumeCalled int32
agent := &myAgenticAgent{
name: "mcp-agent",
runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer gen.Close()
gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{
AgentName: "mcp-agent",
Output: &TypedAgentOutput[*schema.AgenticMessage]{
MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: mcpMsg},
},
})
gen.Send(TypedInterrupt[*schema.AgenticMessage](ctx, "approve tools"))
}()
return iter
},
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
atomic.StoreInt32(&resumeCalled, 1)
iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer gen.Close()
gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{
AgentName: "mcp-agent",
Output: &TypedAgentOutput[*schema.AgenticMessage]{
MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("tools approved")},
},
})
}()
return iter
},
}
store := newMyStore()
runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{
Agent: agent,
CheckPointStore: store,
})
iter := runner.Query(ctx, "list tools", WithCheckPointID("mcp-1"))
var interruptEvent *TypedAgentEvent[*schema.AgenticMessage]
for {
ev, ok := iter.Next()
if !ok {
break
}
require.Nil(t, ev.Err)
if ev.Action != nil && ev.Action.Interrupted != nil {
interruptEvent = ev
}
}
require.NotNil(t, interruptEvent)
interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID
resumeIter, err := runner.ResumeWithParams(ctx, "mcp-1", &ResumeParams{
Targets: map[string]any{interruptID: nil},
})
require.NoError(t, err)
var outputs []string
for {
ev, ok := resumeIter.Next()
if !ok {
break
}
require.Nil(t, ev.Err)
if ev.Output != nil && ev.Output.MessageOutput != nil && ev.Output.MessageOutput.Message != nil {
outputs = append(outputs, agenticTextContent(ev.Output.MessageOutput.Message))
}
}
assert.Equal(t, int32(1), atomic.LoadInt32(&resumeCalled))
assert.Contains(t, outputs, "tools approved")
}
File diff suppressed because it is too large Load Diff
+1745
View File
File diff suppressed because it is too large Load Diff
+211
View File
@@ -0,0 +1,211 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import "github.com/cloudwego/eino/callbacks"
type options struct {
sharedParentSession bool
sessionValues map[string]any
checkPointID *string
skipTransferMessages bool
handlers []callbacks.Handler
cancelCtx *cancelContext
}
// AgentRunOption is the call option for adk Agent.
type AgentRunOption struct {
implSpecificOptFn any
// specify which Agent can see this AgentRunOption, if empty, all Agents can see this AgentRunOption
agentNames []string
}
func (o AgentRunOption) DesignateAgent(name ...string) AgentRunOption {
o.agentNames = append(o.agentNames, name...)
return o
}
func getCommonOptions(base *options, opts ...AgentRunOption) *options {
if base == nil {
base = &options{}
}
return GetImplSpecificOptions(base, opts...)
}
// WithSessionValues sets session-scoped values for the agent run.
func WithSessionValues(v map[string]any) AgentRunOption {
return WrapImplSpecificOptFn(func(o *options) {
o.sessionValues = v
})
}
// WithSkipTransferMessages disables forwarding transfer messages during execution.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func WithSkipTransferMessages() AgentRunOption {
return WrapImplSpecificOptFn(func(t *options) {
t.skipTransferMessages = true
})
}
func withSharedParentSession() AgentRunOption {
return WrapImplSpecificOptFn(func(o *options) {
o.sharedParentSession = true
})
}
// WithCallbacks adds callback handlers to receive agent lifecycle events.
// Handlers receive OnStart with AgentCallbackInput and OnEnd with AgentCallbackOutput.
// Multiple handlers can be added; each receives an independent copy of the event stream.
func WithCallbacks(handlers ...callbacks.Handler) AgentRunOption {
return WrapImplSpecificOptFn(func(o *options) {
o.handlers = append(o.handlers, handlers...)
})
}
// WrapImplSpecificOptFn is the option to wrap the implementation specific option function.
func WrapImplSpecificOptFn[T any](optFn func(*T)) AgentRunOption {
return AgentRunOption{
implSpecificOptFn: optFn,
}
}
// GetImplSpecificOptions extract the implementation specific options from AgentRunOption list, optionally providing a base options with default values.
// e.g.
//
// myOption := &MyOption{
// Field1: "default_value",
// }
//
// myOption := model.GetImplSpecificOptions(myOption, opts...)
func GetImplSpecificOptions[T any](base *T, opts ...AgentRunOption) *T {
if base == nil {
base = new(T)
}
for i := range opts {
opt := opts[i]
if opt.implSpecificOptFn != nil {
optFn, ok := opt.implSpecificOptFn.(func(*T))
if ok {
optFn(base)
}
}
}
return base
}
// filterCallbackHandlersForNestedAgents removes callback handlers that have already been applied
// to the current agent before passing opts to nested inner agents.
//
// This is necessary for workflow agents (LoopAgent, SequentialAgent, ParallelAgent) because:
// 1. Callback handlers designated for the current agent are applied via initAgentCallbacks(),
// which stores them in the context.
// 2. Nested inner agents inherit this context, so they automatically receive these callbacks.
// 3. If we also pass these handlers in opts to inner agents, they would be applied twice,
// causing duplicate callback invocations.
//
// Note: This only applies to workflow agents where inner agents inherit context from the parent.
// For flowAgent's sub-agents (which are peer agents that transfer to each other), the full opts
// are passed since they don't inherit the parent's callback context.
func filterCallbackHandlersForNestedAgents(currentAgentName string, opts []AgentRunOption) []AgentRunOption {
if len(opts) == 0 {
return nil
}
var filteredOpts []AgentRunOption
for i := range opts {
opt := opts[i]
if opt.implSpecificOptFn == nil {
filteredOpts = append(filteredOpts, opt)
continue
}
if _, isCallbackOpt := opt.implSpecificOptFn.(func(*options)); isCallbackOpt {
testOpt := &options{}
opt.implSpecificOptFn.(func(*options))(testOpt)
if len(testOpt.handlers) > 0 {
if len(opt.agentNames) == 0 {
continue
}
matched := false
for _, name := range opt.agentNames {
if name == currentAgentName {
matched = true
break
}
}
if matched {
continue
}
}
}
filteredOpts = append(filteredOpts, opt)
}
return filteredOpts
}
// filterCancelOption removes any AgentRunOption that sets a cancelCtx on *options.
// This prevents inner (nested) agents from receiving the cancel option when the
// outer flowAgent owns the cancel lifecycle. Inner agents access the cancelContext
// via the Go context (getCancelContext) instead.
func filterCancelOption(opts []AgentRunOption) []AgentRunOption {
if len(opts) == 0 {
return nil
}
var filteredOpts []AgentRunOption
for i := range opts {
opt := opts[i]
if opt.implSpecificOptFn == nil {
filteredOpts = append(filteredOpts, opt)
continue
}
if _, isCommonOpt := opt.implSpecificOptFn.(func(*options)); isCommonOpt {
testOpt := &options{}
opt.implSpecificOptFn.(func(*options))(testOpt)
if testOpt.cancelCtx != nil {
continue
}
}
filteredOpts = append(filteredOpts, opt)
}
return filteredOpts
}
func filterOptions(agentName string, opts []AgentRunOption) []AgentRunOption {
if len(opts) == 0 {
return nil
}
var filteredOpts []AgentRunOption
for i := range opts {
opt := opts[i]
if len(opt.agentNames) == 0 {
filteredOpts = append(filteredOpts, opt)
continue
}
for j := range opt.agentNames {
if opt.agentNames[j] == agentName {
filteredOpts = append(filteredOpts, opt)
break
}
}
}
return filteredOpts
}
+202
View File
@@ -0,0 +1,202 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"github.com/cloudwego/eino/callbacks"
"github.com/cloudwego/eino/components"
icb "github.com/cloudwego/eino/internal/callbacks"
)
// AgentCallbackInput represents the input passed to agent callbacks during OnStart.
// Use ConvAgentCallbackInput to safely convert from callbacks.CallbackInput.
type AgentCallbackInput struct {
// Input contains the agent input for a new run. Nil when resuming.
Input *AgentInput
// ResumeInfo contains resume information when resuming from an interrupt. Nil for new runs.
ResumeInfo *ResumeInfo
}
// AgentCallbackOutput represents the output passed to agent callbacks during OnEnd.
// Use ConvAgentCallbackOutput to safely convert from callbacks.CallbackOutput.
//
// Important: The Events iterator should be consumed asynchronously to avoid blocking
// the agent execution. Each callback handler receives an independent copy of the iterator.
type AgentCallbackOutput struct {
// Events provides a stream of agent events. Each handler receives its own copy.
Events *AsyncIterator[*AgentEvent]
}
func copyTypedEventIterator[M MessageType](iter *AsyncIterator[*TypedAgentEvent[M]], n int) []*AsyncIterator[*TypedAgentEvent[M]] {
if n <= 0 {
return nil
}
if n == 1 {
return []*AsyncIterator[*TypedAgentEvent[M]]{iter}
}
iterators := make([]*AsyncIterator[*TypedAgentEvent[M]], n)
generators := make([]*AsyncGenerator[*TypedAgentEvent[M]], n)
for i := 0; i < n; i++ {
iterators[i], generators[i] = NewAsyncIteratorPair[*TypedAgentEvent[M]]()
}
go func() {
defer func() {
for _, g := range generators {
g.Close()
}
}()
for {
event, ok := iter.Next()
if !ok {
break
}
for i := 0; i < n-1; i++ {
generators[i].Send(copyTypedAgentEvent(event))
}
generators[n-1].Send(event)
}
}()
return iterators
}
func copyAgentCallbackOutput(out *AgentCallbackOutput, n int) []*AgentCallbackOutput {
if out == nil || out.Events == nil {
result := make([]*AgentCallbackOutput, n)
for i := 0; i < n; i++ {
result[i] = out
}
return result
}
iters := copyTypedEventIterator(out.Events, n)
result := make([]*AgentCallbackOutput, n)
for i, iter := range iters {
result[i] = &AgentCallbackOutput{Events: iter}
}
return result
}
// ConvAgentCallbackInput converts a generic CallbackInput to AgentCallbackInput.
// Returns nil if the input is not an AgentCallbackInput.
func ConvAgentCallbackInput(input callbacks.CallbackInput) *AgentCallbackInput {
if v, ok := input.(*AgentCallbackInput); ok {
return v
}
return nil
}
// ConvAgentCallbackOutput converts a generic CallbackOutput to AgentCallbackOutput.
// Returns nil if the output is not an AgentCallbackOutput.
func ConvAgentCallbackOutput(output callbacks.CallbackOutput) *AgentCallbackOutput {
if v, ok := output.(*AgentCallbackOutput); ok {
return v
}
return nil
}
func initAgentCallbacks(ctx context.Context, agentName, agentType string, opts ...AgentRunOption) context.Context {
ri := &callbacks.RunInfo{
Name: agentName,
Type: agentType,
Component: ComponentOfAgent,
}
o := getCommonOptions(nil, opts...)
if len(o.handlers) == 0 {
return icb.ReuseHandlers(ctx, ri)
}
return icb.AppendHandlers(ctx, ri, o.handlers...)
}
func getAgentType(agent Agent) string {
if typer, ok := agent.(components.Typer); ok {
return typer.GetType()
}
return ""
}
// TypedAgentCallbackInput represents the input passed to typed agent callbacks during OnStart.
// Use ConvTypedCallbackInput to safely convert from callbacks.CallbackInput.
type TypedAgentCallbackInput[M MessageType] struct {
// Input contains the agent input for a new run. Nil when resuming.
Input *TypedAgentInput[M]
// ResumeInfo contains resume information when resuming from an interrupt. Nil for new runs.
ResumeInfo *ResumeInfo
}
// TypedAgentCallbackOutput represents the output passed to typed agent callbacks during OnEnd.
// Use ConvTypedCallbackOutput to safely convert from callbacks.CallbackOutput.
//
// Important: The Events iterator should be consumed asynchronously to avoid blocking
// the agent execution. Each callback handler receives an independent copy of the iterator.
type TypedAgentCallbackOutput[M MessageType] struct {
// Events provides a stream of agent events. Each handler receives its own copy.
Events *AsyncIterator[*TypedAgentEvent[M]]
}
// ConvTypedCallbackInput converts a callbacks.CallbackInput to *TypedAgentCallbackInput[M].
// Returns nil if the input is not of the expected type.
func ConvTypedCallbackInput[M MessageType](input callbacks.CallbackInput) *TypedAgentCallbackInput[M] {
if v, ok := input.(*TypedAgentCallbackInput[M]); ok {
return v
}
return nil
}
// ConvTypedCallbackOutput converts a callbacks.CallbackOutput to *TypedAgentCallbackOutput[M].
// Returns nil if the output is not of the expected type.
func ConvTypedCallbackOutput[M MessageType](output callbacks.CallbackOutput) *TypedAgentCallbackOutput[M] {
if v, ok := output.(*TypedAgentCallbackOutput[M]); ok {
return v
}
return nil
}
func copyTypedCallbackOutput[M MessageType](out *TypedAgentCallbackOutput[M], n int) []*TypedAgentCallbackOutput[M] {
if out == nil || out.Events == nil {
result := make([]*TypedAgentCallbackOutput[M], n)
for i := 0; i < n; i++ {
result[i] = out
}
return result
}
iters := copyTypedEventIterator(out.Events, n)
result := make([]*TypedAgentCallbackOutput[M], n)
for i, iter := range iters {
result[i] = &TypedAgentCallbackOutput[M]{Events: iter}
}
return result
}
func initAgenticCallbacks(ctx context.Context, agentName, agentType string, opts ...AgentRunOption) context.Context {
ri := &callbacks.RunInfo{
Name: agentName,
Type: agentType,
Component: ComponentOfAgenticAgent,
}
o := getCommonOptions(nil, opts...)
if len(o.handlers) == 0 {
return icb.ReuseHandlers(ctx, ri)
}
return icb.AppendHandlers(ctx, ri, o.handlers...)
}
File diff suppressed because it is too large Load Diff
+388
View File
@@ -0,0 +1,388 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/cloudwego/eino/callbacks"
"github.com/cloudwego/eino/schema"
)
func TestCopyTypedEventIterator(t *testing.T) {
t.Run("n=0 returns nil", func(t *testing.T) {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
gen.Send(&AgentEvent{AgentName: "test"})
gen.Close()
}()
result := copyTypedEventIterator(iter, 0)
assert.Nil(t, result)
})
t.Run("n=1 returns original iterator", func(t *testing.T) {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
gen.Send(&AgentEvent{AgentName: "test"})
gen.Close()
}()
result := copyTypedEventIterator(iter, 1)
assert.Len(t, result, 1)
assert.Equal(t, iter, result[0])
})
t.Run("n>1 creates n independent copies", func(t *testing.T) {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
events := []*AgentEvent{
{AgentName: "agent1", Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.AssistantMessage("msg1", nil)}}},
{AgentName: "agent2", Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.AssistantMessage("msg2", nil)}}},
}
go func() {
for _, e := range events {
gen.Send(e)
}
gen.Close()
}()
n := 3
copies := copyTypedEventIterator(iter, n)
assert.Len(t, copies, n)
var wg sync.WaitGroup
receivedEvents := make([][]*AgentEvent, n)
for i := 0; i < n; i++ {
wg.Add(1)
go func(idx int) {
defer wg.Done()
for {
event, ok := copies[idx].Next()
if !ok {
break
}
receivedEvents[idx] = append(receivedEvents[idx], event)
}
}(i)
}
wg.Wait()
for i := 0; i < n; i++ {
assert.Len(t, receivedEvents[i], len(events), "iterator %d should receive all events", i)
for j, e := range receivedEvents[i] {
assert.Equal(t, events[j].AgentName, e.AgentName)
}
}
})
}
func TestCopyAgentCallbackOutput(t *testing.T) {
t.Run("nil output", func(t *testing.T) {
result := copyAgentCallbackOutput(nil, 3)
assert.Len(t, result, 3)
for _, r := range result {
assert.Nil(t, r)
}
})
t.Run("output with nil Events", func(t *testing.T) {
out := &AgentCallbackOutput{Events: nil}
result := copyAgentCallbackOutput(out, 3)
assert.Len(t, result, 3)
for _, r := range result {
assert.Equal(t, out, r)
}
})
t.Run("valid output with events", func(t *testing.T) {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
gen.Send(&AgentEvent{AgentName: "test"})
gen.Close()
}()
out := &AgentCallbackOutput{Events: iter}
result := copyAgentCallbackOutput(out, 2)
assert.Len(t, result, 2)
for i, r := range result {
require.NotNil(t, r, "result[%d] should not be nil", i)
assert.NotNil(t, r.Events, "result[%d].Events should not be nil", i)
}
})
}
func TestConvAgentCallbackInput(t *testing.T) {
t.Run("valid AgentCallbackInput", func(t *testing.T) {
input := &AgentCallbackInput{
Input: &AgentInput{Messages: []Message{schema.UserMessage("test")}},
}
result := ConvAgentCallbackInput(input)
assert.Equal(t, input, result)
})
t.Run("invalid type returns nil", func(t *testing.T) {
result := ConvAgentCallbackInput("invalid")
assert.Nil(t, result)
})
t.Run("nil returns nil", func(t *testing.T) {
result := ConvAgentCallbackInput(nil)
assert.Nil(t, result)
})
}
func TestConvAgentCallbackOutput(t *testing.T) {
t.Run("valid AgentCallbackOutput", func(t *testing.T) {
iter, _ := NewAsyncIteratorPair[*AgentEvent]()
output := &AgentCallbackOutput{Events: iter}
result := ConvAgentCallbackOutput(output)
assert.Equal(t, output, result)
})
t.Run("invalid type returns nil", func(t *testing.T) {
result := ConvAgentCallbackOutput("invalid")
assert.Nil(t, result)
})
t.Run("nil returns nil", func(t *testing.T) {
result := ConvAgentCallbackOutput(nil)
assert.Nil(t, result)
})
}
type mockTyperAgent struct {
name string
agentType string
}
func (a *mockTyperAgent) Name(_ context.Context) string { return a.name }
func (a *mockTyperAgent) Description(_ context.Context) string { return "mock agent" }
func (a *mockTyperAgent) GetType() string { return a.agentType }
func (a *mockTyperAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
gen.Close()
return iter
}
type mockNonTyperAgent struct {
name string
}
func (a *mockNonTyperAgent) Name(_ context.Context) string { return a.name }
func (a *mockNonTyperAgent) Description(_ context.Context) string { return "mock agent" }
func (a *mockNonTyperAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
gen.Close()
return iter
}
func TestGetAgentType(t *testing.T) {
t.Run("agent implementing Typer", func(t *testing.T) {
agent := &mockTyperAgent{name: "test", agentType: "CustomType"}
result := getAgentType(agent)
assert.Equal(t, "CustomType", result)
})
t.Run("agent not implementing Typer", func(t *testing.T) {
agent := &mockNonTyperAgent{name: "test"}
result := getAgentType(agent)
assert.Equal(t, "", result)
})
}
func TestWithCallbacksOption(t *testing.T) {
handler := callbacks.NewHandlerBuilder().
OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
return ctx
}).
Build()
opt := WithCallbacks(handler)
opts := getCommonOptions(nil, opt)
assert.Len(t, opts.handlers, 1)
}
func TestWithMultipleCallbacksOption(t *testing.T) {
handler1 := callbacks.NewHandlerBuilder().Build()
handler2 := callbacks.NewHandlerBuilder().Build()
opt := WithCallbacks(handler1, handler2)
opts := getCommonOptions(nil, opt)
assert.Len(t, opts.handlers, 2)
}
func TestCopyTypedEventIteratorAgentic(t *testing.T) {
t.Run("n=0 returns nil", func(t *testing.T) {
iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"})
gen.Close()
}()
result := copyTypedEventIterator(iter, 0)
assert.Nil(t, result)
})
t.Run("n=1 returns original iterator", func(t *testing.T) {
iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"})
gen.Close()
}()
result := copyTypedEventIterator(iter, 1)
assert.Len(t, result, 1)
assert.Equal(t, iter, result[0])
})
t.Run("n>1 creates n independent copies", func(t *testing.T) {
iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
events := []*TypedAgentEvent[*schema.AgenticMessage]{
{AgentName: "agent1", Output: &TypedAgentOutput[*schema.AgenticMessage]{
MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("msg1")},
}},
{AgentName: "agent2", Output: &TypedAgentOutput[*schema.AgenticMessage]{
MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("msg2")},
}},
}
go func() {
for _, e := range events {
gen.Send(e)
}
gen.Close()
}()
n := 3
copies := copyTypedEventIterator(iter, n)
assert.Len(t, copies, n)
var wg sync.WaitGroup
receivedEvents := make([][]*TypedAgentEvent[*schema.AgenticMessage], n)
for i := 0; i < n; i++ {
wg.Add(1)
go func(idx int) {
defer wg.Done()
for {
event, ok := copies[idx].Next()
if !ok {
break
}
receivedEvents[idx] = append(receivedEvents[idx], event)
}
}(i)
}
wg.Wait()
for i := 0; i < n; i++ {
assert.Len(t, receivedEvents[i], len(events), "iterator %d should receive all events", i)
for j, e := range receivedEvents[i] {
assert.Equal(t, events[j].AgentName, e.AgentName)
}
}
})
}
func TestCopyTypedCallbackOutput(t *testing.T) {
t.Run("nil output", func(t *testing.T) {
result := copyTypedCallbackOutput[*schema.AgenticMessage](nil, 3)
assert.Len(t, result, 3)
for _, r := range result {
assert.Nil(t, r)
}
})
t.Run("output with nil Events", func(t *testing.T) {
out := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: nil}
result := copyTypedCallbackOutput(out, 3)
assert.Len(t, result, 3)
for _, r := range result {
assert.Equal(t, out, r)
}
})
t.Run("valid output with events", func(t *testing.T) {
iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"})
gen.Close()
}()
out := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: iter}
result := copyTypedCallbackOutput(out, 2)
assert.Len(t, result, 2)
for i, r := range result {
require.NotNil(t, r, "result[%d] should not be nil", i)
assert.NotNil(t, r.Events, "result[%d].Events should not be nil", i)
}
})
}
func TestConvTypedCallbackInput(t *testing.T) {
t.Run("valid TypedAgentCallbackInput", func(t *testing.T) {
input := &TypedAgentCallbackInput[*schema.AgenticMessage]{
Input: &TypedAgentInput[*schema.AgenticMessage]{
Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")},
},
}
result := ConvTypedCallbackInput[*schema.AgenticMessage](input)
assert.Equal(t, input, result)
})
t.Run("invalid type returns nil", func(t *testing.T) {
result := ConvTypedCallbackInput[*schema.AgenticMessage]("invalid")
assert.Nil(t, result)
})
t.Run("nil returns nil", func(t *testing.T) {
result := ConvTypedCallbackInput[*schema.AgenticMessage](nil)
assert.Nil(t, result)
})
}
func TestConvTypedCallbackOutput(t *testing.T) {
t.Run("valid TypedAgentCallbackOutput", func(t *testing.T) {
iter, _ := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
output := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: iter}
result := ConvTypedCallbackOutput[*schema.AgenticMessage](output)
assert.Equal(t, output, result)
})
t.Run("invalid type returns nil", func(t *testing.T) {
result := ConvTypedCallbackOutput[*schema.AgenticMessage]("invalid")
assert.Nil(t, result)
})
t.Run("nil returns nil", func(t *testing.T) {
result := ConvTypedCallbackOutput[*schema.AgenticMessage](nil)
assert.Nil(t, result)
})
}
+975
View File
@@ -0,0 +1,975 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"errors"
"fmt"
"io"
"sync"
"sync/atomic"
"time"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/compose"
"github.com/cloudwego/eino/schema"
)
func init() {
schema.RegisterName[*CancelError]("_eino_adk_cancel_error")
schema.RegisterName[*AgentCancelInfo]("_eino_adk_agent_cancel_info")
schema.RegisterName[*StreamCanceledError]("_eino_adk_stream_cancelled_error")
}
// CancelMode specifies when an agent should be canceled.
// Modes can be combined with bitwise OR to cancel at multiple safe-points.
// For example, CancelAfterChatModel | CancelAfterToolCalls cancels the agent
// after whichever safe-point is reached first.
type CancelMode int
const (
// CancelImmediate cancels the agent as soon as the signal is received,
// without waiting for a ChatModel or ToolCalls safe-point.
// By default, only the root agent is interrupted; descendant agents inside
// AgentTools are torn down via context cancellation as a side effect.
// Use WithRecursive to propagate explicit immediate-cancel signals to
// descendants for clean teardown with grace period.
CancelImmediate CancelMode = 0
// CancelAfterChatModel cancels after the root agent's next chat model call
// completes. By default, only the root agent checks this safe-point;
// nested sub-agents inside AgentTools are unaware of the cancel.
// Use WithRecursive to propagate the cancel to all descendants — whichever
// ChatModel finishes first triggers the cancel.
CancelAfterChatModel CancelMode = 1 << iota
// CancelAfterToolCalls cancels after the root agent's next set of tool calls
// completes. By default, only the root agent checks this safe-point.
// Use WithRecursive to propagate to all descendants.
CancelAfterToolCalls
)
// CancelHandle represents a cancel operation that can be waited on.
type CancelHandle struct {
wait func() error
}
// Wait blocks until the cancel request reaches a terminal outcome.
//
// It reports the result of the cancel operation itself, not the agent's final
// business error:
// - nil: cancellation succeeded, including the case where a business interrupt
// was absorbed into CancelError while cancellation was active
// - ErrCancelTimeout: the requested safe-point cancellation timed out and was
// escalated to immediate cancellation
// - ErrExecutionEnded: the execution ended before cancellation took effect,
// meaning the stream drained to completion without any interrupt
func (h *CancelHandle) Wait() error {
return h.wait()
}
// AgentCancelFunc is called to request cancellation of a running agent.
// It returns after the cancel request is committed; use the returned handle's
// Wait to block for completion and outcome.
//
// The returned bool reports whether this call contributed to the CancelError
// for the current execution. "Contributed" means this call's cancel options
// were included before cancellation was finalized. It is false when cancellation
// was already finalized (handled or execution completed).
type AgentCancelFunc func(...AgentCancelOption) (*CancelHandle, bool)
type agentCancelConfig struct {
Mode CancelMode
Recursive bool
Timeout *time.Duration
}
// AgentCancelOption configures cancel behavior.
type AgentCancelOption func(*agentCancelConfig)
// WithAgentCancelMode sets the cancel mode for the agent cancel operation.
func WithAgentCancelMode(mode CancelMode) AgentCancelOption {
return func(config *agentCancelConfig) {
config.Mode = mode
}
}
// WithAgentCancelTimeout sets a timeout for the cancel operation.
// This only applies to safe-point modes (CancelAfterChatModel, CancelAfterToolCalls):
// if the safe-point hasn't fired within this duration, the cancel escalates to
// CancelImmediate. The escalated cancel still saves a checkpoint, so the execution
// can be resumed via Runner.Resume or Runner.ResumeWithParams.
// For CancelImmediate this timeout is ignored — the cancel fires immediately.
func WithAgentCancelTimeout(timeout time.Duration) AgentCancelOption {
return func(config *agentCancelConfig) {
config.Timeout = &timeout
}
}
// WithRecursive opts into recursive cancel propagation. By default, cancel
// modes only affect the root agent; descendant agents inside AgentTools are
// not notified. WithRecursive makes the cancel propagate to all descendants:
// - CancelAfterChatModel / CancelAfterToolCalls: descendants check their own safe-points.
// - CancelImmediate: descendants receive explicit immediate-cancel signals for
// clean teardown; the root uses a grace period to collect child interrupts.
//
// With recursive cancellation, each descendant agent also triggers cancellation
// and cascades its interrupt information upward. The root agent ultimately
// produces a complete checkpoint that includes descendant checkpoints, enabling
// resumption from the exact point where each descendant was interrupted.
//
// Once any cancel call includes WithRecursive, the flag stays set for the
// entire cancel lifecycle (monotonic escalation).
func WithRecursive() AgentCancelOption {
return func(config *agentCancelConfig) {
config.Recursive = true
}
}
// AgentCancelInfo contains information about a cancel operation.
type AgentCancelInfo struct {
Mode CancelMode
Escalated bool
Timeout bool
}
// CancelError is sent via AgentEvent.Err when an agent is canceled.
// Use errors.As to match and extract *CancelError from event errors.
//
// Interrupt absorption: when a cancel is active (shouldCancel() == true), ANY
// interrupt — whether from a cancel safe-point node or from business logic
// (e.g. tool.Interrupt in a tool) — is converted to a CancelError. The
// cancel "absorbs" the business interrupt. This is intentional:
//
// - In concurrent execution (parallel workflows, concurrent tool calls),
// cancel-induced and business interrupts can arrive as a single composite
// signal that cannot be split apart.
// - Even in sequential execution, treating business interrupts as CancelError
// during active cancel gives consistent semantics.
// - The business interrupt is NOT lost — the checkpoint preserves the full
// interrupt hierarchy. On resume (Runner.Resume or Runner.ResumeWithParams),
// the agent re-executes the interrupting code path and the business
// interrupt re-fires naturally.
type CancelError struct {
Info *AgentCancelInfo
// InterruptContexts provides the interrupt contexts needed for targeted
// resumption via Runner.ResumeWithParams. Each context represents a step
// in the agent hierarchy that was interrupted. This is a slice because
// composite agents (e.g. parallel workflows) may interrupt at multiple
// points simultaneously, matching the shape of AgentAction.Interrupted.InterruptContexts.
// Use each InterruptCtx.ID as a key in ResumeParams.Targets.
InterruptContexts []*InterruptCtx
interruptSignal *InterruptSignal // unexported — only Runner needs it for checkpoint
}
func (e *CancelError) Error() string {
return fmt.Sprintf("agent canceled: mode=%v, escalated=%v", e.Info.Mode, e.Info.Escalated)
}
// Sentinel errors for cancel outcomes.
var (
// ErrCancelTimeout is returned by CancelHandle.Wait when the cancel operation timed out.
ErrCancelTimeout = errors.New("cancel timed out")
// ErrExecutionEnded is returned by CancelHandle.Wait when the agent ended
// before the cancel took effect. "Ended" means the event stream was fully
// drained without any interrupt — normal completion or a fatal error.
//
// Note: business interrupts that occur while cancel is active are absorbed
// into CancelError (see CancelError doc), so they result in nil (cancel
// succeeded), NOT ErrExecutionEnded. Only execution that completes with
// no interrupt at all produces this error.
ErrExecutionEnded = errors.New("execution already ended")
// ErrStreamCanceled is the error sent through the stream when CancelImmediate aborts it.
// It is a *StreamCanceledError so it can be gob-serialized during checkpoint save
// (when stored as agentEventWrapper.StreamErr).
ErrStreamCanceled error = &StreamCanceledError{}
)
// StreamCanceledError is the concrete error type for ErrStreamCanceled.
// It is exported so that gob can serialize it during checkpoint save when the error
// is stored in agentEventWrapper.StreamErr.
type StreamCanceledError struct{}
func (e *StreamCanceledError) Error() string {
return "stream canceled"
}
// WithCancel creates an AgentRunOption that enables cancellation for an agent run.
// It returns the option to pass to Run/Resume and a cancel function.
// Cancel options (mode, timeout) are passed to the returned AgentCancelFunc at call time.
func WithCancel() (AgentRunOption, AgentCancelFunc) {
cc := newCancelContext()
opt := WrapImplSpecificOptFn(func(o *options) {
o.cancelCtx = cc
})
cancelFn := cc.buildCancelFunc()
return opt, cancelFn
}
// cancelContext state constants (for int32 CAS).
//
// State transition rules:
//
// stateRunning -> stateCancelling (cancel requested by AgentCancelFunc)
// stateRunning -> stateDone (execution finished without interrupt)
// stateCancelling -> stateCancelHandled (ANY interrupt absorbed as CancelError)
// stateCancelling -> stateDone (execution finished without interrupt while cancel pending)
//
// Terminal states: stateDone, stateCancelHandled.
//
// Note: We intentionally do NOT distinguish between "completed" and "errored"
// terminal states. End-users get the actual outcome from AgentEvent.
// This simplification keeps the state machine minimal — only the cancel/non-cancel
// distinction matters for the AgentCancelFunc return value.
//
// Business interrupt handling: when cancel is active (stateCancelling) and any
// interrupt arrives — cancel-induced OR business — wrapIterWithCancelCtx absorbs
// it as a CancelError and transitions to stateCancelHandled. The business interrupt
// data is preserved in the checkpoint for re-emission on resume.
const (
// stateRunning is the initial state: agent is executing normally.
stateRunning int32 = 0
// stateCancelling means AgentCancelFunc has been called and cancelChan is
// closed, but the cancel has not yet been handled by the runFunc.
stateCancelling int32 = 1
// stateDone means execution has finished through any non-cancel path:
// normal completion, business interrupt, or error. The specific outcome
// is conveyed through AgentEvent, not through the cancel state machine.
stateDone int32 = 2
// stateCancelHandled means the cancel was processed by the runFunc and a
// CancelError was emitted through the event stream. This is the success
// terminal state for cancellation.
stateCancelHandled int32 = 5
)
// interruptSent constants (for int32 CAS).
//
// Transition rules:
//
// interruptNotSent -> interruptImmediate (CancelImmediate or escalation)
const (
// interruptNotSent means no compose graph interrupt has been sent.
interruptNotSent int32 = 0
// interruptImmediate means an immediate graph interrupt was sent with
// timeout=0, forcing the graph to stop as soon as possible.
interruptImmediate int32 = 1
)
// defaultCancelImmediateGracePeriod is the bounded time a recursive
// AgentTool cancel waits before forcing the current level's graph interrupt.
// This gives deeper AgentTool/internal-agent interrupts a chance to bubble up
// as CompositeInterrupts. If this proves insufficient for deeply nested
// structures or too slow for latency-sensitive use cases, consider making it
// configurable via an AgentCancelOption.
const defaultCancelImmediateGracePeriod = 1 * time.Second
type cancelContextKey struct{}
// withCancelContext stores a cancelContext in the Go context.
func withCancelContext(ctx context.Context, cc *cancelContext) context.Context {
if cc == nil {
return ctx
}
return context.WithValue(ctx, cancelContextKey{}, cc)
}
// getCancelContext retrieves the cancelContext from the Go context, or nil.
func getCancelContext(ctx context.Context) *cancelContext {
if v := ctx.Value(cancelContextKey{}); v != nil {
return v.(*cancelContext)
}
return nil
}
type cancelContext struct {
mode int32 // atomic, CancelMode
cancelChan chan struct{} // closed when cancel is requested (all modes, not just safe-point)
immediateChan chan struct{} // closed when an immediate graph interrupt fires
doneChan chan struct{} // closed when execution completes (by any mark* method)
doneOnce sync.Once // ensures doneChan is closed exactly once
state int32 // stateRunning, stateCancelling, stateDone, stateCancelHandled
interruptSent int32 // interruptNotSent, interruptImmediate
escalated int32 // 1 if escalated from safe-point to immediate
timeoutEscalated int32 // 1 if escalation was triggered by timeout
startedMode int32 // atomic, mode when state transitioned to cancelling
deadlineUnixNano int64 // atomic, 0 means no deadline
recursive int32 // atomic; 1 if cancel should propagate into AgentTool internal agents
recursiveChan chan struct{} // closed when recursive transitions from 0 to 1
root bool // true for the original cancelContext created by WithCancel(); false for AgentTool internal agents
parent *cancelContext // non-nil for AgentTool internal agents; used to propagate AgentTool boundary markers upward
agentToolDescendant int32 // atomic; 1 once an AgentTool runs under this cancel context
cancelMu sync.Mutex
timeoutOnce sync.Once
timeoutNotify chan struct{}
mu sync.Mutex
graphInterruptFuncs []func(...compose.GraphInterruptOption)
}
func newCancelContext() *cancelContext {
return &cancelContext{
cancelChan: make(chan struct{}),
immediateChan: make(chan struct{}),
doneChan: make(chan struct{}),
timeoutNotify: make(chan struct{}, 1),
recursiveChan: make(chan struct{}),
root: true,
}
}
func (cc *cancelContext) isRoot() bool {
return cc != nil && cc.root
}
func (cc *cancelContext) isRecursive() bool {
return cc != nil && atomic.LoadInt32(&cc.recursive) == 1
}
// setRecursive(false) is a no-op; recursive is monotonically escalating:
// once set to true, it cannot be reverted.
func (cc *cancelContext) setRecursive(v bool) {
if v && atomic.CompareAndSwapInt32(&cc.recursive, 0, 1) {
close(cc.recursiveChan)
}
}
// deriveAgentToolCancelContext creates the cancelContext used by an AgentTool's
// internal agent. It receives recursive cancel propagation from the parent
// AgentTool call. The caller MUST ensure the child's markDone() is eventually
// called (e.g., via wrapIterWithCancelCtx's defer) or that ctx is canceled;
// otherwise the two propagation goroutines will leak.
func (cc *cancelContext) deriveAgentToolCancelContext(ctx context.Context) *cancelContext {
if cc == nil {
return nil
}
child := newCancelContext()
child.root = false
child.parent = cc
// Each goroutine below propagates one signal class (cancel / immediate) to
// the child. The pattern is a two-phase select:
// Phase 1: wait for the parent signal (or child/ctx completion).
// Phase 2: if the signal fired but recursive mode is not active yet,
// enter a second select waiting for either recursive escalation
// (recursiveChan) or child/ctx completion. This ensures
// non-recursive cancels leave children unaware, while a late
// escalation to recursive still propagates.
go func() {
select {
case <-cc.cancelChan:
if cc.isRecursive() {
child.setRecursive(true)
child.triggerCancel(cc.getMode())
return
}
select {
case <-cc.recursiveChan:
child.setRecursive(true)
child.triggerCancel(cc.getMode())
case <-child.doneChan:
case <-ctx.Done():
}
case <-child.doneChan:
case <-ctx.Done():
}
}()
go func() {
select {
case <-cc.immediateChan:
if cc.isRecursive() {
child.setRecursive(true)
child.triggerImmediateCancel()
return
}
select {
case <-cc.recursiveChan:
child.setRecursive(true)
child.triggerImmediateCancel()
case <-child.doneChan:
case <-ctx.Done():
}
case <-child.doneChan:
case <-ctx.Done():
}
}()
return child
}
func (cc *cancelContext) triggerCancel(mode CancelMode) {
cc.setMode(mode)
if atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) {
close(cc.cancelChan)
}
}
func (cc *cancelContext) triggerImmediateCancel() {
atomic.StoreInt32(&cc.escalated, 1)
cc.setMode(CancelImmediate)
if atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) {
close(cc.cancelChan)
}
cc.sendImmediateInterrupt()
}
func (cc *cancelContext) getMode() CancelMode {
if cc == nil {
return CancelImmediate
}
return CancelMode(atomic.LoadInt32(&cc.mode))
}
func (cc *cancelContext) setMode(mode CancelMode) {
atomic.StoreInt32(&cc.mode, int32(mode))
}
func (cc *cancelContext) getDeadlineUnixNano() int64 {
return atomic.LoadInt64(&cc.deadlineUnixNano)
}
func (cc *cancelContext) setDeadlineUnixNano(v int64) {
atomic.StoreInt64(&cc.deadlineUnixNano, v)
}
func (cc *cancelContext) wakeTimeoutController() {
select {
case cc.timeoutNotify <- struct{}{}:
default:
}
}
// shouldCancel returns true if a cancel has been requested (cancelChan is closed).
func (cc *cancelContext) shouldCancel() bool {
if cc == nil {
return false
}
select {
case <-cc.cancelChan:
return true
default:
return false
}
}
// isImmediateCancelled returns true if an immediate graph interrupt has been
// fired (CancelImmediate or timeout escalation). This is stronger than
// shouldCancel: it means the compose graph is being torn down right now and
// orphaned goroutines should not attempt to send events.
func (cc *cancelContext) isImmediateCancelled() bool {
if cc == nil {
return false
}
select {
case <-cc.immediateChan:
return true
default:
return false
}
}
// sendImmediateInterrupt sends the compose graph interrupt signal via graphInterruptFuncs.
// Also closes immediateChan (used by cancelMonitoredModel to abort an in-progress stream).
// Returns false if an interrupt was already sent or if no graphInterruptFuncs have been
// registered yet (the deferred fire in setGraphInterruptFunc will handle that case).
func (cc *cancelContext) sendImmediateInterrupt() bool {
cc.mu.Lock()
if !atomic.CompareAndSwapInt32(&cc.interruptSent, interruptNotSent, interruptImmediate) {
cc.mu.Unlock()
return false
}
close(cc.immediateChan)
fns := make([]func(...compose.GraphInterruptOption), len(cc.graphInterruptFuncs))
copy(fns, cc.graphInterruptFuncs)
if cc.isRecursive() && cc.hasAgentToolDescendant() {
select {
case <-cc.doneChan:
cc.mu.Unlock()
return true
case <-time.After(defaultCancelImmediateGracePeriod):
}
}
if len(fns) == 0 {
cc.mu.Unlock()
return false
}
for _, fn := range fns {
fn(compose.WithGraphInterruptTimeout(0))
}
cc.mu.Unlock()
return true
}
// setGraphInterruptFunc appends a graph interrupt function to the list.
// If an immediate cancel was already requested, fires it retroactively.
// Multiple functions can be registered (e.g. one per parallel sub-agent).
//
// Both this method and sendImmediateInterrupt hold cc.mu across the entire
// check-and-fire sequence, ensuring each interrupt function is called exactly
// once (compose.WithGraphInterrupt returns a non-idempotent closure that panics
// on double-call).
func (cc *cancelContext) setGraphInterruptFunc(interrupt func(...compose.GraphInterruptOption)) {
cc.mu.Lock()
cc.graphInterruptFuncs = append(cc.graphInterruptFuncs, interrupt)
shouldFire := atomic.LoadInt32(&cc.interruptSent) == interruptImmediate
if shouldFire {
interrupt(compose.WithGraphInterruptTimeout(0))
}
cc.mu.Unlock()
}
// markDone marks the execution as finished through any non-cancel path
// (normal completion, business interrupt, or error).
// This is safe to call even if a cancel is in progress — it allows the
// cancel func to detect that execution finished before cancel took effect.
func (cc *cancelContext) markDone() {
if cc == nil {
return
}
if atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateDone) ||
atomic.CompareAndSwapInt32(&cc.state, stateCancelling, stateDone) {
cc.doneOnce.Do(func() { close(cc.doneChan) })
}
}
func (cc *cancelContext) hasAgentToolDescendant() bool {
return cc != nil && atomic.LoadInt32(&cc.agentToolDescendant) == 1
}
func (cc *cancelContext) markAgentToolDescendant() {
for cur := cc; cur != nil; cur = cur.parent {
atomic.StoreInt32(&cur.agentToolDescendant, 1)
}
}
// markCancelHandled signals that the cancel path in the runFunc has created
// and sent a CancelError. Transitions state to stateCancelHandled so that:
// 1. The deferred markDone() becomes a no-op (CAS from cancelling fails).
// 2. buildCancelFunc sees stateCancelHandled and returns nil (cancel succeeded).
// Returns true if the transition succeeded, false if cancel was already handled
// (e.g., by a sub-agent). This prevents duplicate CancelError emission.
func (cc *cancelContext) markCancelHandled() bool {
if cc == nil {
return false
}
if atomic.CompareAndSwapInt32(&cc.state, stateCancelling, stateCancelHandled) {
cc.doneOnce.Do(func() { close(cc.doneChan) })
return true
}
return false
}
// createCancelError creates a CancelError based on the current cancel state.
func (cc *cancelContext) createCancelError() *CancelError {
info := &AgentCancelInfo{}
info.Mode = cc.getMode()
if atomic.LoadInt32(&cc.escalated) == 1 {
info.Escalated = true
info.Timeout = atomic.LoadInt32(&cc.timeoutEscalated) == 1
}
return &CancelError{
Info: info,
}
}
func (cc *cancelContext) createAndMarkCancelHandled() (*CancelError, bool) {
cc.cancelMu.Lock()
defer cc.cancelMu.Unlock()
cancelErr := cc.createCancelError()
ok := cc.markCancelHandled()
return cancelErr, ok
}
// buildCancelFunc builds the AgentCancelFunc for external use.
func (cc *cancelContext) buildCancelFunc() AgentCancelFunc {
joinMode := func(a, b CancelMode) CancelMode {
if a == CancelImmediate || b == CancelImmediate {
return CancelImmediate
}
return a | b
}
parseReq := func(callOpts ...AgentCancelOption) *agentCancelConfig {
cfg := &agentCancelConfig{Mode: CancelImmediate}
for _, opt := range callOpts {
opt(cfg)
}
return cfg
}
startTimeoutController := func() {
cc.timeoutOnce.Do(func() {
go func() {
for {
select {
case <-cc.doneChan:
return
default:
}
mode := cc.getMode()
if mode == CancelImmediate {
return
}
deadline := cc.getDeadlineUnixNano()
if deadline == 0 {
select {
case <-cc.timeoutNotify:
continue
case <-cc.doneChan:
return
}
}
now := time.Now().UnixNano()
wait := time.Duration(deadline - now)
if wait <= 0 {
atomic.StoreInt32(&cc.escalated, 1)
atomic.StoreInt32(&cc.timeoutEscalated, 1)
cc.sendImmediateInterrupt()
return
}
timer := time.NewTimer(wait)
select {
case <-timer.C:
timer.Stop()
atomic.StoreInt32(&cc.escalated, 1)
atomic.StoreInt32(&cc.timeoutEscalated, 1)
cc.sendImmediateInterrupt()
return
case <-cc.timeoutNotify:
timer.Stop()
continue
case <-cc.doneChan:
timer.Stop()
return
}
}
}()
})
}
newHandle := func(wait func() error) *CancelHandle {
return &CancelHandle{wait: wait}
}
waitForCompletion := func() error {
<-cc.doneChan
st := atomic.LoadInt32(&cc.state)
switch st {
case stateDone:
return ErrExecutionEnded
default:
if atomic.LoadInt32(&cc.timeoutEscalated) == 1 {
return ErrCancelTimeout
}
return nil
}
}
return func(callOpts ...AgentCancelOption) (*CancelHandle, bool) {
req := parseReq(callOpts...)
st := atomic.LoadInt32(&cc.state)
switch st {
case stateCancelHandled:
return newHandle(func() error { return nil }), false
case stateDone:
return newHandle(func() error { return ErrExecutionEnded }), false
}
var needImmediate, needTimeoutCtl bool
cc.cancelMu.Lock()
st = atomic.LoadInt32(&cc.state)
switch st {
case stateCancelHandled:
cc.cancelMu.Unlock()
return newHandle(func() error { return nil }), false
case stateDone:
cc.cancelMu.Unlock()
return newHandle(func() error { return ErrExecutionEnded }), false
}
curMode := cc.getMode()
if st == stateRunning {
if !atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) {
st = atomic.LoadInt32(&cc.state)
cc.cancelMu.Unlock()
if st == stateDone {
return newHandle(func() error { return ErrExecutionEnded }), false
}
return newHandle(waitForCompletion), true
}
curMode = req.Mode
cc.setMode(curMode)
atomic.StoreInt32(&cc.startedMode, int32(curMode))
cc.setRecursive(req.Recursive)
close(cc.cancelChan)
} else {
// Recursive is monotonic: once set, cannot be unset. The first
// cancel call uses the bool directly; subsequent calls only
// escalate (false → true) — setRecursive(false) is a no-op.
curMode = joinMode(curMode, req.Mode)
cc.setMode(curMode)
if req.Recursive {
cc.setRecursive(true)
}
}
if curMode == CancelImmediate {
cc.setDeadlineUnixNano(0)
needImmediate = true
} else if req.Timeout != nil && *req.Timeout > 0 {
proposed := time.Now().Add(*req.Timeout).UnixNano()
existing := cc.getDeadlineUnixNano()
if existing == 0 || proposed < existing {
cc.setDeadlineUnixNano(proposed)
cc.wakeTimeoutController()
}
needTimeoutCtl = cc.getDeadlineUnixNano() != 0
}
cc.cancelMu.Unlock()
if needImmediate {
if atomic.LoadInt32(&cc.startedMode) != int32(CancelImmediate) {
atomic.StoreInt32(&cc.escalated, 1)
}
cc.sendImmediateInterrupt()
}
if needTimeoutCtl {
startTimeoutController()
}
return newHandle(waitForCompletion), true
}
}
// wrapIterWithCancelCtx wraps an iterator with cancel lifecycle management.
// It calls markDone when the inner iterator is fully drained, ensuring the
// cancelContext's doneChan is closed and propagation goroutines can exit.
//
// For root cancelContexts (created by WithCancel, not deriveAgentToolCancelContext), it also
// converts interrupt ACTION events to CancelError when cancel is active.
// This is the single point of interrupt-to-CancelError conversion in the
// system — Runner.handleIter only enriches the resulting CancelError with
// checkpoint metadata.
//
// Interrupt absorption: ALL interrupts are converted when cancel is active,
// including business interrupts (compose.Interrupt from user code). Cancel and
// business interrupts cannot be reliably distinguished in concurrent execution
// (parallel workflows, concurrent tool calls) where they merge into a single
// composite signal. The business interrupt data is preserved in the checkpoint
// and re-fires naturally on resume.
//
// This conversion MUST happen in this wrapper (not deferred to Runner.handleIter)
// because markDone runs as a defer in this goroutine — if the interrupt event
// were passed through unconverted, markDone would transition stateCancelling→stateDone
// before the Runner goroutine could call createAndMarkCancelHandled, causing it
// to fail the CAS.
func wrapIterWithCancelCtx[M MessageType](iter *AsyncIterator[*TypedAgentEvent[M]], cancelCtx *cancelContext) *AsyncIterator[*TypedAgentEvent[M]] {
if cancelCtx == nil {
return iter
}
it, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
go func() {
defer cancelCtx.markDone()
defer gen.Close()
for {
event, ok := iter.Next()
if !ok {
break
}
if cancelCtx.isRoot() && event.Action != nil && event.Action.internalInterrupted != nil {
if cancelCtx.shouldCancel() {
cancelErr, ok := cancelCtx.createAndMarkCancelHandled()
if ok {
cancelErr.interruptSignal = event.Action.internalInterrupted
gen.Send(&TypedAgentEvent[M]{Err: cancelErr})
}
return
}
}
gen.Send(event)
}
}()
return it
}
// typedCancelMonitoredModel wraps a model with cancel monitoring.
// Generate: pure delegate to the inner model (CancelAfterChatModel is handled
// by a dedicated node after the ChatModel in the compose graph).
// Stream: pipes chunks through a goroutine that selects on immediateChan for
// CancelImmediate abort.
type typedCancelMonitoredModel[M MessageType] struct {
inner model.BaseModel[M]
cancelContext *cancelContext
}
type recvResult[T any] struct {
data T
err error
}
func (m *typedCancelMonitoredModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
return m.inner.Generate(ctx, input, opts...)
}
func (m *typedCancelMonitoredModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
stream, err := m.inner.Stream(ctx, input, opts...)
if err != nil {
return nil, err
}
wrapped := wrapStreamWithCancelMonitoring(stream, m.cancelContext)
return wrapped, nil
}
// wrapStreamWithCancelMonitoring wraps a stream with cancel monitoring.
// When immediateChan fires (CancelImmediate or timeout escalation), the output
// stream is terminated with ErrStreamCanceled.
func wrapStreamWithCancelMonitoring[T any](stream *schema.StreamReader[T], cc *cancelContext) *schema.StreamReader[T] {
if cc == nil {
return stream
}
// Already canceled — terminate immediately
select {
case <-cc.immediateChan:
stream.Close()
r, w := schema.Pipe[T](1)
var zero T
w.Send(zero, ErrStreamCanceled)
w.Close()
return r
default:
}
reader, writer := schema.Pipe[T](1)
go func() {
done := make(chan struct{})
defer close(done)
defer writer.Close()
defer stream.Close()
ch := make(chan recvResult[T])
go func() {
defer close(ch)
for {
chunk, recvErr := stream.Recv()
select {
case ch <- recvResult[T]{chunk, recvErr}:
case <-done:
return
}
if recvErr != nil {
return
}
}
}()
for {
select {
case <-cc.immediateChan:
var zero T
writer.Send(zero, ErrStreamCanceled)
return
case r, ok := <-ch:
if !ok {
return
}
if r.err != nil {
if r.err == io.EOF {
return
}
var zero T
writer.Send(zero, r.err)
return
}
if closed := writer.Send(r.data, nil); closed {
return
}
}
}
}()
return reader
}
// cancelMonitoredToolHandler wraps streamable tool calls with cancel monitoring.
// When CancelImmediate fires, the tool output stream is terminated with ErrStreamCanceled.
// This handler reads the cancelContext from the Go context via getCancelContext.
type cancelMonitoredToolHandler struct{}
func (h *cancelMonitoredToolHandler) WrapStreamableToolCall(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
output, err := next(ctx, input)
if err != nil {
return nil, err
}
cc := getCancelContext(ctx)
if cc == nil {
return output, nil
}
wrapped := wrapStreamWithCancelMonitoring(output.Result, cc)
return &compose.StreamToolOutput{Result: wrapped}, nil
}
}
func (h *cancelMonitoredToolHandler) WrapEnhancedStreamableToolCall(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint {
return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {
output, err := next(ctx, input)
if err != nil {
return nil, err
}
cc := getCancelContext(ctx)
if cc == nil {
return output, nil
}
wrapped := wrapStreamWithCancelMonitoring(output.Result, cc)
return &compose.EnhancedStreamableToolOutput{Result: wrapped}, nil
}
}
File diff suppressed because it is too large Load Diff
+125
View File
@@ -0,0 +1,125 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/cloudwego/eino/compose"
)
func TestAgentCancelFunc_MultiCall_EscalateToImmediate(t *testing.T) {
cc := newCancelContext()
var interruptCalls int32
cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) {
atomic.AddInt32(&interruptCalls, 1)
})
cancelFn := cc.buildCancelFunc()
handle1, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel))
handle2, _ := cancelFn(WithAgentCancelMode(CancelImmediate))
assert.Equal(t, int32(1), atomic.LoadInt32(&interruptCalls))
cancelErr := cc.createCancelError()
assert.Equal(t, CancelImmediate, cancelErr.Info.Mode)
assert.True(t, cancelErr.Info.Escalated)
assert.False(t, cancelErr.Info.Timeout)
assert.True(t, cc.markCancelHandled())
assert.NoError(t, handle1.Wait())
assert.NoError(t, handle2.Wait())
}
func TestAgentCancelFunc_MultiCall_JoinSafePointModes(t *testing.T) {
cc := newCancelContext()
cancelFn := cc.buildCancelFunc()
handle1, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel))
handle2, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls))
want := CancelAfterChatModel | CancelAfterToolCalls
assert.Equal(t, want, cc.getMode())
assert.True(t, cc.markCancelHandled())
assert.NoError(t, handle1.Wait())
assert.NoError(t, handle2.Wait())
}
func TestAgentCancelFunc_MultiCall_TimeoutDeadlineJoinUsesAbsoluteTime(t *testing.T) {
cc := newCancelContext()
cancelFn := cc.buildCancelFunc()
handle1, _ := cancelFn(
WithAgentCancelMode(CancelAfterChatModel),
WithAgentCancelTimeout(200*time.Millisecond),
)
firstDeadline := cc.getDeadlineUnixNano()
assert.NotZero(t, firstDeadline)
time.Sleep(50 * time.Millisecond)
handle2, _ := cancelFn(
WithAgentCancelMode(CancelAfterToolCalls),
WithAgentCancelTimeout(60*time.Millisecond),
)
secondDeadline := cc.getDeadlineUnixNano()
assert.NotZero(t, secondDeadline)
assert.Less(t, secondDeadline, firstDeadline)
assert.True(t, cc.markCancelHandled())
assert.NoError(t, handle1.Wait())
assert.NoError(t, handle2.Wait())
}
func TestAgentCancelFunc_MultiCall_TimeoutEscalationReturnsErrCancelTimeout(t *testing.T) {
cc := newCancelContext()
var interruptCalls int32
interruptCh := make(chan struct{}, 1)
cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) {
atomic.AddInt32(&interruptCalls, 1)
select {
case interruptCh <- struct{}{}:
default:
}
})
cancelFn := cc.buildCancelFunc()
handle, _ := cancelFn(
WithAgentCancelMode(CancelAfterChatModel),
WithAgentCancelTimeout(30*time.Millisecond),
)
select {
case <-interruptCh:
case <-time.After(1 * time.Second):
t.Fatal("timeout escalation did not interrupt")
}
assert.Equal(t, int32(1), atomic.LoadInt32(&interruptCalls))
cancelErr := cc.createCancelError()
assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode)
assert.True(t, cancelErr.Info.Escalated)
assert.True(t, cancelErr.Info.Timeout)
assert.True(t, cc.markCancelHandled())
assert.Equal(t, ErrCancelTimeout, handle.Wait())
}
+387
View File
@@ -0,0 +1,387 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"runtime"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func assertNotClosedWithin(t *testing.T, ch <-chan struct{}, d time.Duration) {
t.Helper()
select {
case <-ch:
t.Fatal("channel was closed but should not have been")
case <-time.After(d):
}
}
func setupParentChild(t *testing.T) (parent, child *cancelContext, cleanup func()) {
parent = newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
child = parent.deriveAgentToolCancelContext(ctx)
cleanup = func() {
child.markDone()
cancel()
}
t.Cleanup(cleanup)
return parent, child, cleanup
}
func TestDeriveAgentToolCancelContext(t *testing.T) {
t.Run("Shallow", func(t *testing.T) {
t.Run("DoesNotPropagateSafePoint", func(t *testing.T) {
parent, child, _ := setupParentChild(t)
parent.triggerCancel(CancelAfterChatModel)
assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond)
})
t.Run("ImmediateDoesNotPropagate", func(t *testing.T) {
parent, child, _ := setupParentChild(t)
parent.triggerImmediateCancel()
assertNotClosedWithin(t, child.immediateChan, 50*time.Millisecond)
})
t.Run("GrandchildNoPropagation", func(t *testing.T) {
a := newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
b := a.deriveAgentToolCancelContext(ctx)
c := b.deriveAgentToolCancelContext(ctx)
t.Cleanup(func() {
c.markDone()
b.markDone()
cancel()
})
a.triggerCancel(CancelAfterChatModel)
assertNotClosedWithin(t, b.cancelChan, 50*time.Millisecond)
assertNotClosedWithin(t, c.cancelChan, 50*time.Millisecond)
})
t.Run("NeverRecursive_GoroutineCleanup", func(t *testing.T) {
runtime.GC()
time.Sleep(50 * time.Millisecond)
before := runtime.NumGoroutine()
parent := newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
child := parent.deriveAgentToolCancelContext(ctx)
parent.triggerCancel(CancelAfterChatModel)
time.Sleep(100 * time.Millisecond)
child.markDone()
cancel()
time.Sleep(200 * time.Millisecond)
runtime.GC()
time.Sleep(50 * time.Millisecond)
after := runtime.NumGoroutine()
assert.InDelta(t, before, after, 5, "goroutine leak detected: before=%d after=%d", before, after)
})
})
t.Run("Recursive", func(t *testing.T) {
t.Run("PropagatesSafePoint", func(t *testing.T) {
parent, child, _ := setupParentChild(t)
parent.setRecursive(true)
parent.triggerCancel(CancelAfterChatModel)
select {
case <-child.cancelChan:
case <-time.After(1 * time.Second):
t.Fatal("child did not receive cancel within 1s")
}
assert.True(t, child.shouldCancel())
})
t.Run("ImmediatePropagates", func(t *testing.T) {
parent, child, _ := setupParentChild(t)
parent.setRecursive(true)
parent.triggerImmediateCancel()
select {
case <-child.immediateChan:
case <-time.After(1 * time.Second):
t.Fatal("child did not receive immediate cancel within 1s")
}
assert.True(t, child.isImmediateCancelled())
})
t.Run("GrandchildPropagation", func(t *testing.T) {
a := newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
b := a.deriveAgentToolCancelContext(ctx)
c := b.deriveAgentToolCancelContext(ctx)
t.Cleanup(func() {
c.markDone()
b.markDone()
cancel()
})
a.setRecursive(true)
a.triggerCancel(CancelAfterChatModel)
select {
case <-b.cancelChan:
case <-time.After(1 * time.Second):
t.Fatal("B did not receive cancel within 1s")
}
select {
case <-c.cancelChan:
case <-time.After(1 * time.Second):
t.Fatal("C did not receive cancel within 1s")
}
assert.True(t, b.shouldCancel())
assert.True(t, c.shouldCancel())
})
t.Run("SetBeforeCancel", func(t *testing.T) {
parent, child, _ := setupParentChild(t)
parent.setRecursive(true)
parent.triggerCancel(CancelAfterChatModel)
select {
case <-child.cancelChan:
case <-time.After(1 * time.Second):
t.Fatal("child did not receive cancel within 1s")
}
assert.True(t, child.shouldCancel())
})
t.Run("AfterRecursiveAndCancelAlreadySet", func(t *testing.T) {
parent := newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
parent.setRecursive(true)
parent.triggerCancel(CancelAfterChatModel)
child := parent.deriveAgentToolCancelContext(ctx)
t.Cleanup(func() {
child.markDone()
cancel()
})
select {
case <-child.cancelChan:
case <-time.After(1 * time.Second):
t.Fatal("child did not immediately receive cancel")
}
assert.True(t, child.shouldCancel())
})
})
t.Run("Escalation", func(t *testing.T) {
t.Run("EscalateFromNonRecursive", func(t *testing.T) {
parent, child, _ := setupParentChild(t)
parent.triggerCancel(CancelAfterChatModel)
assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond)
parent.setRecursive(true)
select {
case <-child.cancelChan:
case <-time.After(1 * time.Second):
t.Fatal("child did not receive cancel after escalation within 1s")
}
assert.True(t, child.shouldCancel())
})
t.Run("EscalateImmediate", func(t *testing.T) {
parent, child, _ := setupParentChild(t)
parent.triggerImmediateCancel()
assertNotClosedWithin(t, child.immediateChan, 50*time.Millisecond)
parent.setRecursive(true)
select {
case <-child.immediateChan:
case <-time.After(1 * time.Second):
t.Fatal("child did not receive immediate cancel after escalation within 1s")
}
assert.True(t, child.isImmediateCancelled())
})
})
}
func TestDeriveAgentToolCancelContext_Race(t *testing.T) {
t.Run("SetRecursiveConcurrentWithCancelChan", func(t *testing.T) {
for i := 0; i < 100; i++ {
parent := newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
child := parent.deriveAgentToolCancelContext(ctx)
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
parent.setRecursive(true)
}()
go func() {
defer wg.Done()
parent.triggerCancel(CancelAfterChatModel)
}()
wg.Wait()
select {
case <-child.cancelChan:
case <-time.After(1 * time.Second):
t.Fatalf("iteration %d: child did not receive cancel within 1s", i)
}
assert.True(t, child.shouldCancel())
child.markDone()
cancel()
}
})
t.Run("ChildCompletesBeforeEscalation", func(t *testing.T) {
parent := newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
child := parent.deriveAgentToolCancelContext(ctx)
parent.triggerCancel(CancelAfterChatModel)
time.Sleep(50 * time.Millisecond)
child.markDone()
time.Sleep(50 * time.Millisecond)
parent.setRecursive(true)
assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond)
})
t.Run("MultipleChildren_PartialCompletion", func(t *testing.T) {
parent := newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
child1 := parent.deriveAgentToolCancelContext(ctx)
child2 := parent.deriveAgentToolCancelContext(ctx)
parent.triggerCancel(CancelAfterChatModel)
time.Sleep(50 * time.Millisecond)
child1.markDone()
time.Sleep(50 * time.Millisecond)
parent.setRecursive(true)
select {
case <-child2.cancelChan:
case <-time.After(1 * time.Second):
t.Fatal("running child did not receive cancel within 1s")
}
assert.True(t, child2.shouldCancel())
assert.False(t, child1.shouldCancel())
child2.markDone()
})
t.Run("ContextCancelConcurrentWithRecursive", func(t *testing.T) {
done := make(chan struct{})
go func() {
defer close(done)
parent := newCancelContext()
ctx, cancel := context.WithCancel(context.Background())
child := parent.deriveAgentToolCancelContext(ctx)
parent.triggerCancel(CancelAfterChatModel)
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
cancel()
}()
go func() {
defer wg.Done()
parent.setRecursive(true)
}()
wg.Wait()
child.markDone()
}()
select {
case <-done:
case <-time.After(1 * time.Second):
t.Fatal("deadlock detected")
}
})
t.Run("ConcurrentSetRecursive", func(t *testing.T) {
parent := newCancelContext()
var wg sync.WaitGroup
for i := 0; i < 10; i++ {
wg.Add(1)
go func() {
defer wg.Done()
parent.setRecursive(true)
}()
}
done := make(chan struct{})
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
case <-time.After(1 * time.Second):
t.Fatal("deadlock or panic in concurrent setRecursive")
}
assert.True(t, parent.isRecursive())
})
}
+3794
View File
File diff suppressed because it is too large Load Diff
+1714
View File
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+35
View File
@@ -0,0 +1,35 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import "github.com/cloudwego/eino/adk/internal"
// Language represents the language setting for the ADK built-in prompts.
type Language = internal.Language
const (
// LanguageEnglish represents English language.
LanguageEnglish Language = internal.LanguageEnglish
// LanguageChinese represents Chinese language.
LanguageChinese Language = internal.LanguageChinese
)
// SetLanguage sets the language for the ADK built-in prompts.
// The default language is English if not explicitly set.
func SetLanguage(lang Language) error {
return internal.SetLanguage(lang)
}
+301
View File
@@ -0,0 +1,301 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"errors"
"runtime/debug"
"sync"
"github.com/cloudwego/eino/components"
"github.com/cloudwego/eino/internal/safe"
"github.com/cloudwego/eino/schema"
)
func init() {
schema.RegisterName[*deterministicTransferState]("_eino_adk_deterministic_transfer_state")
}
type deterministicTransferState struct {
EventList []*agentEventWrapper
}
// AgentWithDeterministicTransferTo wraps an agent to transfer to given agents deterministically.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func AgentWithDeterministicTransferTo(_ context.Context, config *DeterministicTransferConfig) Agent {
if ra, ok := config.Agent.(ResumableAgent); ok {
return &resumableAgentWithDeterministicTransferTo{
agent: ra,
toAgentNames: config.ToAgentNames,
}
}
return &agentWithDeterministicTransferTo{
agent: config.Agent,
toAgentNames: config.ToAgentNames,
}
}
type agentWithDeterministicTransferTo struct {
agent Agent
toAgentNames []string
}
func (a *agentWithDeterministicTransferTo) Description(ctx context.Context) string {
return a.agent.Description(ctx)
}
func (a *agentWithDeterministicTransferTo) Name(ctx context.Context) string {
return a.agent.Name(ctx)
}
func (a *agentWithDeterministicTransferTo) GetType() string {
if typer, ok := a.agent.(components.Typer); ok {
return typer.GetType()
}
return "DeterministicTransfer"
}
func (a *agentWithDeterministicTransferTo) Run(ctx context.Context,
input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
if fa, ok := a.agent.(*flowAgent); ok {
return runFlowAgentWithIsolatedSession(ctx, fa, input, a.toAgentNames, options...)
}
aIter := a.agent.Run(ctx, input, options...)
iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
go forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames)
return iterator
}
type resumableAgentWithDeterministicTransferTo struct {
agent ResumableAgent
toAgentNames []string
}
func (a *resumableAgentWithDeterministicTransferTo) Description(ctx context.Context) string {
return a.agent.Description(ctx)
}
func (a *resumableAgentWithDeterministicTransferTo) Name(ctx context.Context) string {
return a.agent.Name(ctx)
}
func (a *resumableAgentWithDeterministicTransferTo) GetType() string {
if typer, ok := a.agent.(components.Typer); ok {
return typer.GetType()
}
return "DeterministicTransfer"
}
func (a *resumableAgentWithDeterministicTransferTo) Run(ctx context.Context,
input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
if fa, ok := a.agent.(*flowAgent); ok {
return runFlowAgentWithIsolatedSession(ctx, fa, input, a.toAgentNames, options...)
}
aIter := a.agent.Run(ctx, input, options...)
iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
go forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames)
return iterator
}
func (a *resumableAgentWithDeterministicTransferTo) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
if fa, ok := a.agent.(*flowAgent); ok {
return resumeFlowAgentWithIsolatedSession(ctx, fa, info, a.toAgentNames, opts...)
}
aIter := a.agent.Resume(ctx, info, opts...)
iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
go forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames)
return iterator
}
func forwardEventsAndAppendTransfer(iter *AsyncIterator[*AgentEvent],
generator *AsyncGenerator[*AgentEvent], toAgentNames []string) {
defer func() {
if panicErr := recover(); panicErr != nil {
generator.Send(&AgentEvent{Err: safe.NewPanicErr(panicErr, debug.Stack())})
}
generator.Close()
}()
var lastEvent *AgentEvent
for {
event, ok := iter.Next()
if !ok {
break
}
generator.Send(event)
lastEvent = event
}
if lastEvent != nil && lastEvent.Action != nil && (lastEvent.Action.Interrupted != nil || lastEvent.Action.Exit) {
return
}
sendTransferEvents(generator, toAgentNames)
}
func runFlowAgentWithIsolatedSession(ctx context.Context, fa *flowAgent, input *AgentInput,
toAgentNames []string, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
parentSession := getSession(ctx)
parentRunCtx := getRunCtx(ctx)
isolatedSession := &runSession{
Values: parentSession.Values,
valuesMtx: parentSession.valuesMtx,
}
if isolatedSession.valuesMtx == nil {
isolatedSession.valuesMtx = &sync.Mutex{}
}
if isolatedSession.Values == nil {
isolatedSession.Values = make(map[string]any)
}
ctx = setRunCtx(ctx, &runContext{
Session: isolatedSession,
RootInput: parentRunCtx.RootInput,
RunPath: parentRunCtx.RunPath,
})
iter := fa.Run(ctx, input, options...)
iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
go handleFlowAgentEvents(ctx, iter, generator, isolatedSession, parentSession, toAgentNames)
return iterator
}
func resumeFlowAgentWithIsolatedSession(ctx context.Context, fa *flowAgent, info *ResumeInfo,
toAgentNames []string, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
state, ok := info.InterruptState.(*deterministicTransferState)
if !ok || state == nil {
return genErrorIter(errors.New("invalid interrupt state for flowAgent resume in deterministic transfer"))
}
parentSession := getSession(ctx)
parentRunCtx := getRunCtx(ctx)
isolatedSession := &runSession{
Values: parentSession.Values,
valuesMtx: parentSession.valuesMtx,
Events: state.EventList,
}
if isolatedSession.valuesMtx == nil {
isolatedSession.valuesMtx = &sync.Mutex{}
}
if isolatedSession.Values == nil {
isolatedSession.Values = make(map[string]any)
}
ctx = setRunCtx(ctx, &runContext{
Session: isolatedSession,
RootInput: parentRunCtx.RootInput,
RunPath: parentRunCtx.RunPath,
})
iter := fa.Resume(ctx, info, opts...)
iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
go handleFlowAgentEvents(ctx, iter, generator, isolatedSession, parentSession, toAgentNames)
return iterator
}
func handleFlowAgentEvents(ctx context.Context, iter *AsyncIterator[*AgentEvent],
generator *AsyncGenerator[*AgentEvent], isolatedSession, parentSession *runSession, toAgentNames []string) {
defer func() {
if panicErr := recover(); panicErr != nil {
generator.Send(&AgentEvent{Err: safe.NewPanicErr(panicErr, debug.Stack())})
}
generator.Close()
}()
var lastEvent *AgentEvent
for {
event, ok := iter.Next()
if !ok {
break
}
if parentSession != nil && (event.Action == nil || event.Action.Interrupted == nil) {
copied := copyTypedAgentEvent(event)
setAutomaticClose(copied)
setAutomaticClose(event)
parentSession.addEvent(copied)
}
if event.Action != nil && event.Action.internalInterrupted != nil {
lastEvent = event
continue
}
generator.Send(event)
lastEvent = event
}
if lastEvent != nil && lastEvent.Action != nil {
if lastEvent.Action.internalInterrupted != nil {
events := isolatedSession.getEvents()
state := &deterministicTransferState{EventList: events}
compositeEvent := CompositeInterrupt(ctx, "deterministic transfer wrapper interrupted",
state, lastEvent.Action.internalInterrupted)
generator.Send(compositeEvent)
return
}
if lastEvent.Action.Exit {
return
}
}
sendTransferEvents(generator, toAgentNames)
}
func sendTransferEvents(generator *AsyncGenerator[*AgentEvent], toAgentNames []string) {
for _, toAgentName := range toAgentNames {
aMsg, tMsg := GenTransferMessages(context.Background(), toAgentName)
aEvent := EventFromMessage(aMsg, nil, schema.Assistant, "")
generator.Send(aEvent)
tEvent := EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)
tEvent.Action = &AgentAction{
TransferToAgent: &TransferToAgentAction{
DestAgentName: toAgentName,
},
}
generator.Send(tEvent)
}
}
+731
View File
@@ -0,0 +1,731 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/cloudwego/eino/schema"
)
type dtTestStore struct {
data map[string][]byte
}
func newDTTestStore() *dtTestStore {
return &dtTestStore{data: make(map[string][]byte)}
}
func (s *dtTestStore) Set(_ context.Context, key string, value []byte) error {
s.data[key] = value
return nil
}
func (s *dtTestStore) Get(_ context.Context, key string) ([]byte, bool, error) {
v, ok := s.data[key]
return v, ok, nil
}
type dtTestAgent struct {
name string
runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]
resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]
}
func (a *dtTestAgent) Name(_ context.Context) string {
return a.name
}
func (a *dtTestAgent) Description(_ context.Context) string {
return a.name + " description"
}
func (a *dtTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
return a.runFn(ctx, input, options...)
}
func (a *dtTestAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
if a.resumeFn != nil {
return a.resumeFn(ctx, info, opts...)
}
return a.runFn(ctx, &AgentInput{}, opts...)
}
func TestDeterministicTransferFlowAgentInterruptResume(t *testing.T) {
ctx := context.Background()
store := newDTTestStore()
interruptData := "interrupt_data"
var runCount int
innerAgent := &dtTestAgent{
name: "inner",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
runCount++
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
gen.Send(EventFromMessage(schema.AssistantMessage("before interrupt", nil), nil, schema.Assistant, ""))
intEvent := Interrupt(ctx, interruptData)
intEvent.Action.Interrupted.Data = interruptData
gen.Send(intEvent)
}()
return iter
},
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
runCount++
assert.True(t, info.WasInterrupted, "innerAgent resumeFn: should be interrupted")
assert.True(t, info.IsResumeTarget, "innerAgent resumeFn: should be resume target")
runCtx := getRunCtx(ctx)
assert.NotNil(t, runCtx, "innerAgent resumeFn: runCtx should not be nil")
assert.NotNil(t, runCtx.Session, "innerAgent resumeFn: runCtx.Session should not be nil")
var agentEvents []*AgentEvent
for _, ev := range runCtx.Session.Events {
if ev.AgentEvent != nil {
agentEvents = append(agentEvents, ev.AgentEvent)
}
}
assert.Len(t, agentEvents, 1, "innerAgent resumeFn: should have exactly 1 agent event")
if len(agentEvents) == 1 {
ev := agentEvents[0]
assert.Equal(t, "inner", ev.AgentName, "innerAgent resumeFn: event should be from inner agent")
assert.Equal(t, "before interrupt", ev.Output.MessageOutput.Message.Content, "innerAgent resumeFn: event content should be 'before interrupt'")
assert.Len(t, ev.RunPath, 2, "innerAgent resumeFn: RunPath should have 2 steps (outer agent, inner agent)")
if len(ev.RunPath) == 2 {
assert.Equal(t, "outer", ev.RunPath[0].agentName, "innerAgent resumeFn: RunPath[0] should be outer agent")
assert.Equal(t, "inner", ev.RunPath[1].agentName, "innerAgent resumeFn: RunPath[1] should be inner agent")
}
}
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
gen.Send(EventFromMessage(schema.AssistantMessage("after resume", nil), nil, schema.Assistant, ""))
}()
return iter
},
}
innerFlowAgent := toFlowAgent(ctx, innerAgent)
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: innerFlowAgent,
ToAgentNames: []string{"next_agent"},
})
outerAgent := &dtTestAgent{
name: "outer",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
return wrapped.Run(ctx, input, options...)
},
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
assert.True(t, info.WasInterrupted, "outerAgent resumeFn: should be interrupted")
runCtx := getRunCtx(ctx)
assert.NotNil(t, runCtx, "outerAgent resumeFn: runCtx should not be nil")
assert.NotNil(t, runCtx.Session, "outerAgent resumeFn: runCtx.Session should not be nil")
var agentEvents []*AgentEvent
for _, ev := range runCtx.Session.Events {
if ev.AgentEvent != nil {
agentEvents = append(agentEvents, ev.AgentEvent)
}
}
assert.Len(t, agentEvents, 1, "outerAgent resumeFn: should have exactly 1 agent event")
if len(agentEvents) == 1 {
ev := agentEvents[0]
assert.Equal(t, "inner", ev.AgentName, "outerAgent resumeFn: event should be from inner agent (preserved original)")
assert.Equal(t, "before interrupt", ev.Output.MessageOutput.Message.Content, "outerAgent resumeFn: event content should be 'before interrupt'")
assert.Len(t, ev.RunPath, 2, "outerAgent resumeFn: RunPath should have 2 steps")
if len(ev.RunPath) == 2 {
assert.Equal(t, "outer", ev.RunPath[0].agentName, "outerAgent resumeFn: RunPath[0] should be outer agent")
assert.Equal(t, "inner", ev.RunPath[1].agentName, "outerAgent resumeFn: RunPath[1] should be inner agent")
}
}
ra := wrapped.(ResumableAgent)
return ra.Resume(ctx, info, opts...)
},
}
outerFlowAgent := toFlowAgent(ctx, outerAgent)
runner := NewRunner(ctx, RunnerConfig{
Agent: outerFlowAgent,
EnableStreaming: true,
CheckPointStore: store,
})
iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, WithCheckPointID("cp1"))
var events []*AgentEvent
var interrupted bool
var interruptEvent *AgentEvent
for {
ev, ok := iter.Next()
if !ok {
break
}
events = append(events, ev)
if ev.Action != nil && ev.Action.Interrupted != nil {
interrupted = true
interruptEvent = ev
}
}
assert.Equal(t, 1, runCount, "run should have been called once")
assert.True(t, interrupted, "should have interrupted")
assert.Greater(t, len(events), 0, "should have events")
if interruptEvent == nil {
t.Fatal("should have interrupt event")
}
assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, "should have interrupt contexts")
_, exists, err := store.Get(ctx, "cp1")
assert.NoError(t, err)
assert.True(t, exists, "checkpoint should have been saved")
var hasDeterministicTransferContext bool
for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {
t.Logf("InterruptContext: ID=%s, Info=%v, IsRootCause=%v, Addr=%v", intCtx.ID, intCtx.Info, intCtx.IsRootCause, intCtx.Address)
if intCtx.Info == "deterministic transfer wrapper interrupted" {
hasDeterministicTransferContext = true
}
for parent := intCtx.Parent; parent != nil; parent = parent.Parent {
t.Logf(" Parent: ID=%s, Info=%v, Addr=%v", parent.ID, parent.Info, parent.Address)
if parent.Info == "deterministic transfer wrapper interrupted" {
hasDeterministicTransferContext = true
}
}
}
assert.True(t, hasDeterministicTransferContext, "should have deterministic transfer interrupt context")
var rootCauseID string
for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {
if intCtx.IsRootCause {
rootCauseID = intCtx.ID
break
}
}
assert.NotEmpty(t, rootCauseID, "should have root cause interrupt ID")
resumeIter, err := runner.ResumeWithParams(ctx, "cp1", &ResumeParams{
Targets: map[string]any{rootCauseID: nil},
})
assert.NoError(t, err)
var resumeEvents []*AgentEvent
var resumeErr error
var hasTransfer bool
for {
ev, ok := resumeIter.Next()
if !ok {
break
}
if ev.Err != nil {
resumeErr = ev.Err
}
if ev.Action != nil && ev.Action.TransferToAgent != nil {
hasTransfer = true
}
resumeEvents = append(resumeEvents, ev)
}
assert.Equal(t, 2, runCount, "inner agent should be called twice (once for initial, once for resume)")
assert.NotEmpty(t, resumeEvents, "should have resume events")
assert.True(t, hasTransfer, "should have transfer action after resume")
assert.Error(t, resumeErr, "transfer should fail because next_agent doesn't exist")
assert.Contains(t, resumeErr.Error(), "next_agent", "error should mention the missing agent")
}
func TestDeterministicTransferRunPathPreserved(t *testing.T) {
ctx := context.Background()
store := newDTTestStore()
var collectedRunPaths [][]RunStep
innerAgent := &dtTestAgent{
name: "inner",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
ev := EventFromMessage(schema.AssistantMessage("from inner", nil), nil, schema.Assistant, "")
gen.Send(ev)
}()
return iter
},
}
innerFlowAgent := toFlowAgent(ctx, innerAgent)
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: innerFlowAgent,
ToAgentNames: []string{},
})
outerAgent := &dtTestAgent{
name: "outer",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
innerIter := wrapped.Run(ctx, input, options...)
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
for {
ev, ok := innerIter.Next()
if !ok {
break
}
collectedRunPaths = append(collectedRunPaths, ev.RunPath)
gen.Send(ev)
}
}()
return iter
},
}
outerFlowAgent := toFlowAgent(ctx, outerAgent)
runner := NewRunner(ctx, RunnerConfig{
Agent: outerFlowAgent,
EnableStreaming: true,
CheckPointStore: store,
})
iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, WithCheckPointID("cp1"))
for {
_, ok := iter.Next()
if !ok {
break
}
}
assert.NotEmpty(t, collectedRunPaths, "should have collected RunPaths")
for _, rp := range collectedRunPaths {
assert.Len(t, rp, 2, "RunPath should have 2 steps (outer agent, inner agent)")
if len(rp) == 2 {
assert.Equal(t, "outer", rp[0].agentName, "RunPath[0] should be outer agent")
assert.Equal(t, "inner", rp[1].agentName, "RunPath[1] should be inner agent")
}
}
}
func TestDeterministicTransferExitSkipsTransfer(t *testing.T) {
ctx := context.Background()
store := newDTTestStore()
innerAgent := &dtTestAgent{
name: "inner",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
ev := EventFromMessage(schema.AssistantMessage("inner exits", nil), nil, schema.Assistant, "")
ev.Action = &AgentAction{Exit: true}
gen.Send(ev)
}()
return iter
},
}
innerFlowAgent := toFlowAgent(ctx, innerAgent)
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: innerFlowAgent,
ToAgentNames: []string{"next_agent"},
})
var outerSawExit bool
var transferGenerated bool
outerAgent := &dtTestAgent{
name: "outer",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
innerIter := wrapped.Run(ctx, input, options...)
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
for {
ev, ok := innerIter.Next()
if !ok {
break
}
if ev.Action != nil && ev.Action.Exit {
outerSawExit = true
}
if ev.Action != nil && ev.Action.TransferToAgent != nil {
transferGenerated = true
}
gen.Send(ev)
}
}()
return iter
},
}
outerFlowAgent := toFlowAgent(ctx, outerAgent)
runner := NewRunner(ctx, RunnerConfig{
Agent: outerFlowAgent,
EnableStreaming: true,
CheckPointStore: store,
})
iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, WithCheckPointID("cp1"))
for {
_, ok := iter.Next()
if !ok {
break
}
}
assert.True(t, outerSawExit, "outer should see exit event from inner")
assert.False(t, transferGenerated, "transfer should not be generated when inner exits")
}
type nonFlowTestAgent struct {
name string
runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]
resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]
}
func (a *nonFlowTestAgent) Name(_ context.Context) string {
return a.name
}
func (a *nonFlowTestAgent) Description(_ context.Context) string {
return a.name + " description"
}
func (a *nonFlowTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
return a.runFn(ctx, input, options...)
}
func (a *nonFlowTestAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
if a.resumeFn != nil {
return a.resumeFn(ctx, info, opts...)
}
return a.runFn(ctx, &AgentInput{}, opts...)
}
type nonResumableTestAgent struct {
name string
runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]
}
func (a *nonResumableTestAgent) Name(_ context.Context) string {
return a.name
}
func (a *nonResumableTestAgent) Description(_ context.Context) string {
return a.name + " description"
}
func (a *nonResumableTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
return a.runFn(ctx, input, options...)
}
func TestDeterministicTransferNonFlowAgent_ExitSkipsTransfer(t *testing.T) {
ctx := context.Background()
agent := &nonFlowTestAgent{
name: "test_agent",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
ev := EventFromMessage(schema.AssistantMessage("exiting", nil), nil, schema.Assistant, "")
ev.Action = &AgentAction{Exit: true}
gen.Send(ev)
}()
return iter
},
}
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: agent,
ToAgentNames: []string{"next_agent"},
})
iter := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}})
var events []*AgentEvent
var sawExit bool
var sawTransfer bool
for {
ev, ok := iter.Next()
if !ok {
break
}
events = append(events, ev)
if ev.Action != nil && ev.Action.Exit {
sawExit = true
}
if ev.Action != nil && ev.Action.TransferToAgent != nil {
sawTransfer = true
}
}
assert.True(t, sawExit, "should see exit event")
assert.False(t, sawTransfer, "should NOT see transfer when exit is last event")
assert.Len(t, events, 1, "should have exactly 1 event (exit)")
}
func TestDeterministicTransferNonFlowAgent_AppendsTransfer(t *testing.T) {
ctx := context.Background()
agent := &nonFlowTestAgent{
name: "test_agent",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
ev := EventFromMessage(schema.AssistantMessage("normal output", nil), nil, schema.Assistant, "")
gen.Send(ev)
}()
return iter
},
}
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: agent,
ToAgentNames: []string{"next_agent"},
})
iter := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}})
var events []*AgentEvent
var sawTransfer bool
var transferTarget string
for {
ev, ok := iter.Next()
if !ok {
break
}
events = append(events, ev)
if ev.Action != nil && ev.Action.TransferToAgent != nil {
sawTransfer = true
transferTarget = ev.Action.TransferToAgent.DestAgentName
}
}
assert.True(t, sawTransfer, "should see transfer event after normal completion")
assert.Equal(t, "next_agent", transferTarget, "transfer target should be next_agent")
assert.Greater(t, len(events), 1, "should have more than 1 event (output + transfer messages)")
}
func TestDeterministicTransferNonFlowAgent_InterruptSkipsTransfer(t *testing.T) {
ctx := context.Background()
agent := &nonFlowTestAgent{
name: "test_agent",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
ev := &AgentEvent{
Action: &AgentAction{
Interrupted: &InterruptInfo{Data: "test interrupt"},
},
}
gen.Send(ev)
}()
return iter
},
}
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: agent,
ToAgentNames: []string{"next_agent"},
})
iter := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}})
var events []*AgentEvent
var sawInterrupt bool
var sawTransfer bool
for {
ev, ok := iter.Next()
if !ok {
break
}
events = append(events, ev)
if ev.Action != nil && ev.Action.Interrupted != nil {
sawInterrupt = true
}
if ev.Action != nil && ev.Action.TransferToAgent != nil {
sawTransfer = true
}
}
assert.True(t, sawInterrupt, "should see interrupt event")
assert.False(t, sawTransfer, "should NOT see transfer when interrupted")
assert.Len(t, events, 1, "should have exactly 1 event (interrupt)")
}
func TestDeterministicTransferNonFlowAgent_Resume(t *testing.T) {
ctx := context.Background()
var resumeCalled bool
agent := &nonFlowTestAgent{
name: "test_agent",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
ev := EventFromMessage(schema.AssistantMessage("from run", nil), nil, schema.Assistant, "")
gen.Send(ev)
}()
return iter
},
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
resumeCalled = true
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
ev := EventFromMessage(schema.AssistantMessage("from resume", nil), nil, schema.Assistant, "")
gen.Send(ev)
}()
return iter
},
}
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: agent,
ToAgentNames: []string{"next_agent"},
})
ra, ok := wrapped.(ResumableAgent)
assert.True(t, ok, "wrapped agent should be ResumableAgent")
iter := ra.Resume(ctx, &ResumeInfo{WasInterrupted: true})
var sawTransfer bool
for {
ev, ok := iter.Next()
if !ok {
break
}
if ev.Action != nil && ev.Action.TransferToAgent != nil {
sawTransfer = true
}
}
assert.True(t, resumeCalled, "resume should have been called on inner agent")
assert.True(t, sawTransfer, "should see transfer event after resume completion")
}
func TestDeterministicTransferFlowAgent_ResumeWithInvalidState(t *testing.T) {
ctx := context.Background()
innerAgent := &dtTestAgent{
name: "inner",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
gen.Send(EventFromMessage(schema.AssistantMessage("test", nil), nil, schema.Assistant, ""))
}()
return iter
},
}
innerFlowAgent := toFlowAgent(ctx, innerAgent)
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: innerFlowAgent,
ToAgentNames: []string{"next_agent"},
})
ra, ok := wrapped.(ResumableAgent)
assert.True(t, ok, "wrapped flowAgent should be ResumableAgent")
iter := ra.Resume(ctx, &ResumeInfo{
WasInterrupted: true,
InterruptState: nil,
})
var gotError bool
var errorMsg string
for {
ev, ok := iter.Next()
if !ok {
break
}
if ev.Err != nil {
gotError = true
errorMsg = ev.Err.Error()
}
}
assert.True(t, gotError, "should get error for invalid state")
assert.Contains(t, errorMsg, "invalid interrupt state", "error should mention invalid state")
}
func TestDeterministicTransferNonResumableAgent(t *testing.T) {
ctx := context.Background()
agent := &nonResumableTestAgent{
name: "non_resumable",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
ev := EventFromMessage(schema.AssistantMessage("output", nil), nil, schema.Assistant, "")
gen.Send(ev)
}()
return iter
},
}
wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{
Agent: agent,
ToAgentNames: []string{"next_agent"},
})
_, isResumable := wrapped.(ResumableAgent)
assert.False(t, isResumable, "wrapped non-resumable agent should NOT be ResumableAgent")
assert.Equal(t, "non_resumable", wrapped.Name(ctx), "Name should delegate to inner agent")
assert.Equal(t, "non_resumable description", wrapped.Description(ctx), "Description should delegate to inner agent")
iter := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}})
var sawTransfer bool
for {
ev, ok := iter.Next()
if !ok {
break
}
if ev.Action != nil && ev.Action.TransferToAgent != nil {
sawTransfer = true
}
}
assert.True(t, sawTransfer, "should see transfer event")
}
+490
View File
@@ -0,0 +1,490 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"errors"
"fmt"
"io"
"log"
"github.com/cloudwego/eino/components"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/compose"
"github.com/cloudwego/eino/schema"
)
type failoverCurrentModelKey struct{}
func typedSetFailoverCurrentModel[M MessageType](ctx context.Context, currentModel model.BaseModel[M]) context.Context {
return context.WithValue(ctx, failoverCurrentModelKey{}, currentModel)
}
func typedGetFailoverCurrentModel[M MessageType](ctx context.Context) (model.BaseModel[M], bool) {
m, ok := ctx.Value(failoverCurrentModelKey{}).(model.BaseModel[M])
return m, ok
}
type failoverHasMoreAttemptsKey struct{}
// withFailoverHasMoreAttempts sets a flag in context indicating whether additional failover
// attempts remain after the current one. This is read by buildErrWrapper to decide whether
// stream errors should be wrapped as WillRetryError.
func withFailoverHasMoreAttempts(ctx context.Context, hasMore bool) context.Context {
return context.WithValue(ctx, failoverHasMoreAttemptsKey{}, hasMore)
}
// getFailoverHasMoreAttempts returns true if the current failover attempt has more attempts
// after it, false otherwise (including when no failover context is present).
func getFailoverHasMoreAttempts(ctx context.Context) bool {
v, _ := ctx.Value(failoverHasMoreAttemptsKey{}).(bool)
return v
}
type typedFailoverProxyModel[M MessageType] struct {
}
func (m *typedFailoverProxyModel[M]) prepareTarget(ctx context.Context) (model.BaseModel[M], error) {
target, ok := typedGetFailoverCurrentModel[M](ctx)
if !ok {
return nil, errors.New("failover current model not found in context")
}
if !components.IsCallbacksEnabled(target) {
target = typedCallbackInjectionModelWrapper[M]{}.wrapModel(target)
}
return target, nil
}
func (m *typedFailoverProxyModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
target, err := m.prepareTarget(ctx)
if err != nil {
var zero M
return zero, err
}
return target.Generate(ctx, input, opts...)
}
func (m *typedFailoverProxyModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
target, err := m.prepareTarget(ctx)
if err != nil {
return nil, err
}
return target.Stream(ctx, input, opts...)
}
func (m *typedFailoverProxyModel[M]) IsCallbacksEnabled() bool {
return true
}
func (m *typedFailoverProxyModel[M]) GetType() string {
return "FailoverProxyModel"
}
type failoverProxyModel = typedFailoverProxyModel[*schema.Message]
// FailoverContext contains context information during failover process.
type FailoverContext[M MessageType] struct {
// FailoverAttempt is the current failover attempt number, starting from 1.
FailoverAttempt uint
// InputMessages is the original input messages before any transformation.
InputMessages []M
// LastOutputMessage is the output message from the last failed attempt.
// May be nil if no output was produced. For streaming, this may be a partial message
// already received before the stream error.
LastOutputMessage M
// LastErr is the error from the last failed attempt that triggered this failover.
//
// Note: When ModelRetryConfig is also configured, LastErr will be a *RetryExhaustedError
// (if retries were exhausted) rather than the original model error. The original error
// can be retrieved via RetryExhaustedError.LastErr.
LastErr error
}
// ModelFailoverConfig configures failover behavior for ChatModel.
// When configured, each ChatModel call first tries the last successful model (initially the configured Model),
// and if that fails, calls GetFailoverModel to select alternate models.
type ModelFailoverConfig[M MessageType] struct {
// MaxRetries specifies the maximum number of failover attempts.
//
// When failover is triggered, GetFailoverModel will be called up to MaxRetries times
// (FailoverAttempt starts from 1). If GetFailoverModel returns an error, failover
// stops immediately and that error is returned.
//
// A value of 0 means no failover (GetFailoverModel will not be called).
// A value of 1 means GetFailoverModel may be called once.
//
// Note: if lastSuccessModel is set (from a previous successful call), it will be tried
// first before calling GetFailoverModel.
MaxRetries uint
// ShouldFailover determines whether to fail over to the next model when an error occurs.
// It receives the output message (may be nil/zero if no output is available) and the error (non-nil on failure).
// For streaming errors, outputMessage can carry a partial message accumulated before the error.
//
// Note: When ModelRetryConfig is also configured, outputErr will be a *RetryExhaustedError
// (if retries were exhausted) rather than the original model error. Use errors.As to extract
// the RetryExhaustedError and access RetryExhaustedError.LastErr for the original error.
//
// Note: When the context itself is cancelled (ctx.Err() != nil), failover will stop immediately
// regardless of this function. However, if the model returns context.Canceled or context.DeadlineExceeded
// as an error while the context is still active, this function will still be called.
// Should not be nil when ModelFailoverConfig is set.
// Return true to fail over to the next model, false to stop and return the current result/error.
ShouldFailover func(ctx context.Context, outputMessage M, outputErr error) bool
// GetFailoverModel is called when a model call fails and ShouldFailover returns true.
// It selects the next model to use for the failover attempt and optionally transforms input messages.
// It receives the failover context containing attempt number (starting from 1), original input, and last result.
// Return values:
// - failoverModel: The model to use for this failover attempt.
// - failoverModelInputMessages: The transformed input messages for the failover model. If nil, will use original input.
// - failoverErr: If non-nil, failover stops and this error is returned.
// Should not be nil when ModelFailoverConfig is set via ChatModelAgentConfig.
GetFailoverModel func(ctx context.Context, failoverCtx *FailoverContext[M]) (
failoverModel model.BaseModel[M], failoverModelInputMessages []M, failoverErr error)
}
func typedGetFailoverLastSuccessModel[M MessageType](ctx context.Context) model.BaseModel[M] {
execCtx := getTypedChatModelAgentExecCtx[M](ctx)
if execCtx == nil {
return nil
}
return execCtx.failoverLastSuccessModel
}
func typedSetFailoverLastSuccessModel[M MessageType](ctx context.Context, m model.BaseModel[M]) {
if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil {
execCtx.failoverLastSuccessModel = m
}
}
type failoverModelWrapper[M MessageType] struct {
config *ModelFailoverConfig[M]
inner model.BaseModel[M]
}
func newFailoverModelWrapper[M MessageType](inner model.BaseModel[M], config *ModelFailoverConfig[M]) *failoverModelWrapper[M] {
return &failoverModelWrapper[M]{
config: config,
inner: inner,
}
}
func (f *failoverModelWrapper[M]) needFailover(ctx context.Context, outputMessage M, outputErr error) bool {
if ctx.Err() != nil {
return false
}
if _, ok := compose.ExtractInterruptInfo(outputErr); ok {
return false
}
// ErrStreamCanceled means the caller voluntarily abandoned the stream;
// never retry or fail over in this case.
if errors.Is(outputErr, ErrStreamCanceled) {
return false
}
// ShouldFailover is validated at agent construction; nil here indicates a programmer error.
return f.config.ShouldFailover(ctx, outputMessage, outputErr)
}
func (f *failoverModelWrapper[M]) getFailoverModel(ctx context.Context, failoverCtx *FailoverContext[M]) (model.BaseModel[M], []M, error) {
currentModel, msgs, err := f.config.GetFailoverModel(ctx, failoverCtx)
if err != nil {
return nil, nil, err
}
if currentModel == nil {
return nil, nil, nil
}
return currentModel, msgs, nil
}
func (f *failoverModelWrapper[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
// Defensive: GetFailoverModel is validated non-nil at agent construction.
if f.config.GetFailoverModel == nil {
return f.inner.Generate(ctx, input, opts...)
}
var lastOutputMessage M
var lastErr error
// Try lastSuccessModel first if available.
if lastSuccess := typedGetFailoverLastSuccessModel[M](ctx); lastSuccess != nil {
if err := ctx.Err(); err != nil {
var zero M
return zero, err
}
modelCtx := typedSetFailoverCurrentModel(ctx, lastSuccess)
modelCtx = withFailoverHasMoreAttempts(modelCtx, f.config.MaxRetries > 0)
result, err := f.inner.Generate(modelCtx, input, opts...)
if err == nil {
return result, nil
}
lastOutputMessage = result
lastErr = err
if !f.needFailover(ctx, result, err) {
return result, err
}
log.Printf("failover ChatModel.Generate lastSuccessModel failed: %v", err)
}
for attempt := uint(1); attempt <= f.config.MaxRetries; attempt++ {
if err := ctx.Err(); err != nil {
var zero M
return zero, err
}
failoverCtx := &FailoverContext[M]{
FailoverAttempt: attempt,
InputMessages: input,
LastOutputMessage: lastOutputMessage,
LastErr: lastErr,
}
currentModel, currentInput, err := f.getFailoverModel(ctx, failoverCtx)
if err != nil {
var zero M
return zero, err
}
if currentModel == nil {
var zero M
return zero, fmt.Errorf("failover GetFailoverModel returned nil model at attempt %d", attempt)
}
if currentInput == nil {
currentInput = input
}
modelCtx := typedSetFailoverCurrentModel(ctx, currentModel)
modelCtx = withFailoverHasMoreAttempts(modelCtx, attempt < f.config.MaxRetries)
result, err := f.inner.Generate(modelCtx, currentInput, opts...)
lastOutputMessage = result
lastErr = err
if err == nil {
typedSetFailoverLastSuccessModel(ctx, currentModel)
return result, nil
}
if !f.needFailover(ctx, result, err) {
return result, err
}
if attempt < f.config.MaxRetries {
log.Printf("failover ChatModel.Generate attempt %d failed: %v", attempt, err)
}
}
return lastOutputMessage, lastErr
}
func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (
*schema.StreamReader[M], error) {
// Defensive: GetFailoverModel is validated non-nil at agent construction.
if f.config.GetFailoverModel == nil {
return f.inner.Stream(ctx, input, opts...)
}
var lastOutputMessage M
var lastErr error
// Try lastSuccessModel first if available.
if lastSuccess := typedGetFailoverLastSuccessModel[M](ctx); lastSuccess != nil {
if err := ctx.Err(); err != nil {
return nil, err
}
modelCtx := typedSetFailoverCurrentModel(ctx, lastSuccess)
modelCtx = withFailoverHasMoreAttempts(modelCtx, f.config.MaxRetries > 0)
stream, err := f.inner.Stream(modelCtx, input, opts...)
if err != nil {
lastErr = err
var zero M
if !f.needFailover(ctx, zero, err) {
return nil, err
}
log.Printf("failover ChatModel.Stream lastSuccessModel failed: %v", err)
} else {
copies := stream.Copy(2)
checkCopy := copies[0]
returnCopy := copies[1]
outMsg, streamErr := typedConsumeStream(checkCopy)
if streamErr != nil {
lastOutputMessage = outMsg
lastErr = streamErr
returnCopy.Close()
if !f.needFailover(ctx, outMsg, streamErr) {
return nil, streamErr
}
log.Printf("failover ChatModel.Stream lastSuccessModel failed: %v", streamErr)
} else {
return returnCopy, nil
}
}
}
for attempt := uint(1); attempt <= f.config.MaxRetries; attempt++ {
if err := ctx.Err(); err != nil {
return nil, err
}
failoverCtx := &FailoverContext[M]{
FailoverAttempt: attempt,
InputMessages: input,
LastOutputMessage: lastOutputMessage,
LastErr: lastErr,
}
currentModel, currentInput, err := f.getFailoverModel(ctx, failoverCtx)
if err != nil {
return nil, err
}
if currentModel == nil {
return nil, fmt.Errorf("failover GetFailoverModel returned nil model at attempt %d", attempt)
}
if currentInput == nil {
currentInput = input
}
modelCtx := typedSetFailoverCurrentModel(ctx, currentModel)
modelCtx = withFailoverHasMoreAttempts(modelCtx, attempt < f.config.MaxRetries)
stream, err := f.inner.Stream(modelCtx, currentInput, opts...)
if err != nil {
lastErr = err
var zero M
lastOutputMessage = zero
if !f.needFailover(ctx, zero, err) {
return nil, err
}
if attempt < f.config.MaxRetries {
log.Printf("failover ChatModel.Stream attempt %d failed: %v", attempt, err)
}
continue
}
// The stream returned by f.inner.Stream is already Copy'd by the inner eventSender layer: one
// copy is forwarded to the client in real time via events. Therefore consuming a copy here does
// NOT block client-side streaming.
//
// We Copy the stream into two readers:
// - checkCopy: consumed synchronously to surface mid-stream errors and decide whether to fail over.
// - returnCopy: returned to the caller (stateModelWrapper), which also consumes synchronously to
// build state (AfterModelRewriteState), so waiting here adds no extra latency.
//
// If checkCopy errors and failover is allowed, we close returnCopy and retry with the next model.
// Otherwise we return returnCopy.
//
// NOTE on duplicate events during failover: when a retry happens, events from the failed attempt
// may already have been emitted to the client, and the retry will emit a new stream. Client-side
// handlers are expected to handle multiple rounds (e.g., reset on retry or deduplicate by attempt
// metadata).
copies := stream.Copy(2)
checkCopy := copies[0]
returnCopy := copies[1]
outMsg, streamErr := typedConsumeStream(checkCopy)
if streamErr != nil {
lastOutputMessage = outMsg
lastErr = streamErr
returnCopy.Close()
if !f.needFailover(ctx, outMsg, streamErr) {
return nil, streamErr
}
if attempt < f.config.MaxRetries {
log.Printf("failover ChatModel.Stream attempt %d failed: %v", attempt, streamErr)
}
continue
}
typedSetFailoverLastSuccessModel(ctx, currentModel)
return returnCopy, nil
}
return nil, lastErr
}
func typedConsumeStream[M MessageType](stream *schema.StreamReader[M]) (M, error) {
var zero M
defer stream.Close()
switch s := any(stream).(type) {
case *schema.StreamReader[*schema.Message]:
chunks := make([]*schema.Message, 0)
for {
chunk, err := s.Recv()
if err == io.EOF {
break
}
if err != nil {
msg, _ := schema.ConcatMessages(chunks)
if msg != nil {
return any(msg).(M), err
}
return zero, err
}
chunks = append(chunks, chunk)
}
msg, _ := schema.ConcatMessages(chunks)
if msg != nil {
return any(msg).(M), nil
}
return zero, nil
case *schema.StreamReader[*schema.AgenticMessage]:
chunks := make([]*schema.AgenticMessage, 0)
for {
chunk, err := s.Recv()
if err == io.EOF {
break
}
if err != nil {
msg, _ := schema.ConcatAgenticMessages(chunks)
if msg != nil {
return any(msg).(M), err
}
return zero, err
}
chunks = append(chunks, chunk)
}
msg, _ := schema.ConcatAgenticMessages(chunks)
if msg != nil {
return any(msg).(M), nil
}
return zero, nil
default:
panic("unreachable: unknown MessageType")
}
}
+742
View File
@@ -0,0 +1,742 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"errors"
"io"
"sync/atomic"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/schema"
)
type fakeChatModel struct {
generate func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error)
stream func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error)
callbacksEnabled bool
}
func (m *fakeChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
return m.generate(ctx, input, opts...)
}
func (m *fakeChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return m.stream(ctx, input, opts...)
}
func (m *fakeChatModel) IsCallbacksEnabled() bool {
return m.callbacksEnabled
}
func drainMessageStream(sr *schema.StreamReader[*schema.Message]) ([]*schema.Message, error) {
defer sr.Close()
var out []*schema.Message
for {
chunk, err := sr.Recv()
if err == io.EOF {
return out, nil
}
if err != nil {
return out, err
}
out = append(out, chunk)
}
}
func streamWithMidError(chunks []*schema.Message, err error) *schema.StreamReader[*schema.Message] {
sr, sw := schema.Pipe[*schema.Message](2)
go func() {
defer sw.Close()
for _, c := range chunks {
sw.Send(c, nil)
}
sw.Send(nil, err)
}()
return sr
}
func streamWithMidErrorControlled(chunks []*schema.Message, err error, firstSent chan struct{}, release chan struct{}) *schema.StreamReader[*schema.Message] {
sr, sw := schema.Pipe[*schema.Message](2)
go func() {
defer sw.Close()
for i, c := range chunks {
sw.Send(c, nil)
if i == 0 && firstSent != nil {
close(firstSent)
if release != nil {
<-release
}
}
}
sw.Send(nil, err)
}()
return sr
}
func TestFailoverCurrentModelContext(t *testing.T) {
t.Run("set and get", func(t *testing.T) {
ctx := context.Background()
m := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return schema.AssistantMessage("ok", nil), nil
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil
},
}
ctx = typedSetFailoverCurrentModel[*schema.Message](ctx, m)
got, ok := typedGetFailoverCurrentModel[*schema.Message](ctx)
require.True(t, ok)
require.Same(t, m, got)
})
t.Run("wrong type", func(t *testing.T) {
ctx := context.WithValue(context.Background(), failoverCurrentModelKey{}, "bad")
_, ok := typedGetFailoverCurrentModel[*schema.Message](ctx)
require.False(t, ok)
})
t.Run("missing", func(t *testing.T) {
_, ok := typedGetFailoverCurrentModel[*schema.Message](context.Background())
require.False(t, ok)
})
}
func TestFailoverProxyModel(t *testing.T) {
t.Run("generate missing context", func(t *testing.T) {
p := &failoverProxyModel{}
_, err := p.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
require.Error(t, err)
})
t.Run("stream missing context", func(t *testing.T) {
p := &failoverProxyModel{}
_, err := p.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
require.Error(t, err)
})
t.Run("generate routes to current model", func(t *testing.T) {
var called int32
target := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
atomic.AddInt32(&called, 1)
return schema.AssistantMessage("routed", nil), nil
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("routed", nil)}), nil
},
}
ctx := typedSetFailoverCurrentModel[*schema.Message](context.Background(), target)
p := &failoverProxyModel{}
msg, err := p.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
require.NoError(t, err)
require.Equal(t, "routed", msg.Content)
require.Equal(t, int32(1), atomic.LoadInt32(&called))
})
}
func TestFailoverModelWrapper_Generate(t *testing.T) {
t.Run("delegates when GetFailoverModel nil", func(t *testing.T) {
var called int32
inner := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
atomic.AddInt32(&called, 1)
return schema.AssistantMessage("inner", nil), nil
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("inner", nil)}), nil
},
}
w := newFailoverModelWrapper[*schema.Message](inner, &ModelFailoverConfig[*schema.Message]{
MaxRetries: 2,
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
GetFailoverModel: nil,
})
msg, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
require.NoError(t, err)
require.Equal(t, "inner", msg.Content)
require.Equal(t, int32(1), atomic.LoadInt32(&called))
})
t.Run("failover to second model", func(t *testing.T) {
wantErr := errors.New("first failed")
var shouldCalls int32
var m1Calls int32
var m2Calls int32
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
atomic.AddInt32(&m1Calls, 1)
return nil, wantErr
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return nil, errors.New("unused")
},
}
m2 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
atomic.AddInt32(&m2Calls, 1)
return schema.AssistantMessage("ok", nil), nil
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return nil, errors.New("unused")
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 1,
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
atomic.AddInt32(&shouldCalls, 1)
return errors.Is(err, wantErr)
},
GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
require.Equal(t, uint(1), failoverCtx.FailoverAttempt)
return m2, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
msg, err := w.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
require.NoError(t, err)
require.Equal(t, "ok", msg.Content)
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls))
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
})
t.Run("canceled error delegates to ShouldFailover", func(t *testing.T) {
var shouldCalls int32
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, context.Canceled
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return nil, errors.New("unused")
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 5,
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
atomic.AddInt32(&shouldCalls, 1)
// User decides to stop on canceled error
return !errors.Is(err, context.Canceled)
},
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return m1, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
_, err := w.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
require.ErrorIs(t, err, context.Canceled)
// ShouldFailover is called once and returns false, stopping failover
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
})
t.Run("stops when GetFailoverModel returns error", func(t *testing.T) {
wantErr := errors.New("get model failed")
var called int32
inner := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
atomic.AddInt32(&called, 1)
return schema.AssistantMessage("unused", nil), nil
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return nil, errors.New("unused")
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 3,
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return nil, nil, wantErr
},
}
w := newFailoverModelWrapper[*schema.Message](inner, cfg)
_, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
require.ErrorIs(t, err, wantErr)
require.Equal(t, int32(0), atomic.LoadInt32(&called))
})
t.Run("stops when GetFailoverModel returns nil model", func(t *testing.T) {
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 1,
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return nil, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
msg, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
require.Nil(t, msg)
require.Error(t, err)
require.ErrorContains(t, err, "GetFailoverModel returned nil model")
})
}
func TestFailoverModelWrapper_Stream(t *testing.T) {
t.Run("returns stream when first attempt succeeds", func(t *testing.T) {
var shouldCalls int32
in := schema.UserMessage("hi")
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
require.Len(t, input, 1)
require.Same(t, in, input[0])
return schema.StreamReaderFromArray([]*schema.Message{
schema.AssistantMessage("a", nil),
schema.AssistantMessage("b", nil),
}), nil
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 0,
ShouldFailover: func(context.Context, *schema.Message, error) bool {
atomic.AddInt32(&shouldCalls, 1)
return false
},
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return m1, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
sr, err := w.Stream(ctx, []*schema.Message{in})
require.NoError(t, err)
msgs, err := drainMessageStream(sr)
require.NoError(t, err)
require.Len(t, msgs, 2)
require.Equal(t, "a", msgs[0].Content)
require.Equal(t, "b", msgs[1].Content)
require.Equal(t, int32(0), atomic.LoadInt32(&shouldCalls))
})
t.Run("failover when Stream returns error immediately", func(t *testing.T) {
wantErr := errors.New("stream init failed")
var shouldCalls int32
var m1Calls int32
var m2Calls int32
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
atomic.AddInt32(&m1Calls, 1)
return nil, wantErr
},
}
m2 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
atomic.AddInt32(&m2Calls, 1)
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 1,
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
atomic.AddInt32(&shouldCalls, 1)
return errors.Is(err, wantErr)
},
GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
require.Equal(t, uint(1), failoverCtx.FailoverAttempt)
return m2, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
require.NoError(t, err)
msgs, err := drainMessageStream(sr)
require.NoError(t, err)
require.Len(t, msgs, 1)
require.Equal(t, "ok", msgs[0].Content)
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls))
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
})
t.Run("failover when stream errors mid-way", func(t *testing.T) {
streamErr := errors.New("mid error")
var shouldCalls int32
var seenOutput atomic.Value
var m1Calls int32
var m2Calls int32
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
atomic.AddInt32(&m1Calls, 1)
return streamWithMidError([]*schema.Message{
schema.AssistantMessage("p1", nil),
schema.AssistantMessage("p2", nil),
}, streamErr), nil
},
}
m2 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
atomic.AddInt32(&m2Calls, 1)
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("final", nil)}), nil
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 1,
ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool {
atomic.AddInt32(&shouldCalls, 1)
if errors.Is(err, streamErr) && out != nil {
seenOutput.Store(out.Content)
}
return errors.Is(err, streamErr)
},
GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
require.Equal(t, uint(1), failoverCtx.FailoverAttempt)
return m2, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
require.NoError(t, err)
msgs, err := drainMessageStream(sr)
require.NoError(t, err)
require.Len(t, msgs, 1)
require.Equal(t, "final", msgs[0].Content)
require.Equal(t, "p1p2", seenOutput.Load())
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls))
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
})
t.Run("stop when ShouldFailover returns false for mid-way error", func(t *testing.T) {
streamErr := errors.New("mid error")
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return streamWithMidError([]*schema.Message{schema.AssistantMessage("p", nil)}, streamErr), nil
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 3,
ShouldFailover: func(context.Context, *schema.Message, error) bool {
return false
},
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return m1, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
require.Nil(t, sr)
require.ErrorIs(t, err, streamErr)
})
t.Run("canceled mid-way error delegates to ShouldFailover", func(t *testing.T) {
var shouldCalls int32
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
return streamWithMidError([]*schema.Message{schema.AssistantMessage("p", nil)}, context.Canceled), nil
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 3,
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
atomic.AddInt32(&shouldCalls, 1)
// User decides to stop on canceled error
return !errors.Is(err, context.Canceled)
},
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return m1, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
require.Nil(t, sr)
require.ErrorIs(t, err, context.Canceled)
// ShouldFailover is called once and returns false, stopping failover
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
})
t.Run("stop when Stream returns error immediately and ShouldFailover returns false", func(t *testing.T) {
wantErr := errors.New("stream init failed")
var shouldCalls int32
var m1Calls int32
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
atomic.AddInt32(&m1Calls, 1)
return nil, wantErr
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 3,
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
atomic.AddInt32(&shouldCalls, 1)
require.ErrorIs(t, err, wantErr)
return false
},
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return m1, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
require.Nil(t, sr)
require.ErrorIs(t, err, wantErr)
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
})
t.Run("stops when GetFailoverModel returns nil model", func(t *testing.T) {
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 1,
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return nil, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
sr, err := w.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
require.Nil(t, sr)
require.Error(t, err)
require.ErrorContains(t, err, "GetFailoverModel returned nil model")
})
t.Run("stops when GetFailoverModel returns error", func(t *testing.T) {
wantErr := errors.New("get model failed")
var called int32
inner := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
atomic.AddInt32(&called, 1)
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 3,
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
return nil, nil, wantErr
},
}
w := newFailoverModelWrapper[*schema.Message](inner, cfg)
sr, err := w.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
require.Nil(t, sr)
require.ErrorIs(t, err, wantErr)
require.Equal(t, int32(0), atomic.LoadInt32(&called))
})
t.Run("stops when ctx canceled during mid-way error handling", func(t *testing.T) {
midErr := errors.New("mid error")
var shouldCalls int32
var m1Calls int32
var m2Calls int32
firstSent := make(chan struct{})
release := make(chan struct{})
m1 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
atomic.AddInt32(&m1Calls, 1)
return streamWithMidErrorControlled(
[]*schema.Message{schema.AssistantMessage("p", nil)},
midErr,
firstSent,
release,
), nil
},
}
m2 := &fakeChatModel{
callbacksEnabled: true,
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
return nil, errors.New("unused")
},
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
atomic.AddInt32(&m2Calls, 1)
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil
},
}
cfg := &ModelFailoverConfig[*schema.Message]{
MaxRetries: 1,
ShouldFailover: func(context.Context, *schema.Message, error) bool {
atomic.AddInt32(&shouldCalls, 1)
return true
},
GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
require.Equal(t, uint(1), failoverCtx.FailoverAttempt)
return m2, nil, nil
},
}
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
baseCtx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
failoverLastSuccessModel: m1,
})
ctx, cancel := context.WithCancel(baseCtx)
type result struct {
sr *schema.StreamReader[*schema.Message]
err error
}
ch := make(chan result, 1)
go func() {
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
ch <- result{sr: sr, err: err}
}()
<-firstSent
cancel()
close(release)
res := <-ch
if res.sr != nil {
res.sr.Close()
}
require.Nil(t, res.sr)
require.ErrorIs(t, res.err, midErr)
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
require.Equal(t, int32(0), atomic.LoadInt32(&m2Calls))
require.Equal(t, int32(0), atomic.LoadInt32(&shouldCalls))
})
}
func TestTypedConsumeStream_EmptyAgenticStream(t *testing.T) {
sr, sw := schema.Pipe[*schema.AgenticMessage](1)
sw.Close()
msg, err := typedConsumeStream(sr)
assert.Nil(t, err, "empty stream should not return error")
assert.NotNil(t, msg, "empty stream should return non-nil message from ConcatAgenticMessages")
}
func TestTypedConsumeStream_AgenticMidStreamError(t *testing.T) {
midErr := errors.New("mid-stream failure")
sr := streamWithMidErrorAgentic(
[]*schema.AgenticMessage{agenticChunk("chunk1"), agenticChunk("chunk2")},
midErr,
)
msg, err := typedConsumeStream(sr)
assert.ErrorIs(t, err, midErr, "should return the mid-stream error")
assert.NotNil(t, msg, "should return concatenated partial message from received chunks")
}
func streamWithMidErrorAgentic(chunks []*schema.AgenticMessage, err error) *schema.StreamReader[*schema.AgenticMessage] {
sr, sw := schema.Pipe[*schema.AgenticMessage](len(chunks) + 1)
go func() {
defer sw.Close()
for _, c := range chunks {
sw.Send(c, nil)
}
sw.Send(nil, err)
}()
return sr
}
func agenticChunk(text string) *schema.AgenticMessage {
return &schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.AssistantGenText{Text: text}),
},
}
}
+304
View File
@@ -0,0 +1,304 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Package filesystem provides file system operations.
package filesystem
import (
"context"
"github.com/cloudwego/eino/schema"
)
// FileInfo represents basic file metadata information.
type FileInfo struct {
// Path is the path of the file or directory, which can be a filename, relative path, or absolute path.
Path string
// IsDir indicates whether the entry is a directory.
// true for directories, false for regular files.
IsDir bool
// Size is the file size in bytes.
// For directories, this value may be 0 or platform-dependent.
Size int64
// ModifiedAt is the last modification time in ISO 8601 format.
// Example: "2025-01-15T10:30:00Z"
ModifiedAt string
}
// GrepMatch represents a single pattern match result.
type GrepMatch struct {
Content string
// Path is the file path where the match was found.
Path string
// Line is the 1-based line number of the match.
Line int
}
// LsInfoRequest contains parameters for listing file information.
type LsInfoRequest struct {
// Path specifies the directory path to list.
Path string
}
// ReadRequest contains parameters for reading file content.
type ReadRequest struct {
// FilePath is the path to the file to be read.
FilePath string
// Offset specifies the starting line number (1-based) for reading.
// Line 1 is the first line of the file.
// Use this when the file is too large to read at once.
// Defaults to 1 (start from the first line).
// Values < 1 will be treated as 1.
Offset int
// Limit specifies the maximum number of lines to read.
// When Limit is 0 (default), the entire file content is returned.
Limit int
}
// MultiModalReadRequest extends ReadRequest with parameters only applicable
// to MultiModalReader implementations (e.g. PDF page ranges).
type MultiModalReadRequest struct {
ReadRequest
// Pages specifies the page range for PDF files (e.g. "1-5", "3", "10-20").
Pages string
}
// GrepRequest contains parameters for searching file content.
type GrepRequest struct {
// ===== Search Parameters =====
// Pattern is the search pattern, supports full regular expression syntax.
// Uses ripgrep syntax (not grep). Examples:
// - "log.*Error" matches lines with "log" followed by "Error"
// - "function\\s+\\w+" matches "function" followed by whitespace and word characters
// - Literal braces need escaping: "interface\\{\\}" matches "interface{}"
Pattern string
// Path is an optional directory path to limit the search scope.
Path string
// ===== File Filtering =====
// Glob is an optional pattern to filter the files to be searched.
// It filters by file path, not content. If empty, no files are filtered.
// Supports standard glob wildcards:
// - `*` matches any characters except path separators.
// - `**` matches any directories recursively.
// - `?` matches a single character.
// - `[abc]` matches one character from the set.
Glob string
// FileType is the file type filter, e.g., "js", "py", "rust".
// More efficient than Glob for standard file types.
FileType string
// ===== Search Options =====
// CaseInsensitive enables case insensitive search.
CaseInsensitive bool
// EnableMultiline enables multiline mode where patterns can span lines.
// Default: false (patterns match within single lines only).
EnableMultiline bool
// ===== Context Display (Content mode only) =====
// AfterLines shows N lines after each match.
// Only applicable when OutputMode is "content".
// Values <= 0 are treated as unset.
AfterLines int
// BeforeLines shows N lines before each match.
// Only applicable when OutputMode is "content".
// Values <= 0 are treated as unset.
BeforeLines int
}
// GlobInfoRequest contains parameters for glob pattern matching.
type GlobInfoRequest struct {
// Pattern is the glob expression used to match file paths.
// It supports standard glob syntax:
// - `*` matches any characters except path separators.
// - `**` matches any directories recursively.
// - `?` matches a single character.
// - `[abc]` matches one character from the set.
Pattern string
// Path is the base directory from which to start the search.
Path string
}
// WriteRequest contains parameters for writing file content.
type WriteRequest struct {
// FilePath is the path of the file to write.
FilePath string
// Content is the data to be written to the file.
Content string
}
// EditRequest contains parameters for editing file content.
type EditRequest struct {
// FilePath is the path of the file to edit.
FilePath string
// OldString is the exact string to be replaced. It must be non-empty and will be matched literally, including whitespace.
OldString string
// NewString is the string that will replace OldString.
// It must be different from OldString.
// An empty string can be used to effectively delete OldString.
NewString string
// ReplaceAll controls the replacement behavior.
// If true, all occurrences of OldString are replaced.
// If false, the operation fails unless OldString appears exactly once in the file.
ReplaceAll bool
}
// FileContentPartType defines the type of a multimodal file content part.
type FileContentPartType string
const (
// FileContentPartTypeImage represents an image part (e.g. PNG, JPG).
FileContentPartTypeImage FileContentPartType = "image"
// FileContentPartTypePDF represents a file part (e.g. PDF).
FileContentPartTypePDF FileContentPartType = "pdf"
)
// FileContentPart represents a multimodal part of file content.
// Data holds raw bytes; encoding (e.g. base64) is handled by the consumer.
type FileContentPart struct {
// Type is the kind of content this part represents.
// Required.
Type FileContentPartType
// MIMEType is the MIME type of the content (e.g. "image/png", "application/pdf").
// Required.
MIMEType string
// Data is the raw binary content.
// Required.
Data []byte
}
// FileContent holds the result of a Read operation.
type FileContent struct {
// Content holds the plain text content of the file.
Content string
}
// MultiFileContent holds the result of a MultiModalRead operation.
//
// FileContent and Parts are mutually exclusive (one-of):
// - Set FileContent for plain text results (same as a normal Read).
// - Set Parts for multimodal results (images, PDFs, etc.).
//
// When Parts is non-empty, FileContent is ignored.
type MultiFileContent struct {
*FileContent
// Parts holds multimodal output parts (e.g. image, PDF).
Parts []FileContentPart
}
// MultiModalReader is an optional extension interface for Backend.
// Backends that implement this interface support multimodal file reading,
// returning structured parts (images, PDFs) instead of plain text.
//
// For large file handling, there are two approaches to control output size:
// - Implement size control within MultiModalRead (e.g. reject files exceeding a threshold,
// downsample images, or limit PDF page counts at the backend level).
// - Use ToolMiddleware's EnhancedInvokable to customize result transformation,
// or use the built-in reduction middleware with configurable policies.
type MultiModalReader interface {
MultiModalRead(ctx context.Context, req *MultiModalReadRequest) (*MultiFileContent, error)
}
// Backend is a pluggable, unified file backend protocol interface.
//
// All methods use struct-based parameters to allow future extensibility
// without breaking backward compatibility.
type Backend interface {
// LsInfo lists file information under the given path.
//
// Returns:
// - []FileInfo: List of matching file information
// - error: Error if the operation fails
LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error)
// Read reads file content with support for line-based offset and limit.
//
// Returns:
// - string: The file content read
// - error: Error if file does not exist or read fails
Read(ctx context.Context, req *ReadRequest) (*FileContent, error)
// GrepRaw searches for content matching the specified pattern in files.
//
// Returns:
// - []GrepMatch: List of all matching results
// - error: Error if the search fails
GrepRaw(ctx context.Context, req *GrepRequest) ([]GrepMatch, error)
// GlobInfo returns file information matching the glob pattern.
//
// Returns:
// - []FileInfo: List of matching file information
// - error: Error if the pattern is invalid or operation fails
GlobInfo(ctx context.Context, req *GlobInfoRequest) ([]FileInfo, error)
// Write creates or updates file content.
//
// Returns:
// - error: Error if the write operation fails
Write(ctx context.Context, req *WriteRequest) error
// Edit replaces string occurrences in a file.
//
// Returns:
// - error: Error if file does not exist, OldString is empty, or OldString is not found
Edit(ctx context.Context, req *EditRequest) error
}
// ExecuteRequest contains parameters for executing a command.
type ExecuteRequest struct {
Command string // The command to execute
RunInBackendGround bool
}
// ExecuteResponse contains the response result of command execution.
type ExecuteResponse struct {
Output string // Command output content
ExitCode *int // Command exit code
Truncated bool // Whether the output was truncated
}
type Shell interface {
Execute(ctx context.Context, input *ExecuteRequest) (result *ExecuteResponse, err error)
}
type StreamingShell interface {
ExecuteStreaming(ctx context.Context, input *ExecuteRequest) (result *schema.StreamReader[*ExecuteResponse], err error)
}
+774
View File
@@ -0,0 +1,774 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package filesystem
import (
"context"
"fmt"
"path/filepath"
"regexp"
"strings"
"sync"
"time"
"github.com/bmatcuk/doublestar/v4"
)
type fileEntry struct {
content string
modifiedAt time.Time
}
// InMemoryBackend is an in-memory implementation of the Backend interface.
// It stores files in a map and is safe for concurrent use.
type InMemoryBackend struct {
mu sync.RWMutex
files map[string]*fileEntry
}
// NewInMemoryBackend creates a new in-memory backend.
func NewInMemoryBackend() *InMemoryBackend {
return &InMemoryBackend{
files: make(map[string]*fileEntry),
}
}
// LsInfo lists file information under the given path.
func (b *InMemoryBackend) LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error) {
b.mu.RLock()
defer b.mu.RUnlock()
// Normalize path
path := normalizePath(req.Path)
var result []FileInfo
seen := make(map[string]bool)
dirInfo := make(map[string]*FileInfo)
for filePath, entry := range b.files {
normalizedFilePath := normalizePath(filePath)
// Check if file is under the given path
if path == "/" || strings.HasPrefix(normalizedFilePath, path+"/") || normalizedFilePath == path {
// For directory listing, we want to show immediate children
relativePath := strings.TrimPrefix(normalizedFilePath, path)
relativePath = strings.TrimPrefix(relativePath, "/")
if relativePath == "" {
// The path itself is a file
if !seen[normalizedFilePath] {
result = append(result, FileInfo{
Path: filepath.Base(normalizedFilePath),
IsDir: false,
Size: int64(len(entry.content)),
ModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano),
})
seen[normalizedFilePath] = true
}
continue
}
// Get the first segment (immediate child)
parts := strings.SplitN(relativePath, "/", 2)
if len(parts) > 0 {
childPath := path
if path != "/" {
childPath += "/"
}
childPath += parts[0]
isDir := len(parts) > 1
if !seen[childPath] {
if isDir {
dirInfo[childPath] = &FileInfo{
Path: parts[0],
IsDir: true,
Size: 0,
ModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano),
}
} else {
result = append(result, FileInfo{
Path: parts[0],
IsDir: false,
Size: int64(len(entry.content)),
ModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano),
})
}
seen[childPath] = true
} else if isDir {
if info, ok := dirInfo[childPath]; ok {
if entry.modifiedAt.After(mustParseTime(info.ModifiedAt)) {
info.ModifiedAt = entry.modifiedAt.Format(time.RFC3339Nano)
}
}
}
}
}
}
for _, info := range dirInfo {
result = append(result, *info)
}
return result, nil
}
func mustParseTime(s string) time.Time {
t, _ := time.Parse(time.RFC3339Nano, s)
return t
}
// Read reads file content with offset and limit.
func (b *InMemoryBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) {
b.mu.RLock()
defer b.mu.RUnlock()
filePath := normalizePath(req.FilePath)
entry, exists := b.files[filePath]
if !exists {
return nil, fmt.Errorf("file not found: %s", filePath)
}
// Convert 1-based offset to 0-based index; values < 1 default to line 1
offset := req.Offset - 1
if offset < 0 {
offset = 0
}
limit := req.Limit
content := entry.content
// Fast path: no offset and no limit — return as-is
if offset == 0 && limit <= 0 {
return &FileContent{Content: content}, nil
}
// Fast path: no offset, content fits within limit — return as-is
if offset == 0 {
lineCount := strings.Count(content, "\n") + 1
if lineCount <= limit {
return &FileContent{Content: content}, nil
}
}
// Skip `offset` lines by scanning for newlines directly
start := 0
for i := 0; i < offset; i++ {
idx := strings.IndexByte(content[start:], '\n')
if idx == -1 {
// offset exceeds total lines
return &FileContent{}, nil
}
start += idx + 1
}
// No limit: return everything from start
if limit <= 0 {
return &FileContent{Content: content[start:]}, nil
}
// Find the end position after `limit` lines
end := start
for i := 0; i < limit; i++ {
idx := strings.IndexByte(content[end:], '\n')
if idx == -1 {
// Reached the end of content
return &FileContent{Content: content[start:]}, nil
}
end += idx + 1
}
// Trim the trailing newline from the last included line
return &FileContent{Content: content[start : end-1]}, nil
}
// GrepRaw returns matches for the given pattern.
func (b *InMemoryBackend) GrepRaw(ctx context.Context, req *GrepRequest) ([]GrepMatch, error) {
b.mu.RLock()
defer b.mu.RUnlock()
if req.Pattern == "" {
return nil, fmt.Errorf("pattern cannot be empty")
}
re, err := b.compilePattern(req)
if err != nil {
return nil, err
}
searchPath := "/"
if req.Path != "" {
searchPath = normalizePath(req.Path)
}
filteredFiles, err := b.filterFiles(searchPath, req)
if err != nil {
return nil, err
}
if len(filteredFiles) == 0 {
return []GrepMatch{}, nil
}
if len(filteredFiles) == 1 {
collector := newGrepCollector()
entry := b.files[filteredFiles[0]]
collector.processFile(filteredFiles[0], entry.content, re, req)
return collector.buildResults(b, req)
}
matches, err := b.grepFilesInParallel(filteredFiles, re, req)
if err != nil {
return nil, err
}
if req.BeforeLines > 0 || req.AfterLines > 0 {
matches = b.applyContext(matches, req)
}
return matches, nil
}
func (b *InMemoryBackend) grepFilesInParallel(filteredFiles []string, re *regexp.Regexp, req *GrepRequest) ([]GrepMatch, error) {
numWorkers := len(filteredFiles)
if numWorkers > 10 {
numWorkers = 10
}
type fileTask struct {
path string
content string
}
tasks := make(chan fileTask, len(filteredFiles))
results := make(chan []GrepMatch, len(filteredFiles))
errChan := make(chan error, numWorkers)
var wg sync.WaitGroup
for i := 0; i < numWorkers; i++ {
wg.Add(1)
go func(workerID int) {
defer wg.Done()
defer func() {
if r := recover(); r != nil {
errChan <- fmt.Errorf("worker %d panic: %v", workerID, r)
}
}()
collector := newGrepCollector()
for task := range tasks {
fileMatches := collector.findMatches(task.path, task.content, re, req)
if len(fileMatches) > 0 {
results <- fileMatches
}
}
}(i)
}
for _, filePath := range filteredFiles {
entry := b.files[filePath]
tasks <- fileTask{
path: filePath,
content: entry.content,
}
}
close(tasks)
go func() {
wg.Wait()
close(results)
close(errChan)
}()
var allMatches []GrepMatch
var errs []error
for {
select {
case matches, ok := <-results:
if !ok {
results = nil
} else {
allMatches = append(allMatches, matches...)
}
case err, ok := <-errChan:
if !ok {
errChan = nil
} else if err != nil {
errs = append(errs, err)
}
}
if results == nil && errChan == nil {
break
}
}
if len(errs) > 0 {
return nil, fmt.Errorf("grep failed with %d error(s): %v", len(errs), errs[0])
}
return allMatches, nil
}
func (b *InMemoryBackend) compilePattern(req *GrepRequest) (*regexp.Regexp, error) {
pattern := req.Pattern
if req.CaseInsensitive {
pattern = "(?i)" + pattern
}
re, err := regexp.Compile(pattern)
if err != nil {
return nil, fmt.Errorf("invalid regex pattern: %w", err)
}
return re, nil
}
func (b *InMemoryBackend) filterFiles(searchPath string, req *GrepRequest) ([]string, error) {
var candidateFiles []string
for filePath := range b.files {
normalizedFilePath := normalizePath(filePath)
if searchPath != "/" && !strings.HasPrefix(normalizedFilePath, searchPath+"/") && normalizedFilePath != searchPath {
continue
}
candidateFiles = append(candidateFiles, normalizedFilePath)
}
if req.Glob != "" {
filtered, err := b.filterByGlob(candidateFiles, searchPath, req.Glob)
if err != nil {
return nil, err
}
candidateFiles = filtered
}
if req.FileType != "" {
candidateFiles = b.filterByFileType(candidateFiles, req.FileType)
}
return candidateFiles, nil
}
func (b *InMemoryBackend) filterByGlob(files []string, searchPath string, globPattern string) ([]string, error) {
var result []string
for _, filePath := range files {
var matchPath string
if strings.Contains(globPattern, "/") || strings.Contains(globPattern, "**") {
if searchPath == "/" {
matchPath = strings.TrimPrefix(filePath, "/")
} else {
matchPath = strings.TrimPrefix(filePath, searchPath+"/")
}
} else {
matchPath = filepath.Base(filePath)
}
matched, err := doublestar.Match(globPattern, matchPath)
if err != nil {
return nil, fmt.Errorf("invalid glob pattern: %w", err)
}
if matched {
result = append(result, filePath)
}
}
return result, nil
}
func (b *InMemoryBackend) filterByFileType(files []string, fileType string) []string {
var result []string
for _, filePath := range files {
ext := strings.TrimPrefix(filepath.Ext(filePath), ".")
if matchFileType(ext, fileType) {
result = append(result, filePath)
}
}
return result
}
// matchFileType checks if the file extension matches the given file type.
func matchFileType(ext, fileType string) bool {
typeMap := map[string][]string{
"ada": {"adb", "ads"},
"agda": {"agda", "lagda"},
"aidl": {"aidl"},
"amake": {"bp", "mk"},
"asciidoc": {"adoc", "asc", "asciidoc"},
"asm": {"S", "asm", "s"},
"asp": {"ascx", "asp", "aspx"},
"ats": {"ats", "dats", "hats", "sats"},
"avro": {"avdl", "avpr", "avsc"},
"awk": {"awk"},
"bat": {"bat"},
"bazel": {"BUILD", "bazel", "bzl"},
"bitbake": {"bb", "bbappend", "bbclass", "conf", "inc"},
"c": {"c", "h", "H", "cats"},
"cabal": {"cabal"},
"cbor": {"cbor"},
"ceylon": {"ceylon"},
"clojure": {"clj", "cljc", "cljs", "cljx"},
"cmake": {"cmake"},
"coffeescript": {"coffee"},
"config": {"cfg", "conf", "config", "ini"},
"coq": {"v"},
"cpp": {"C", "cc", "cpp", "cxx", "c++", "h", "hh", "hpp", "hxx", "h++", "inl"},
"crystal": {"cr", "ecr"},
"cs": {"cs"},
"csharp": {"cs"},
"cshtml": {"cshtml"},
"css": {"css", "scss", "sass", "less"},
"csv": {"csv"},
"cuda": {"cu", "cuh"},
"cython": {"pxd", "pxi", "pyx"},
"d": {"d"},
"dart": {"dart"},
"devicetree": {"dts", "dtsi"},
"dhall": {"dhall"},
"diff": {"diff", "patch"},
"docker": {"dockerfile"},
"go": {"go"},
"groovy": {"gradle", "groovy"},
"haskell": {"c2hs", "cpphs", "hs", "hsc", "lhs"},
"html": {"ejs", "htm", "html"},
"java": {"java", "jsp", "jspx", "properties"},
"js": {"cjs", "js", "jsx", "mjs", "vue"},
"json": {"json", "sarif"},
"jsonl": {"jsonl"},
"julia": {"jl"},
"jupyter": {"ipynb", "jpynb"},
"kotlin": {"kt", "kts"},
"less": {"less"},
"lua": {"lua"},
"make": {"mak", "mk"},
"markdown": {"markdown", "md", "mdown", "mdwn", "mdx", "mkd", "mkdn"},
"md": {"markdown", "md", "mdown", "mdwn", "mdx", "mkd", "mkdn"},
"matlab": {"m"},
"ocaml": {"ml", "mli", "mll", "mly"},
"perl": {"PL", "perl", "pl", "plh", "plx", "pm", "t"},
"php": {"php", "php3", "php4", "php5", "php7", "php8", "pht", "phtml"},
"python": {"py", "pyi"},
"py": {"py", "pyi"},
"ruby": {"gemspec", "rb", "rbw"},
"rust": {"rs"},
"sass": {"sass", "scss"},
"scala": {"sbt", "scala"},
"sh": {"bash", "sh", "zsh"},
"sql": {"psql", "sql"},
"swift": {"swift"},
"toml": {"toml"},
"ts": {"cts", "mts", "ts", "tsx"},
"typescript": {"cts", "mts", "ts", "tsx"},
"txt": {"txt"},
"vue": {"vue"},
"xml": {"dtd", "xml", "xsd", "xsl", "xslt"},
"yaml": {"yaml", "yml"},
"zig": {"zig"},
}
if exts, ok := typeMap[fileType]; ok {
for _, e := range exts {
if ext == e {
return true
}
}
}
return ext == fileType
}
// applyContext adds context lines around matches.
func (b *InMemoryBackend) applyContext(matches []GrepMatch, req *GrepRequest) []GrepMatch {
if len(matches) == 0 {
return matches
}
beforeLines := 0
afterLines := 0
if req.BeforeLines > 0 {
beforeLines = req.BeforeLines
}
if req.AfterLines > 0 {
afterLines = req.AfterLines
}
if beforeLines <= 0 && afterLines <= 0 {
return matches
}
// Group matches by file path for efficient processing
matchesByFile := make(map[string][]GrepMatch)
fileOrder := make([]string, 0)
seenFiles := make(map[string]bool)
for _, match := range matches {
if !seenFiles[match.Path] {
fileOrder = append(fileOrder, match.Path)
seenFiles[match.Path] = true
}
matchesByFile[match.Path] = append(matchesByFile[match.Path], match)
}
var result []GrepMatch
// Process each file once
for _, filePath := range fileOrder {
fileMatches := matchesByFile[filePath]
// Get file content once per file
b.mu.RLock()
entry, exists := b.files[filePath]
b.mu.RUnlock()
if !exists {
// If file doesn't exist, keep original matches
result = append(result, fileMatches...)
continue
}
lines := strings.Split(entry.content, "\n")
processedLines := make(map[int]bool)
// Process all matches for this file
for _, match := range fileMatches {
startLine := match.Line - beforeLines
if startLine < 1 {
startLine = 1
}
endLine := match.Line + afterLines
if endLine > len(lines) {
endLine = len(lines)
}
for lineNum := startLine; lineNum <= endLine; lineNum++ {
if !processedLines[lineNum] {
processedLines[lineNum] = true
result = append(result, GrepMatch{
Path: filePath,
Line: lineNum,
Content: lines[lineNum-1],
})
}
}
}
}
return result
}
// GlobInfo returns file info entries matching the glob pattern.
func (b *InMemoryBackend) GlobInfo(ctx context.Context, req *GlobInfoRequest) ([]FileInfo, error) {
b.mu.RLock()
defer b.mu.RUnlock()
basePath := normalizePath(req.Path)
isAbsolutePattern := strings.HasPrefix(req.Pattern, "/")
var result []FileInfo
for filePath, entry := range b.files {
normalizedFilePath := normalizePath(filePath)
var matchPath string
var resultPath string
if isAbsolutePattern {
matchPath = normalizedFilePath
resultPath = normalizedFilePath
} else {
if basePath != "/" && !strings.HasPrefix(normalizedFilePath, basePath+"/") && normalizedFilePath != basePath {
continue
}
if basePath == "/" {
matchPath = strings.TrimPrefix(normalizedFilePath, "/")
} else {
matchPath = strings.TrimPrefix(normalizedFilePath, basePath+"/")
}
resultPath = matchPath
}
matched, err := doublestar.Match(req.Pattern, matchPath)
if err != nil {
return nil, fmt.Errorf("invalid glob pattern: %w", err)
}
if matched {
result = append(result, FileInfo{
Path: resultPath,
IsDir: false,
Size: int64(len(entry.content)),
ModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano),
})
}
}
return result, nil
}
// Write creates or overwrites file content.
func (b *InMemoryBackend) Write(ctx context.Context, req *WriteRequest) error {
b.mu.Lock()
defer b.mu.Unlock()
filePath := normalizePath(req.FilePath)
b.files[filePath] = &fileEntry{
content: req.Content,
modifiedAt: time.Now(),
}
return nil
}
// Edit replaces string occurrences in a file.
func (b *InMemoryBackend) Edit(ctx context.Context, req *EditRequest) error {
b.mu.Lock()
defer b.mu.Unlock()
filePath := normalizePath(req.FilePath)
entry, exists := b.files[filePath]
if !exists {
return fmt.Errorf("file not found: %s", filePath)
}
if req.OldString == "" {
return fmt.Errorf("oldString must be non-empty")
}
content := entry.content
if !strings.Contains(content, req.OldString) {
return fmt.Errorf("oldString not found in file: %s", filePath)
}
if !req.ReplaceAll {
firstIndex := strings.Index(content, req.OldString)
if firstIndex != -1 {
// Check if there's another occurrence after the first one
if strings.Contains(content[firstIndex+len(req.OldString):], req.OldString) {
return fmt.Errorf("multiple occurrences of oldString found in file %s, but ReplaceAll is false", filePath)
}
}
}
var newContent string
if req.ReplaceAll {
newContent = strings.ReplaceAll(content, req.OldString, req.NewString)
} else {
newContent = strings.Replace(content, req.OldString, req.NewString, 1)
}
b.files[filePath] = &fileEntry{
content: newContent,
modifiedAt: time.Now(),
}
return nil
}
// normalizePath normalizes a file path by ensuring it starts with "/" and removing trailing slashes.
func normalizePath(path string) string {
if path == "" {
return "/"
}
// Ensure path starts with "/"
if !strings.HasPrefix(path, "/") {
path = "/" + path
}
return filepath.Clean(path)
}
type grepCollector struct {
allMatches []GrepMatch
}
func newGrepCollector() *grepCollector {
return &grepCollector{
allMatches: []GrepMatch{},
}
}
func (c *grepCollector) processFile(filePath, content string, re *regexp.Regexp, req *GrepRequest) {
fileMatches := c.findMatches(filePath, content, re, req)
if len(fileMatches) > 0 {
c.allMatches = append(c.allMatches, fileMatches...)
}
}
func (c *grepCollector) findMatches(filePath, content string, re *regexp.Regexp, req *GrepRequest) []GrepMatch {
if req.EnableMultiline {
return c.findMultilineMatches(filePath, content, re)
}
return c.findSingleLineMatches(filePath, content, re)
}
func (c *grepCollector) findMultilineMatches(filePath, content string, re *regexp.Regexp) []GrepMatch {
var fileMatches []GrepMatch
matches := re.FindAllStringIndex(content, -1)
lines := strings.Split(content, "\n")
for _, match := range matches {
matchStart := match[0]
matchEnd := match[1]
startLineNum := 1 + strings.Count(content[:matchStart], "\n")
endLineNum := 1 + strings.Count(content[:matchEnd], "\n")
for lineNum := startLineNum; lineNum <= endLineNum && lineNum <= len(lines); lineNum++ {
fileMatches = append(fileMatches, GrepMatch{
Path: filePath,
Line: lineNum,
Content: lines[lineNum-1],
})
}
}
return fileMatches
}
func (c *grepCollector) findSingleLineMatches(filePath, content string, re *regexp.Regexp) []GrepMatch {
var fileMatches []GrepMatch
lines := strings.Split(content, "\n")
for lineNum, line := range lines {
if re.MatchString(line) {
fileMatches = append(fileMatches, GrepMatch{
Path: filePath,
Line: lineNum + 1,
Content: line,
})
}
}
return fileMatches
}
func (c *grepCollector) buildResults(b *InMemoryBackend, req *GrepRequest) ([]GrepMatch, error) {
return c.buildContentResult(b, req), nil
}
func (c *grepCollector) buildContentResult(b *InMemoryBackend, req *GrepRequest) []GrepMatch {
results := c.allMatches
if req.BeforeLines > 0 || req.AfterLines > 0 {
results = b.applyContext(c.allMatches, req)
}
return results
}
File diff suppressed because it is too large Load Diff
+809
View File
@@ -0,0 +1,809 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"errors"
"fmt"
"log"
"runtime/debug"
"strings"
"github.com/cloudwego/eino/callbacks"
"github.com/cloudwego/eino/compose"
icb "github.com/cloudwego/eino/internal/callbacks"
"github.com/cloudwego/eino/internal/safe"
"github.com/cloudwego/eino/schema"
)
type HistoryEntry struct {
IsUserInput bool
AgentName string
Message Message
}
type HistoryRewriter func(ctx context.Context, entries []*HistoryEntry) ([]Message, error)
type flowAgent struct {
Agent
subAgents []*flowAgent
parentAgent *flowAgent
disallowTransferToParent bool
historyRewriter HistoryRewriter
checkPointStore compose.CheckPointStore
}
func (a *flowAgent) deepCopy() *flowAgent {
ret := &flowAgent{
Agent: a.Agent,
subAgents: make([]*flowAgent, 0, len(a.subAgents)),
parentAgent: a.parentAgent,
disallowTransferToParent: a.disallowTransferToParent,
historyRewriter: a.historyRewriter,
checkPointStore: a.checkPointStore,
}
for _, sa := range a.subAgents {
ret.subAgents = append(ret.subAgents, sa.deepCopy())
}
return ret
}
// SetSubAgents sets sub-agents for the given agent and returns the updated agent.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func SetSubAgents(ctx context.Context, agent Agent, subAgents []Agent) (ResumableAgent, error) {
return setSubAgents(ctx, agent, subAgents)
}
type AgentOption func(options *flowAgent)
// WithDisallowTransferToParent prevents a sub-agent from transferring to its parent.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func WithDisallowTransferToParent() AgentOption {
return func(fa *flowAgent) {
fa.disallowTransferToParent = true
}
}
// WithHistoryRewriter sets a rewriter to transform conversation history
// during agent transfers.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func WithHistoryRewriter(h HistoryRewriter) AgentOption {
return func(fa *flowAgent) {
fa.historyRewriter = h
}
}
func toFlowAgent(ctx context.Context, agent Agent, opts ...AgentOption) *flowAgent {
var fa *flowAgent
var ok bool
if fa, ok = agent.(*flowAgent); !ok {
fa = &flowAgent{Agent: agent}
} else {
fa = fa.deepCopy()
}
for _, opt := range opts {
opt(fa)
}
if fa.historyRewriter == nil {
fa.historyRewriter = buildDefaultHistoryRewriter(agent.Name(ctx))
}
return fa
}
// AgentWithOptions wraps an agent with flow-specific options and returns it.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func AgentWithOptions(ctx context.Context, agent Agent, opts ...AgentOption) Agent {
return toFlowAgent(ctx, agent, opts...)
}
func setSubAgents(ctx context.Context, agent Agent, subAgents []Agent) (*flowAgent, error) {
fa := toFlowAgent(ctx, agent)
if len(fa.subAgents) > 0 {
return nil, errors.New("agent's sub-agents has already been set")
}
if onAgent, ok_ := fa.Agent.(OnSubAgents); ok_ {
err := onAgent.OnSetSubAgents(ctx, subAgents)
if err != nil {
return nil, err
}
}
for _, s := range subAgents {
fsa := toFlowAgent(ctx, s)
if fsa.parentAgent != nil {
return nil, errors.New("agent has already been set as a sub-agent of another agent")
}
fsa.parentAgent = fa
if onAgent, ok__ := fsa.Agent.(OnSubAgents); ok__ {
err := onAgent.OnSetAsSubAgent(ctx, agent)
if err != nil {
return nil, err
}
if fsa.disallowTransferToParent {
err = onAgent.OnDisallowTransferToParent(ctx)
if err != nil {
return nil, err
}
}
}
fa.subAgents = append(fa.subAgents, fsa)
}
return fa, nil
}
func (a *flowAgent) getAgent(ctx context.Context, name string) *flowAgent {
for _, subAgent := range a.subAgents {
if subAgent.Name(ctx) == name {
return subAgent
}
}
if a.parentAgent != nil && a.parentAgent.Name(ctx) == name {
return a.parentAgent
}
return nil
}
func rewriteMessage(msg Message, agentName string) Message {
var sb strings.Builder
sb.WriteString("For context:")
if msg.Role == schema.Assistant {
if msg.Content != "" {
sb.WriteString(fmt.Sprintf(" [%s] said: %s.", agentName, msg.Content))
}
if len(msg.ToolCalls) > 0 {
for i := range msg.ToolCalls {
f := msg.ToolCalls[i].Function
sb.WriteString(fmt.Sprintf(" [%s] called tool: `%s` with arguments: %s.",
agentName, f.Name, f.Arguments))
}
}
} else if msg.Role == schema.Tool && msg.Content != "" {
sb.WriteString(fmt.Sprintf(" [%s] `%s` tool returned result: %s.",
agentName, msg.ToolName, msg.Content))
}
rewritten := schema.UserMessage(sb.String())
if msg.MultiContent != nil {
rewritten.MultiContent = append([]schema.ChatMessagePart{}, msg.MultiContent...) //nolint:staticcheck // backward compat with deprecated MultiContent field
}
if msg.UserInputMultiContent != nil {
rewritten.UserInputMultiContent = append([]schema.MessageInputPart{}, msg.UserInputMultiContent...)
}
// Convert AssistantGenMultiContent to UserInputMultiContent, since the role changes to User.
// Reasoning parts have no user input equivalent and are dropped.
for _, part := range msg.AssistantGenMultiContent {
switch part.Type {
case schema.ChatMessagePartTypeText:
rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{
Type: part.Type,
Text: part.Text,
Extra: part.Extra,
})
case schema.ChatMessagePartTypeImageURL:
if part.Image != nil {
rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{
Type: part.Type,
Image: &schema.MessageInputImage{MessagePartCommon: part.Image.MessagePartCommon},
Extra: part.Extra,
})
}
case schema.ChatMessagePartTypeAudioURL:
if part.Audio != nil {
rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{
Type: part.Type,
Audio: &schema.MessageInputAudio{MessagePartCommon: part.Audio.MessagePartCommon},
Extra: part.Extra,
})
}
case schema.ChatMessagePartTypeVideoURL:
if part.Video != nil {
rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{
Type: part.Type,
Video: &schema.MessageInputVideo{MessagePartCommon: part.Video.MessagePartCommon},
Extra: part.Extra,
})
}
}
}
return rewritten
}
func genMsg(entry *HistoryEntry, agentName string) (Message, error) {
msg := entry.Message
if entry.AgentName != agentName {
msg = rewriteMessage(msg, entry.AgentName)
}
return msg, nil
}
func deepCopyAgentInput(ai *AgentInput) *AgentInput {
copied := &AgentInput{
Messages: make([]Message, len(ai.Messages)),
EnableStreaming: ai.EnableStreaming,
}
copy(copied.Messages, ai.Messages)
return copied
}
func (a *flowAgent) genAgentInput(ctx context.Context, runCtx *runContext, skipTransferMessages bool) (*AgentInput, error) {
input := deepCopyAgentInput(runCtx.RootInput)
events := runCtx.Session.getEvents()
historyEntries := make([]*HistoryEntry, 0)
for _, m := range input.Messages {
historyEntries = append(historyEntries, &HistoryEntry{
IsUserInput: true,
Message: m,
})
}
for _, event := range events {
if skipTransferMessages && event.Action != nil && event.Action.TransferToAgent != nil {
// If skipTransferMessages is true and the event contain transfer action, the message in this event won't be appended to history entries.
if event.Output != nil &&
event.Output.MessageOutput != nil &&
event.Output.MessageOutput.Role == schema.Tool &&
len(historyEntries) > 0 {
// If the skipped message's role is Tool, remove the previous history entry as it's also a transfer message(from ChatModelAgent and GenTransferMessages).
historyEntries = historyEntries[:len(historyEntries)-1]
}
continue
}
msg, err := getMessageFromWrappedEvent(event)
if err != nil {
var retryErr *WillRetryError
if errors.As(err, &retryErr) {
log.Printf("failed to get message from event, but will retry: %v", err)
continue
}
return nil, err
}
if msg == nil {
continue
}
historyEntries = append(historyEntries, &HistoryEntry{
AgentName: event.AgentName,
Message: msg,
})
}
messages, err := a.historyRewriter(ctx, historyEntries)
if err != nil {
return nil, err
}
input.Messages = messages
return input, nil
}
func buildDefaultHistoryRewriter(agentName string) HistoryRewriter {
return func(ctx context.Context, entries []*HistoryEntry) ([]Message, error) {
messages := make([]Message, 0, len(entries))
var err error
for _, entry := range entries {
msg := entry.Message
if !entry.IsUserInput {
msg, err = genMsg(entry, agentName)
if err != nil {
return nil, fmt.Errorf("gen agent input failed: %w", err)
}
}
if msg != nil {
messages = append(messages, msg)
}
}
return messages, nil
}
}
func (a *flowAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
agentName := a.Name(ctx)
var runCtx *runContext
ctx, runCtx = initRunCtx(ctx, agentName, input)
ctx = AppendAddressSegment(ctx, AddressSegmentAgent, agentName)
o := getCommonOptions(nil, opts...)
cancelCtx := o.cancelCtx
processedInput, err := a.genAgentInput(ctx, runCtx, o.skipTransferMessages)
if err != nil {
if cancelCtx != nil {
cancelCtx.markDone()
}
cbInput := &AgentCallbackInput{Input: input}
ctx = callbacks.OnStart(ctx, cbInput)
return wrapIterWithOnEnd(ctx, genErrorIter(err))
}
ctxForSubAgents := ctx
agentType := getAgentType(a.Agent)
ctx = initAgentCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
cbInput := &AgentCallbackInput{Input: processedInput}
ctx = callbacks.OnStart(ctx, cbInput)
input = processedInput
if wf, ok := a.Agent.(*workflowAgent); ok {
ctx = withCancelContext(ctx, cancelCtx)
filteredOpts := filterCancelOption(filterCallbackHandlersForNestedAgents(agentName, opts))
iter := wf.Run(ctx, input, filteredOpts...)
iter = wrapIterWithCancelCtx(iter, cancelCtx)
return wrapIterWithOnEnd(ctx, iter)
}
aIter := a.Agent.Run(withCancelContext(ctx, cancelCtx), input, filterOptions(agentName, opts)...)
iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), runCtx, aIter, generator, filterCancelOption(opts)...)
return wrapIterWithCancelCtx(iterator, cancelCtx)
}
func (a *flowAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
agentName := a.Name(ctx)
ctx, info = buildResumeInfo(ctx, agentName, info)
ctxForSubAgents := ctx
o := getCommonOptions(nil, opts...)
cancelCtx := o.cancelCtx
agentType := getAgentType(a.Agent)
ctx = initAgentCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
cbInput := &AgentCallbackInput{ResumeInfo: info}
ctx = callbacks.OnStart(ctx, cbInput)
if info.WasInterrupted {
if ra, ok := a.Agent.(ResumableAgent); ok {
if _, ok := ra.(*workflowAgent); ok {
ctx = withCancelContext(ctx, cancelCtx)
filteredOpts := filterCancelOption(filterCallbackHandlersForNestedAgents(agentName, opts))
aIter := ra.Resume(ctx, info, filteredOpts...)
aIter = wrapIterWithCancelCtx(aIter, cancelCtx)
return wrapIterWithOnEnd(ctx, aIter)
}
aIter := ra.Resume(withCancelContext(ctx, cancelCtx), info, opts...)
iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), getRunCtx(ctxForSubAgents), aIter, generator, filterCancelOption(opts)...)
return wrapIterWithCancelCtx(iterator, cancelCtx)
}
if cancelCtx != nil {
cancelCtx.markDone()
}
return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf("failed to resume agent: agent '%s' is an interrupt point "+
"but is not a ResumableAgent", agentName)))
}
nextAgentName, err := getNextResumeAgent(ctx, info)
if err != nil {
if cancelCtx != nil {
cancelCtx.markDone()
}
return wrapIterWithOnEnd(ctx, genErrorIter(err))
}
subAgent := a.getAgent(ctxForSubAgents, nextAgentName)
if subAgent == nil {
if len(a.subAgents) == 0 {
if ra, ok := a.Agent.(ResumableAgent); ok {
ctx = withCancelContext(ctx, cancelCtx)
innerIter := ra.Resume(ctx, info, filterCancelOption(opts)...)
return wrapIterWithCancelCtx(wrapIterWithOnEnd(ctx, innerIter), cancelCtx)
}
return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf(
"failed to resume agent: agent '%s' (type %T) has no sub-agents and does not implement ResumableAgent interface. "+
"To support resume, your custom agent wrapper must implement the ResumableAgent interface", agentName, a.Agent)))
}
if cancelCtx != nil {
cancelCtx.markDone()
}
return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf("failed to resume agent: sub-agent '%s' not found in agent '%s'", nextAgentName, agentName)))
}
ctxForSubAgents = withCancelContext(ctxForSubAgents, cancelCtx)
innerIter := subAgent.Resume(ctxForSubAgents, info, filterCancelOption(opts)...)
return wrapIterWithCancelCtx(wrapIterWithOnEnd(ctx, innerIter), cancelCtx)
}
// DeterministicTransferConfig is the configuration for AgentWithDeterministicTransferTo.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
type DeterministicTransferConfig struct {
Agent Agent
ToAgentNames []string
}
func (a *flowAgent) run(
ctx context.Context,
ctxForSubAgents context.Context,
runCtx *runContext,
aIter *AsyncIterator[*AgentEvent],
generator *AsyncGenerator[*AgentEvent],
opts ...AgentRunOption) {
cbIter, cbGen := NewAsyncIteratorPair[*AgentEvent]()
cbOutput := &AgentCallbackOutput{Events: cbIter}
icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyAgentCallbackOutput), callbacks.TimingOnEnd, false)
defer func() {
panicErr := recover()
if panicErr != nil {
e := safe.NewPanicErr(panicErr, debug.Stack())
generator.Send(&AgentEvent{Err: e})
}
cbGen.Close()
generator.Close()
}()
var lastAction *AgentAction
for {
event, ok := aIter.Next()
if !ok {
break
}
// RunPath ownership: the eino framework sets RunPath exactly once.
// If event.RunPath is already set (e.g., by agentTool), we don't modify it.
// If event.RunPath is nil/empty, we set it to the current runCtx.RunPath.
// This ensures RunPath is set exactly once and not duplicated.
if len(event.RunPath) == 0 {
event.AgentName = a.Name(ctx)
event.RunPath = runCtx.RunPath
}
// Recording policy: exact RunPath match (non-interrupt) indicates events belonging to this agent execution.
// This prevents parent recording of child/tool-internal emissions.
if (event.Action == nil || event.Action.Interrupted == nil) && exactRunPathMatch(runCtx.RunPath, event.RunPath) {
// copy the event so that the copied event's stream is exclusive for any potential consumer
// copy before adding to session because once added to session it's stream could be consumed by genAgentInput at any time
// interrupt action are not added to session, because ALL information contained in it
// is either presented to end-user, or made available to agents through other means
copied := copyTypedAgentEvent(event)
setAutomaticClose(copied)
setAutomaticClose(event)
runCtx.Session.addEvent(copied)
}
// Action gating uses exact run-path match as well:
// only actions originating from this agent execution (not child/tool runs)
// should influence parent control flow (exit/transfer/interrupt).
if exactRunPathMatch(runCtx.RunPath, event.RunPath) {
lastAction = event.Action
}
copied := copyTypedAgentEvent(event)
setAutomaticClose(copied)
setAutomaticClose(event)
cbGen.Send(copied)
generator.Send(event)
}
var destName string
if lastAction != nil {
if lastAction.Interrupted != nil {
return
}
if lastAction.Exit {
return
}
if lastAction.TransferToAgent != nil {
destName = lastAction.TransferToAgent.DestAgentName
}
}
// handle transferring to another agent
if destName != "" {
agentToRun := a.getAgent(ctxForSubAgents, destName)
if agentToRun == nil {
e := fmt.Errorf("transfer failed: agent '%s' not found when transferring from '%s'",
destName, a.Name(ctxForSubAgents))
generator.Send(&AgentEvent{Err: e})
return
}
subAIter := agentToRun.Run(ctxForSubAgents, nil /*subagents get input from runCtx*/, opts...)
for {
subEvent, ok_ := subAIter.Next()
if !ok_ {
break
}
setAutomaticClose(subEvent)
generator.Send(subEvent)
}
}
}
func exactRunPathMatch(aPath, bPath []RunStep) bool {
if len(aPath) != len(bPath) {
return false
}
for i := range aPath {
if !aPath[i].Equals(bPath[i]) {
return false
}
}
return true
}
func wrapIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*AgentEvent]) *AsyncIterator[*AgentEvent] {
cbIter, cbGen := NewAsyncIteratorPair[*AgentEvent]()
cbOutput := &AgentCallbackOutput{Events: cbIter}
icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyAgentCallbackOutput), callbacks.TimingOnEnd, false)
outIter, outGen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer func() {
cbGen.Close()
outGen.Close()
}()
for {
event, ok := iter.Next()
if !ok {
break
}
copied := copyTypedAgentEvent(event)
cbGen.Send(copied)
outGen.Send(event)
}
}()
return outIter
}
// ---------------------------------------------------------------------------
// Typed wrapper for the agentic path (TypedAgent[*schema.AgenticMessage]).
//
// typedFlowAgent is a minimal wrapper used exclusively by TypedRunner and
// AgentTool to execute a TypedAgent[*schema.AgenticMessage]. It handles
// callbacks, event recording, and run-path tracking. Transfer, sub-agent
// orchestration, and history rewriting are handled solely by the concrete
// flowAgent (the *schema.Message path).
// ---------------------------------------------------------------------------
type typedFlowAgent[M MessageType] struct {
TypedAgent[M]
checkPointStore compose.CheckPointStore
}
func toTypedFlowAgent[M MessageType](agent TypedAgent[M]) *typedFlowAgent[M] {
if fa, ok := agent.(*typedFlowAgent[M]); ok {
return fa
}
return &typedFlowAgent[M]{TypedAgent: agent}
}
func getTypedAgentType[M MessageType](agent TypedAgent[M]) string {
if msgAgent, ok := any(agent).(Agent); ok {
return getAgentType(msgAgent)
}
if typer, ok := any(agent).(interface{ GetType() string }); ok {
return typer.GetType()
}
return ""
}
func (a *typedFlowAgent[M]) Run(ctx context.Context, input *TypedAgentInput[M], opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
agentName := a.Name(ctx)
var runCtx *runContext
ctx, runCtx = initTypedRunCtx(ctx, agentName, input)
ctx = AppendAddressSegment(ctx, AddressSegmentAgent, agentName)
o := getCommonOptions(nil, opts...)
cancelCtx := o.cancelCtx
ctxForSubAgents := ctx
agentType := getTypedAgentType(a.TypedAgent)
ctx = initAgenticCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{Input: any(input).(*TypedAgentInput[*schema.AgenticMessage])}
ctx = callbacks.OnStart(ctx, cbInput)
aIter := a.TypedAgent.Run(withCancelContext(ctx, cancelCtx), input, filterOptions(agentName, opts)...)
iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), runCtx, aIter, generator, filterCancelOption(opts)...)
return wrapIterWithCancelCtx(iterator, cancelCtx)
}
func (a *typedFlowAgent[M]) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
agentName := a.Name(ctx)
ctx, info = buildResumeInfo(ctx, agentName, info)
ctxForSubAgents := ctx
o := getCommonOptions(nil, opts...)
cancelCtx := o.cancelCtx
agentType := getTypedAgentType(a.TypedAgent)
ctx = initAgenticCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{ResumeInfo: info}
ctx = callbacks.OnStart(ctx, cbInput)
if info.WasInterrupted {
if ra, ok := a.TypedAgent.(TypedResumableAgent[M]); ok {
aIter := ra.Resume(withCancelContext(ctx, cancelCtx), info, opts...)
iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), getRunCtx(ctxForSubAgents), aIter, generator, filterCancelOption(opts)...)
return wrapIterWithCancelCtx(iterator, cancelCtx)
}
if cancelCtx != nil {
cancelCtx.markDone()
}
return typedErrorIterWithOnEnd[M](ctx, fmt.Errorf("failed to resume agent: agent '%s' is an interrupt point "+
"but is not a ResumableAgent", agentName))
}
_, err := getNextResumeAgent(ctx, info)
if err != nil {
if cancelCtx != nil {
cancelCtx.markDone()
}
return typedErrorIterWithOnEnd[M](ctx, err)
}
if ra, ok := a.TypedAgent.(TypedResumableAgent[M]); ok {
ctx = withCancelContext(ctx, cancelCtx)
innerIter := ra.Resume(ctx, info, filterCancelOption(opts)...)
return wrapIterWithCancelCtx(typedWrapIterWithOnEnd(ctx, innerIter), cancelCtx)
}
return typedErrorIterWithOnEnd[M](ctx, fmt.Errorf(
"failed to resume agent: agent '%s' (type %T) does not implement ResumableAgent interface. "+
"To support resume, your custom agent wrapper must implement the ResumableAgent interface", agentName, a.TypedAgent))
}
func (a *typedFlowAgent[M]) run(
ctx context.Context,
_ context.Context,
runCtx *runContext,
aIter *AsyncIterator[*TypedAgentEvent[M]],
generator *AsyncGenerator[*TypedAgentEvent[M]],
_ ...AgentRunOption) {
agenticCbIter, agenticCbGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
cbOutput := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: agenticCbIter}
icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyTypedCallbackOutput[*schema.AgenticMessage]), callbacks.TimingOnEnd, false)
defer func() {
panicErr := recover()
if panicErr != nil {
e := safe.NewPanicErr(panicErr, debug.Stack())
generator.Send(&TypedAgentEvent[M]{Err: e})
}
agenticCbGen.Close()
generator.Close()
}()
for {
event, ok := aIter.Next()
if !ok {
break
}
if len(event.RunPath) == 0 {
event.AgentName = a.Name(ctx)
event.RunPath = runCtx.RunPath
}
if (event.Action == nil || event.Action.Interrupted == nil) && exactRunPathMatch(runCtx.RunPath, event.RunPath) {
copied := copyTypedAgentEvent(event)
typedSetAutomaticClose(copied)
typedSetAutomaticClose(event)
addTypedEvent(runCtx.Session, copied)
}
agenticCopied := copyTypedAgentEvent(event)
typedSetAutomaticClose(agenticCopied)
typedSetAutomaticClose(event)
agenticCbGen.Send(any(agenticCopied).(*TypedAgentEvent[*schema.AgenticMessage]))
generator.Send(event)
}
}
func wrapAgenticIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]]) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
cbIter, cbGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
cbOutput := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: cbIter}
icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyTypedCallbackOutput[*schema.AgenticMessage]), callbacks.TimingOnEnd, false)
outIter, outGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
go func() {
defer func() {
cbGen.Close()
outGen.Close()
}()
for {
event, ok := iter.Next()
if !ok {
break
}
copied := copyTypedAgentEvent(event)
cbGen.Send(copied)
outGen.Send(event)
}
}()
return outIter
}
func genAgenticErrorIter(err error) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err})
gen.Close()
return iter
}
func typedWrapIterWithOnEnd[M MessageType](ctx context.Context, iter *AsyncIterator[*TypedAgentEvent[M]]) *AsyncIterator[*TypedAgentEvent[M]] {
agenticIter := any(iter).(*AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]])
return any(wrapAgenticIterWithOnEnd(ctx, agenticIter)).(*AsyncIterator[*TypedAgentEvent[M]])
}
func typedErrorIterWithOnEnd[M MessageType](ctx context.Context, err error) *AsyncIterator[*TypedAgentEvent[M]] {
return typedWrapIterWithOnEnd(ctx, typedErrorIter[M](err))
}
+389
View File
@@ -0,0 +1,389 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"go.uber.org/mock/gomock"
"github.com/cloudwego/eino/callbacks"
mockModel "github.com/cloudwego/eino/internal/mock/components/model"
"github.com/cloudwego/eino/schema"
)
func strPtr(s string) *string { return &s }
func TestRewriteMessage(t *testing.T) {
imageCommon := schema.MessagePartCommon{URL: strPtr("http://img.example.com")}
audioCommon := schema.MessagePartCommon{URL: strPtr("http://audio.example.com")}
videoCommon := schema.MessagePartCommon{URL: strPtr("http://video.example.com")}
msg := &schema.Message{
Role: schema.Assistant,
Content: "hello",
MultiContent: []schema.ChatMessagePart{ //nolint:staticcheck // testing deprecated field backward compat
{Type: schema.ChatMessagePartTypeText, Text: "legacy"},
},
UserInputMultiContent: []schema.MessageInputPart{
{Type: schema.ChatMessagePartTypeText, Text: "pre-existing"},
},
AssistantGenMultiContent: []schema.MessageOutputPart{
{Type: schema.ChatMessagePartTypeText, Text: "gen-text", Extra: map[string]any{"k": "v"}},
{Type: schema.ChatMessagePartTypeImageURL, Image: &schema.MessageOutputImage{MessagePartCommon: imageCommon}},
{Type: schema.ChatMessagePartTypeAudioURL, Audio: &schema.MessageOutputAudio{MessagePartCommon: audioCommon}},
{Type: schema.ChatMessagePartTypeVideoURL, Video: &schema.MessageOutputVideo{MessagePartCommon: videoCommon}},
{Type: schema.ChatMessagePartTypeReasoning, Reasoning: &schema.MessageOutputReasoning{Text: "secret thoughts"}},
},
}
rewritten := rewriteMessage(msg, "OtherAgent")
assert.Equal(t, schema.User, rewritten.Role)
// MultiContent: copied, not shared
assert.Equal(t, msg.MultiContent, rewritten.MultiContent)
rewritten.MultiContent[0].Text = "mutated"
assert.Equal(t, "legacy", msg.MultiContent[0].Text)
// UserInputMultiContent: pre-existing entry copied, AssistantGenMultiContent appended (reasoning dropped)
assert.Len(t, rewritten.UserInputMultiContent, 5) // 1 pre-existing + 4 converted (text/image/audio/video)
// pre-existing entry is not shared
rewritten.UserInputMultiContent[0].Text = "mutated"
assert.Equal(t, "pre-existing", msg.UserInputMultiContent[0].Text)
// text conversion
assert.Equal(t, schema.ChatMessagePartTypeText, rewritten.UserInputMultiContent[1].Type)
assert.Equal(t, "gen-text", rewritten.UserInputMultiContent[1].Text)
assert.Equal(t, map[string]any{"k": "v"}, rewritten.UserInputMultiContent[1].Extra)
// image conversion
assert.Equal(t, schema.ChatMessagePartTypeImageURL, rewritten.UserInputMultiContent[2].Type)
assert.Equal(t, imageCommon, rewritten.UserInputMultiContent[2].Image.MessagePartCommon)
// audio conversion
assert.Equal(t, schema.ChatMessagePartTypeAudioURL, rewritten.UserInputMultiContent[3].Type)
assert.Equal(t, audioCommon, rewritten.UserInputMultiContent[3].Audio.MessagePartCommon)
// video conversion
assert.Equal(t, schema.ChatMessagePartTypeVideoURL, rewritten.UserInputMultiContent[4].Type)
assert.Equal(t, videoCommon, rewritten.UserInputMultiContent[4].Video.MessagePartCommon)
// reasoning is dropped; AssistantGenMultiContent is not set on rewritten message
assert.Empty(t, rewritten.AssistantGenMultiContent)
}
// TestTransferToAgent tests the TransferToAgent functionality
func TestTransferToAgent(t *testing.T) {
ctx := context.Background()
// Create a mock controller
ctrl := gomock.NewController(t)
defer ctrl.Finish()
// Create mock models for parent and child agents
parentModel := mockModel.NewMockToolCallingChatModel(ctrl)
childModel := mockModel.NewMockToolCallingChatModel(ctrl)
// Set up expectations for the parent model
// First call: parent model generates a message with TransferToAgent tool call
parentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
Return(schema.AssistantMessage("I'll transfer this to the child agent",
[]schema.ToolCall{
{
ID: "tool-call-1",
Function: schema.FunctionCall{
Name: TransferToAgentToolName,
Arguments: `{"agent_name": "ChildAgent"}`,
},
},
}), nil).
Times(1)
// Set up expectations for the child model
// Second call: child model generates a response
childModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
Return(schema.AssistantMessage("Hello from child agent", nil), nil).
Times(1)
// Both models should implement WithTools
parentModel.EXPECT().WithTools(gomock.Any()).Return(parentModel, nil).AnyTimes()
childModel.EXPECT().WithTools(gomock.Any()).Return(childModel, nil).AnyTimes()
// Create parent agent
parentAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
Name: "ParentAgent",
Description: "Parent agent that will transfer to child",
Instruction: "You are a parent agent.",
Model: parentModel,
})
assert.NoError(t, err)
assert.NotNil(t, parentAgent)
// Create child agent
childAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
Name: "ChildAgent",
Description: "Child agent that handles specific tasks",
Instruction: "You are a child agent.",
Model: childModel,
})
assert.NoError(t, err)
assert.NotNil(t, childAgent)
// Set up parent-child relationship
flowAgent, err := SetSubAgents(ctx, parentAgent, []Agent{childAgent})
assert.NoError(t, err)
assert.NotNil(t, flowAgent)
assert.NotNil(t, parentAgent.subAgents)
assert.NotNil(t, childAgent.parentAgent)
// Run the parent agent
input := &AgentInput{
Messages: []Message{
schema.UserMessage("Please transfer this to the child agent"),
},
}
ctx, _ = initRunCtx(ctx, flowAgent.Name(ctx), input)
iterator := flowAgent.Run(ctx, input)
assert.NotNil(t, iterator)
// First event: parent model output with tool call
event1, ok := iterator.Next()
assert.True(t, ok)
assert.NotNil(t, event1)
assert.Nil(t, event1.Err)
assert.NotNil(t, event1.Output)
assert.NotNil(t, event1.Output.MessageOutput)
assert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role)
// Second event: tool output (TransferToAgent)
event2, ok := iterator.Next()
assert.True(t, ok)
assert.NotNil(t, event2)
assert.Nil(t, event2.Err)
assert.NotNil(t, event2.Output)
assert.NotNil(t, event2.Output.MessageOutput)
assert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role)
// Verify the action is TransferToAgent
assert.NotNil(t, event2.Action)
assert.NotNil(t, event2.Action.TransferToAgent)
assert.Equal(t, "ChildAgent", event2.Action.TransferToAgent.DestAgentName)
// Third event: child model output
event3, ok := iterator.Next()
assert.True(t, ok)
assert.NotNil(t, event3)
assert.Nil(t, event3.Err)
assert.NotNil(t, event3.Output)
assert.NotNil(t, event3.Output.MessageOutput)
assert.Equal(t, schema.Assistant, event3.Output.MessageOutput.Role)
// Verify the message content from child agent
msg := event3.Output.MessageOutput.Message
assert.NotNil(t, msg)
assert.Equal(t, "Hello from child agent", msg.Content)
// No more events
_, ok = iterator.Next()
assert.False(t, ok)
}
func TestResumeNonResumableWrapperError(t *testing.T) {
ctx := context.Background()
interruptingAgent := &nonResumableFlowTestAgent{
name: "inner",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
iter, gen := NewAsyncIteratorPair[*AgentEvent]()
go func() {
defer gen.Close()
gen.Send(Interrupt(ctx, "please confirm"))
}()
return iter
},
}
innerFlowAgent := toFlowAgent(ctx, interruptingAgent)
wrapper := &nonResumableFlowTestAgent{
name: "wrapper",
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
return innerFlowAgent.Run(ctx, input, options...)
},
}
store := &flowTestStore{m: map[string][]byte{}}
runner := NewRunner(ctx, RunnerConfig{
Agent: wrapper,
EnableStreaming: true,
CheckPointStore: store,
})
iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, WithCheckPointID("cp1"))
var interruptID string
for {
ev, ok := iter.Next()
if !ok {
break
}
if ev.Action != nil && ev.Action.Interrupted != nil {
for _, intCtx := range ev.Action.Interrupted.InterruptContexts {
if intCtx.IsRootCause {
interruptID = intCtx.ID
}
}
}
}
assert.NotEmpty(t, interruptID, "should have an interrupt ID")
resumeIter, err := runner.ResumeWithParams(ctx, "cp1", &ResumeParams{
Targets: map[string]any{interruptID: nil},
})
assert.NoError(t, err)
var resumeErr error
for {
ev, ok := resumeIter.Next()
if !ok {
break
}
if ev.Err != nil {
resumeErr = ev.Err
}
}
assert.Error(t, resumeErr)
assert.Contains(t, resumeErr.Error(), "does not implement ResumableAgent interface")
assert.Contains(t, resumeErr.Error(), "custom agent wrapper must implement the ResumableAgent interface")
}
type nonResumableFlowTestAgent struct {
name string
runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]
}
func (a *nonResumableFlowTestAgent) Name(_ context.Context) string {
return a.name
}
func (a *nonResumableFlowTestAgent) Description(_ context.Context) string {
return a.name + " description"
}
func (a *nonResumableFlowTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
return a.runFn(ctx, input, options...)
}
type flowTestStore struct {
m map[string][]byte
}
func (s *flowTestStore) Set(_ context.Context, key string, value []byte) error {
s.m[key] = value
return nil
}
func (s *flowTestStore) Get(_ context.Context, key string) ([]byte, bool, error) {
v, ok := s.m[key]
return v, ok, nil
}
func TestTransferToAgentWithDesignatedCallback(t *testing.T) {
ctx := context.Background()
ctrl := gomock.NewController(t)
defer ctrl.Finish()
parentModel := mockModel.NewMockToolCallingChatModel(ctrl)
childModel := mockModel.NewMockToolCallingChatModel(ctrl)
parentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
Return(schema.AssistantMessage("I'll transfer this to the child agent",
[]schema.ToolCall{
{
ID: "tool-call-1",
Function: schema.FunctionCall{
Name: TransferToAgentToolName,
Arguments: `{"agent_name": "ChildAgent"}`,
},
},
}), nil).
Times(1)
childModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
Return(schema.AssistantMessage("Hello from child agent", nil), nil).
Times(1)
parentModel.EXPECT().WithTools(gomock.Any()).Return(parentModel, nil).AnyTimes()
childModel.EXPECT().WithTools(gomock.Any()).Return(childModel, nil).AnyTimes()
parentAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
Name: "ParentAgent",
Description: "Parent agent that will transfer to child",
Instruction: "You are a parent agent.",
Model: parentModel,
})
assert.NoError(t, err)
childAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
Name: "ChildAgent",
Description: "Child agent that handles specific tasks",
Instruction: "You are a child agent.",
Model: childModel,
})
assert.NoError(t, err)
flowAgent, err := SetSubAgents(ctx, parentAgent, []Agent{childAgent})
assert.NoError(t, err)
var childCallbackCount int
var mu sync.Mutex
handler := callbacks.NewHandlerBuilder().OnStartFn(
func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
if info.Component == ComponentOfAgent && info.Name == "ChildAgent" {
mu.Lock()
childCallbackCount++
mu.Unlock()
}
return ctx
}).Build()
input := &AgentInput{
Messages: []Message{
schema.UserMessage("Please transfer this to the child agent"),
},
}
ctx, _ = initRunCtx(ctx, flowAgent.Name(ctx), input)
iterator := flowAgent.Run(ctx, input, WithCallbacks(handler).DesignateAgent("ChildAgent"))
for {
_, ok := iterator.Next()
if !ok {
break
}
}
assert.Equal(t, 1, childCallbackCount, "designated callback for ChildAgent should fire exactly once during transfer")
}
+451
View File
@@ -0,0 +1,451 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"encoding/gob"
"fmt"
"io"
"reflect"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/compose"
"github.com/cloudwego/eino/schema"
)
// InvokableToolCallEndpoint is the function signature for invoking a tool synchronously.
// Middleware authors implement wrappers around this endpoint to add custom behavior.
type InvokableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error)
// StreamableToolCallEndpoint is the function signature for invoking a tool with streaming output.
// Middleware authors implement wrappers around this endpoint to add custom behavior.
type StreamableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error)
type EnhancedInvokableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error)
type EnhancedStreamableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error)
// ToolContext provides metadata about the tool being wrapped.
type ToolContext struct {
Name string
CallID string
}
// ToolCallsContext contains metadata about the tool calls that just completed.
type ToolCallsContext struct {
// ToolCalls contains the tool call metadata from the model's response.
ToolCalls []ToolContext
}
// TypedModelContext contains context information passed to WrapModel.
type TypedModelContext[M MessageType] struct {
// Tools contains the current tool list configured for the agent.
// This is populated at request time with the tools that will be sent to the model.
//
// Deprecated: Use TypedChatModelAgentState.ToolInfos in BeforeModelRewriteState instead.
// ModelContext.Tools remains populated for backward compatibility with existing WrapModel handlers,
// but new code should read and modify state.ToolInfos which is the source of truth for the model call.
Tools []*schema.ToolInfo
// ModelRetryConfig contains the retry configuration for the model.
// This is populated at request time from the agent's ModelRetryConfig.
// Used by EventSenderModelWrapper to wrap stream errors appropriately.
ModelRetryConfig *TypedModelRetryConfig[M]
// ModelFailoverConfig contains the failover configuration for the model.
// This is populated at request time from the agent's ModelFailoverConfig.
// Used by EventSenderModelWrapper to wrap stream errors so that failed failover
// attempts are skipped (not treated as fatal) by the flow event processor.
ModelFailoverConfig *ModelFailoverConfig[M]
cancelContext *cancelContext
}
// ModelContext is the default model context type using *schema.Message.
type ModelContext = TypedModelContext[*schema.Message]
// ChatModelAgentContext contains runtime information passed to handlers before each ChatModelAgent run.
// Handlers can modify Instruction, Tools, and ReturnDirectly to customize agent behavior.
//
// This type is specific to ChatModelAgent. Other agent types may define their own context types.
type ChatModelAgentContext struct {
// Instruction is the current instruction for the Agent execution.
// It includes the instruction configured for the agent, additional instructions appended by framework
// and AgentMiddleware, and modifications applied by previous BeforeAgent handlers.
// The finalized instruction after all BeforeAgent handlers are then passed to GenModelInput,
// to be (optionally) formatted with SessionValues and converted to system message.
Instruction string
// Tools are the raw tools (without any wrapper or tool middleware) currently configured for the Agent execution.
// They includes tools passed in AgentConfig, implicit tools added by framework such as transfer / exit tools,
// and other tools already added by middlewares.
Tools []tool.BaseTool
// ReturnDirectly is the set of tool names currently configured to cause the Agent to return directly.
// This is based on the return directly map configured for the agent, plus any modifications
// by previous BeforeAgent handlers.
ReturnDirectly map[string]bool
// ToolSearchTool is the tool info for the model's native tool search capability.
// When set by a BeforeAgent handler, the framework passes it to the model via model.WithToolSearchTool.
ToolSearchTool *schema.ToolInfo
}
// TypedChatModelAgentMiddleware defines the interface for customizing TypedChatModelAgent behavior.
//
// IMPORTANT: This interface is specifically designed for TypedChatModelAgent and agents built
// on top of it (e.g., DeepAgent).
//
// Why TypedChatModelAgentMiddleware instead of AgentMiddleware?
//
// AgentMiddleware is a struct type, which has inherent limitations:
// - Struct types are closed: users cannot add new methods to extend functionality
// - The framework only recognizes AgentMiddleware's fixed fields, so even if users
// embed AgentMiddleware in a custom struct and add methods, the framework cannot
// call those methods (config.Middlewares is []AgentMiddleware, not a user type)
// - Callbacks in AgentMiddleware only return error, cannot return modified context
//
// TypedChatModelAgentMiddleware is an interface type, which is open for extension:
// - Users can implement custom handlers with arbitrary internal state and methods
// - Hook methods return (context.Context, ..., error) for direct context propagation
// - Wrapper methods (WrapToolCall, WrapModel) enable context propagation through the
// wrapped endpoint chain: wrappers can pass modified context to the next wrapper
// - Configuration is centralized in struct fields rather than scattered in closures
//
// TypedChatModelAgentMiddleware vs AgentMiddleware:
// - Use AgentMiddleware for simple, static additions (extra instruction/tools)
// - Use TypedChatModelAgentMiddleware for dynamic behavior, context modification, or call wrapping
// - AgentMiddleware is kept for backward compatibility with existing users
// - Both can be used together; see AgentMiddleware documentation for execution order
//
// Use *TypedBaseChatModelAgentMiddleware as an embedded struct to provide default no-op
// implementations for all methods.
type TypedChatModelAgentMiddleware[M MessageType] interface {
// BeforeAgent is called before each agent run, allowing modification of
// the agent's instruction and tools configuration.
BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error)
// AfterAgent is called after the agent run reaches a successful terminal state.
// Successful terminal states are: final answer (model response with no tool calls),
// and return-directly tool result.
//
// AfterAgent is NOT called when the agent terminates with an error (e.g.,
// ErrExceedMaxIterations, context cancellation, model errors).
//
// The state parameter contains the final conversation state, including all messages
// from the completed run.
//
// AfterAgent handlers are called in the same order as BeforeAgent handlers
// (first registered = first called). Consistent with all other middleware hooks,
// if any handler returns an error, subsequent handlers are NOT called (fail-fast)
// and the error is sent to the event stream.
AfterAgent(ctx context.Context, state *TypedChatModelAgentState[M]) (context.Context, error)
// BeforeModelRewriteState is called before each model invocation.
// The returned state is persisted to the agent's internal state and passed to the model.
// The returned context is propagated to the model call and subsequent handlers.
//
// The ChatModelAgentState struct provides access to:
// - Messages: the conversation history
// - ToolInfos: the tool list that will be sent to the model (modifiable)
// - DeferredToolInfos: tools for server-side search (modifiable, nil if unused)
//
// This is the recommended place to modify messages and tools before a model call.
// Changes here are persisted in state and reflected in subsequent iterations.
BeforeModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedChatModelAgentState[M], error)
// AfterModelRewriteState is called after each model invocation.
// The input state includes the model's response as the last message.
// The returned state is persisted to the agent's internal state.
//
// The ChatModelAgentState struct provides access to:
// - Messages: the conversation history including the model's response
// - ToolInfos: the tool list that was sent to the model
// - DeferredToolInfos: tools for server-side search (nil if unused)
AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedChatModelAgentState[M], error)
// WrapInvokableToolCall wraps a tool's synchronous execution with custom behavior.
// Return the input endpoint unchanged and nil error if no wrapping is needed.
//
// This method is only called for tools that implement InvokableTool.
// If a tool only implements StreamableTool, this method will not be called for that tool.
//
// This method is called at request time when the tool is about to be executed.
// The tCtx parameter provides metadata about the tool:
// - Name: The name of the tool being wrapped
// - CallID: The unique identifier for this specific tool call
WrapInvokableToolCall(ctx context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error)
// WrapStreamableToolCall wraps a tool's streaming execution with custom behavior.
// Return the input endpoint unchanged and nil error if no wrapping is needed.
//
// This method is only called for tools that implement StreamableTool.
// If a tool only implements InvokableTool, this method will not be called for that tool.
//
// This method is called at request time when the tool is about to be executed.
// The tCtx parameter provides metadata about the tool:
// - Name: The name of the tool being wrapped
// - CallID: The unique identifier for this specific tool call
WrapStreamableToolCall(ctx context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error)
// WrapEnhancedInvokableToolCall wraps an enhanced tool's synchronous execution with custom behavior.
// Return the input endpoint unchanged and nil error if no wrapping is needed.
//
// This method is only called for tools that implement EnhancedInvokableTool.
// If a tool only implements EnhancedStreamableTool, this method will not be called for that tool.
//
// This method is called at request time when the tool is about to be executed.
// The tCtx parameter provides metadata about the tool:
// - Name: The name of the tool being wrapped
// - CallID: The unique identifier for this specific tool call
WrapEnhancedInvokableToolCall(ctx context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error)
// WrapEnhancedStreamableToolCall wraps an enhanced tool's streaming execution with custom behavior.
// Return the input endpoint unchanged and nil error if no wrapping is needed.
//
// This method is only called for tools that implement EnhancedStreamableTool.
// If a tool only implements EnhancedInvokableTool, this method will not be called for that tool.
//
// This method is called at request time when the tool is about to be executed.
// The tCtx parameter provides metadata about the tool:
// - Name: The name of the tool being wrapped
// - CallID: The unique identifier for this specific tool call
WrapEnhancedStreamableToolCall(ctx context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error)
// WrapModel wraps a chat model with custom behavior around the actual model call.
// Return the input model unchanged and nil error if no wrapping is needed.
//
// This method is called at request time when the model is about to be invoked.
// Note: The parameter is model.BaseModel[M] (not ToolCallingChatModel) because wrappers
// only need to intercept Generate/Stream calls. Tool binding (WithTools) is handled
// separately by the framework and does not flow through user wrappers.
//
// Recommended use cases (behavior around the model call itself):
// - Model call retry logic
// - Model failover (switching to a backup model)
// - Sending events (e.g. streaming progress)
// - Processing or transforming the response stream
// - Changing call configurations (temperature, top_p, etc.)
//
// Discouraged use cases (use BeforeModelRewriteState instead):
// - Modifying input messages: changes here are NOT persisted in state, only
// affect a single model call, and break prompt cache across iterations.
// - Modifying the tool list: use state.ToolInfos / state.DeferredToolInfos in
// BeforeModelRewriteState, which is the source of truth for tool configuration.
//
// The mc parameter provides read-only context about the current model call:
// - Tools: The tool infos that will be sent to the model (Deprecated: read state.ToolInfos instead)
WrapModel(ctx context.Context, m model.BaseModel[M], mc *TypedModelContext[M]) (model.BaseModel[M], error)
}
// ChatModelAgentMiddleware is the default middleware type using *schema.Message.
// See TypedChatModelAgentMiddleware for full documentation.
type ChatModelAgentMiddleware = TypedChatModelAgentMiddleware[*schema.Message]
type TypedBaseChatModelAgentMiddleware[M MessageType] struct{}
// BaseChatModelAgentMiddleware provides default no-op implementations for ChatModelAgentMiddleware.
// Embed *BaseChatModelAgentMiddleware in custom handlers to only override the methods you need.
//
// Example:
//
// type MyHandler struct {
// *adk.BaseChatModelAgentMiddleware
// // custom fields
// }
//
// func (h *MyHandler) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {
// // custom logic
// return ctx, state, nil
// }
type BaseChatModelAgentMiddleware = TypedBaseChatModelAgentMiddleware[*schema.Message]
func (b *TypedBaseChatModelAgentMiddleware[M]) WrapInvokableToolCall(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) (InvokableToolCallEndpoint, error) {
return endpoint, nil
}
func (b *TypedBaseChatModelAgentMiddleware[M]) WrapStreamableToolCall(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) (StreamableToolCallEndpoint, error) {
return endpoint, nil
}
func (b *TypedBaseChatModelAgentMiddleware[M]) WrapEnhancedInvokableToolCall(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) (EnhancedInvokableToolCallEndpoint, error) {
return endpoint, nil
}
func (b *TypedBaseChatModelAgentMiddleware[M]) WrapEnhancedStreamableToolCall(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) (EnhancedStreamableToolCallEndpoint, error) {
return endpoint, nil
}
func (b *TypedBaseChatModelAgentMiddleware[M]) WrapModel(_ context.Context, m model.BaseModel[M], _ *TypedModelContext[M]) (model.BaseModel[M], error) {
return m, nil
}
func (b *TypedBaseChatModelAgentMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {
return ctx, runCtx, nil
}
func (b *TypedBaseChatModelAgentMiddleware[M]) AfterAgent(ctx context.Context, state *TypedChatModelAgentState[M]) (context.Context, error) {
return ctx, nil
}
func (b *TypedBaseChatModelAgentMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedChatModelAgentState[M], error) {
return ctx, state, nil
}
func (b *TypedBaseChatModelAgentMiddleware[M]) AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedChatModelAgentState[M], error) {
return ctx, state, nil
}
func processTypedState(ctx context.Context, fn func(extra map[string]any) map[string]any) error {
runCtx := getRunCtx(ctx)
if runCtx != nil && runCtx.AgenticRootInput != nil {
return compose.ProcessState(ctx, func(_ context.Context, st *typedState[*schema.AgenticMessage]) error {
st.Extra = fn(st.Extra)
return nil
})
}
return compose.ProcessState(ctx, func(_ context.Context, st *typedState[*schema.Message]) error {
st.Extra = fn(st.Extra)
return nil
})
}
// SetRunLocalValue sets a key-value pair that persists for the duration of the current agent Run() invocation.
// The value is scoped to this specific execution and is not shared across different Run() calls or agent instances.
//
// Values stored here are compatible with interrupt/resume cycles - they will be serialized and restored
// when the agent is resumed. For custom types, you must register them using schema.RegisterName[T]()
// in an init() function to ensure proper serialization.
//
// This function can only be called from within a ChatModelAgentMiddleware during agent execution.
// Returns an error if called outside of an agent execution context.
func SetRunLocalValue(ctx context.Context, key string, value any) error {
if err := checkGobEncodability(key, value); err != nil {
return err
}
err := processTypedState(ctx, func(extra map[string]any) map[string]any {
if extra == nil {
extra = make(map[string]any)
}
extra[key] = value
return extra
})
if err != nil {
return fmt.Errorf("SetRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err)
}
return nil
}
// GetRunLocalValue retrieves a value that was set during the current agent Run() invocation.
// The value is scoped to this specific execution and is not shared across different Run() calls or agent instances.
//
// Values stored via SetRunLocalValue are compatible with interrupt/resume cycles - they will be serialized
// and restored when the agent is resumed. For custom types, you must register them using schema.RegisterName[T]()
// in an init() function to ensure proper serialization.
//
// This function can only be called from within a ChatModelAgentMiddleware during agent execution.
// Returns the value and true if found, or nil and false if not found or if called outside of an agent execution context.
func GetRunLocalValue(ctx context.Context, key string) (any, bool, error) {
var val any
var found bool
err := processTypedState(ctx, func(extra map[string]any) map[string]any {
if extra != nil {
val, found = extra[key]
}
return extra
})
if err != nil {
return nil, false, fmt.Errorf("GetRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err)
}
return val, found, nil
}
// DeleteRunLocalValue removes a value that was set during the current agent Run() invocation.
//
// This function can only be called from within a ChatModelAgentMiddleware during agent execution.
// Returns an error if called outside of an agent execution context.
func DeleteRunLocalValue(ctx context.Context, key string) error {
err := processTypedState(ctx, func(extra map[string]any) map[string]any {
if extra != nil {
delete(extra, key)
}
return extra
})
if err != nil {
return fmt.Errorf("DeleteRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err)
}
return nil
}
// TypedSendEvent sends a custom TypedAgentEvent to the event stream during agent execution.
// This allows TypedChatModelAgentMiddleware implementations to emit custom events that will be
// received by the caller iterating over the agent's event stream.
//
// Note: TypedSendEvent is a pure transport — it does NOT auto-assign message IDs.
// Framework-created messages (model output, tool results) receive IDs automatically
// via internal wrapper layers. If your middleware constructs its own messages, call
// EnsureMessageID before sending to assign an ID.
//
// This function can only be called from within a TypedChatModelAgentMiddleware during agent execution.
// Returns an error if called outside of an agent execution context.
func TypedSendEvent[M MessageType](ctx context.Context, event *TypedAgentEvent[M]) error {
execCtx := getTypedChatModelAgentExecCtx[M](ctx)
if execCtx == nil || execCtx.generator == nil {
return fmt.Errorf("TypedSendEvent failed: must be called within a ChatModelAgent Run() or Resume() execution context")
}
execCtx.send(event)
return nil
}
// SendEvent sends a custom AgentEvent to the event stream during agent execution.
// This allows ChatModelAgentMiddleware implementations to emit custom events that will be
// received by the caller iterating over the agent's event stream.
//
// This function can only be called from within a ChatModelAgentMiddleware during agent execution.
// Returns an error if called outside of an agent execution context.
func SendEvent(ctx context.Context, event *AgentEvent) error {
return TypedSendEvent(ctx, event)
}
// checkGobEncodability probes whether the value can be gob-encoded as part of
// a map[string]any, which is exactly how State.Extra is serialized during
// checkpoint. This catches unregistered types early at Set time, rather than
// letting them fail at checkpoint/resume time with a confusing error.
func checkGobEncodability(key string, value any) error {
probe := map[string]any{key: value}
if err := gob.NewEncoder(io.Discard).Encode(probe); err != nil {
typeName := reflect.TypeOf(value).String()
return fmt.Errorf("SetRunLocalValue: the value (type %s) for key %q is not gob-serializable, "+
"which means it will fail when the agent checkpoint is saved or resumed.\n\n"+
"To fix this, register the type in an init() function in your package:\n\n"+
" func init() {\n"+
" schema.RegisterName[%s](\"a_unique_name_for_this_type\")\n"+
" }\n\n"+
"This is required because agent state (including values set via SetRunLocalValue) is "+
"persisted using gob encoding for interrupt/resume support. All concrete types stored "+
"in interface-typed fields (like map[string]any) must be registered with gob.\n\n"+
"If this value does not need to survive interrupt/resume, store it on the context instead, "+
"for example via context.WithValue, so you don't need gob registration.\n\n"+
"Underlying error: %w", typeName, key, typeName, err)
}
return nil
}
+2571
View File
File diff suppressed because it is too large Load Diff
+64
View File
@@ -0,0 +1,64 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"context"
"fmt"
"strings"
"github.com/cloudwego/eino/adk/internal"
)
const (
TransferToAgentInstruction = `Available other agents: %s
Decision rule:
- If you're best suited for the question according to your description: ANSWER
- If another agent is better according its description: CALL '%s' function with their agent name
When transferring: OUTPUT ONLY THE FUNCTION CALL`
TransferToAgentInstructionChinese = `可用的其他 agent%s
决策规则:
- 如果根据你的职责描述,你最适合回答这个问题:ANSWER
- 如果根据其职责描述,另一个 agent 更适合:调用 %s 函数,并传入该 agent 的名称
当进行移交时:只输出函数调用,不要输出其他任何内容`
agentDescriptionTpl = "\n- Agent name: %s\n Agent description: %s"
agentDescriptionTplChinese = "\n- Agent 名字: %s\n Agent 描述: %s"
)
func genTransferToAgentInstruction[M MessageType](ctx context.Context, agents []TypedAgent[M]) string {
tpl := internal.SelectPrompt(internal.I18nPrompts{
English: agentDescriptionTpl,
Chinese: agentDescriptionTplChinese,
})
instruction := internal.SelectPrompt(internal.I18nPrompts{
English: TransferToAgentInstruction,
Chinese: TransferToAgentInstructionChinese,
})
var sb strings.Builder
for _, agent := range agents {
sb.WriteString(fmt.Sprintf(tpl, agent.Name(ctx), agent.Description(ctx)))
}
return fmt.Sprintf(instruction, sb.String(), TransferToAgentToolName)
}
+519
View File
@@ -0,0 +1,519 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"bytes"
"context"
"encoding/gob"
"fmt"
"io"
"github.com/cloudwego/eino/components"
"github.com/cloudwego/eino/internal/core"
"github.com/cloudwego/eino/schema"
)
// ComponentOfAgent is the component type identifier for ADK agents in callbacks.
// Use this to filter callback events to only agent-related events.
const ComponentOfAgent components.Component = "Agent"
// ComponentOfAgenticAgent is the component type identifier for ADK agents
// that use *schema.AgenticMessage in callbacks.
const ComponentOfAgenticAgent components.Component = "AgenticAgent"
// MessageType is the sealed type constraint for message types used in ADK.
// Only *schema.Message and *schema.AgenticMessage satisfy this constraint.
// External packages cannot add new types to this union; all generic functions
// in ADK use exhaustive type switches on these two types.
type MessageType interface {
*schema.Message | *schema.AgenticMessage
}
type Message = *schema.Message
type MessageStream = *schema.StreamReader[Message]
type AgenticMessage = *schema.AgenticMessage
type AgenticMessageStream = *schema.StreamReader[AgenticMessage]
// isNilMessage checks whether a generic message value is nil.
// Direct `msg == nil` does not compile for generic pointer types in Go;
// the canonical workaround is to compare through the `any` interface.
func isNilMessage[M MessageType](msg M) bool {
var zero M
return any(msg) == any(zero)
}
// TypedMessageVariant represents a message output from an agent event.
// It carries either a complete message or a streaming reader, along with
// metadata describing the event's origin.
//
// Role and ToolName are only meaningful for *schema.Message events. For
// *schema.AgenticMessage events (created via EventFromAgenticMessage), these
// fields are always zero-valued because AgenticMessage carries tool results as
// ContentBlocks within the message itself and does not support agent transfer.
//
// For *schema.Message events, Role and ToolName exist independently of the inner
// Message because in streaming mode (IsStreaming=true, Message=nil), the message
// has not materialized yet and the consumer needs metadata without consuming the stream.
type TypedMessageVariant[M MessageType] struct {
IsStreaming bool
Message M
MessageStream *schema.StreamReader[M]
// Role indicates the origin of this event within the agent's ReAct loop.
// Only meaningful for *schema.Message events:
// - schema.Assistant: the event carries model output (generation or stream).
// - schema.Tool: the event carries a tool execution result.
// Always zero-valued for *schema.AgenticMessage events; use AgenticRole instead.
Role schema.RoleType
// AgenticRole indicates the role of the agentic message (assistant, user, system).
// Only meaningful for *schema.AgenticMessage events.
// In streaming mode, this is available before consuming the stream.
// Always zero-valued for *schema.Message events; use Role instead.
AgenticRole schema.AgenticRoleType
// ToolName is the name of the tool that produced this event.
// Only meaningful for *schema.Message events: non-empty when Role == schema.Tool.
// In streaming mode, this is the only way to identify the source tool before
// the stream is consumed.
// Always empty for *schema.AgenticMessage events.
ToolName string
}
func (mv *TypedMessageVariant[M]) GetMessage() (M, error) {
if mv.IsStreaming {
return concatMessageStream(mv.MessageStream)
}
return mv.Message, nil
}
type MessageVariant = TypedMessageVariant[*schema.Message]
type messageVariantSerialization struct {
IsStreaming bool
Message Message
MessageStream Message
Role schema.RoleType
ToolName string
}
type agenticMessageVariantSerialization struct {
IsStreaming bool
Message *schema.AgenticMessage
MessageStream *schema.AgenticMessage
Role schema.RoleType
AgenticRole schema.AgenticRoleType
ToolName string
}
func (mv *TypedMessageVariant[M]) GobEncode() ([]byte, error) {
if mvMsg, ok := any(mv).(*TypedMessageVariant[*schema.Message]); ok {
return gobEncodeMessageVariant(mvMsg)
}
if mvAgentic, ok := any(mv).(*TypedMessageVariant[*schema.AgenticMessage]); ok {
return gobEncodeAgenticMessageVariant(mvAgentic)
}
return nil, fmt.Errorf("gob encoding not supported for this message type")
}
func (mv *TypedMessageVariant[M]) GobDecode(b []byte) error {
if mvMsg, ok := any(mv).(*TypedMessageVariant[*schema.Message]); ok {
return gobDecodeMessageVariant(mvMsg, b)
}
if mvAgentic, ok := any(mv).(*TypedMessageVariant[*schema.AgenticMessage]); ok {
return gobDecodeAgenticMessageVariant(mvAgentic, b)
}
return fmt.Errorf("gob decoding not supported for this message type")
}
func gobEncodeMessageVariant(mv *TypedMessageVariant[*schema.Message]) ([]byte, error) {
s := &messageVariantSerialization{
IsStreaming: mv.IsStreaming,
Message: mv.Message,
Role: mv.Role,
ToolName: mv.ToolName,
}
if mv.IsStreaming {
var messages []Message
for {
frame, err := mv.MessageStream.Recv()
if err == io.EOF {
break
}
if err != nil {
return nil, fmt.Errorf("error receiving message stream: %w", err)
}
messages = append(messages, frame)
}
m, err := schema.ConcatMessages(messages)
if err != nil {
return nil, fmt.Errorf("failed to encode message: cannot concat message stream: %w", err)
}
s.MessageStream = m
}
buf := &bytes.Buffer{}
err := gob.NewEncoder(buf).Encode(s)
if err != nil {
return nil, fmt.Errorf("failed to gob encode message variant: %w", err)
}
return buf.Bytes(), nil
}
func gobDecodeMessageVariant(mv *TypedMessageVariant[*schema.Message], b []byte) error {
s := &messageVariantSerialization{}
err := gob.NewDecoder(bytes.NewReader(b)).Decode(s)
if err != nil {
return fmt.Errorf("failed to decoding message variant: %w", err)
}
mv.IsStreaming = s.IsStreaming
mv.Message = s.Message
mv.Role = s.Role
mv.ToolName = s.ToolName
if s.MessageStream != nil {
mv.MessageStream = schema.StreamReaderFromArray([]*schema.Message{s.MessageStream})
}
return nil
}
func gobEncodeAgenticMessageVariant(mv *TypedMessageVariant[*schema.AgenticMessage]) ([]byte, error) {
s := &agenticMessageVariantSerialization{
IsStreaming: mv.IsStreaming,
Message: mv.Message,
Role: mv.Role,
AgenticRole: mv.AgenticRole,
ToolName: mv.ToolName,
}
if mv.IsStreaming {
var messages []*schema.AgenticMessage
for {
frame, err := mv.MessageStream.Recv()
if err == io.EOF {
break
}
if err != nil {
return nil, fmt.Errorf("error receiving agentic message stream: %w", err)
}
messages = append(messages, frame)
}
m, err := schema.ConcatAgenticMessages(messages)
if err != nil {
return nil, fmt.Errorf("failed to encode agentic message: cannot concat message stream: %w", err)
}
s.MessageStream = m
}
buf := &bytes.Buffer{}
err := gob.NewEncoder(buf).Encode(s)
if err != nil {
return nil, fmt.Errorf("failed to gob encode agentic message variant: %w", err)
}
return buf.Bytes(), nil
}
func gobDecodeAgenticMessageVariant(mv *TypedMessageVariant[*schema.AgenticMessage], b []byte) error {
s := &agenticMessageVariantSerialization{}
err := gob.NewDecoder(bytes.NewReader(b)).Decode(s)
if err != nil {
return fmt.Errorf("failed to decode agentic message variant: %w", err)
}
mv.IsStreaming = s.IsStreaming
mv.Message = s.Message
mv.Role = s.Role
mv.AgenticRole = s.AgenticRole
mv.ToolName = s.ToolName
if s.MessageStream != nil {
mv.MessageStream = schema.StreamReaderFromArray([]*schema.AgenticMessage{s.MessageStream})
}
return nil
}
// typedEventFromMessage creates a TypedAgentEvent containing the given message and optional stream.
func typedEventFromMessage[M MessageType](msg M, msgStream *schema.StreamReader[M],
role schema.RoleType, toolName string) *TypedAgentEvent[M] {
return &TypedAgentEvent[M]{
Output: &TypedAgentOutput[M]{
MessageOutput: &TypedMessageVariant[M]{
IsStreaming: msgStream != nil,
Message: msg,
MessageStream: msgStream,
Role: role,
ToolName: toolName,
},
},
}
}
// typedModelOutputEvent creates a model-output event for the generic path.
// For *schema.Message, Role is set to schema.Assistant.
// For *schema.AgenticMessage, AgenticRole is set to schema.AgenticRoleTypeAssistant.
func typedModelOutputEvent[M MessageType](msg M, msgStream *schema.StreamReader[M]) *TypedAgentEvent[M] {
var role schema.RoleType
var agenticRole schema.AgenticRoleType
var zero M
if _, ok := any(zero).(*schema.Message); ok {
role = schema.Assistant
} else {
agenticRole = schema.AgenticRoleTypeAssistant
}
event := typedEventFromMessage(msg, msgStream, role, "")
event.Output.MessageOutput.AgenticRole = agenticRole
return event
}
// EventFromMessage creates an AgentEvent containing the given message and optional stream.
//
// role identifies the origin of this event:
// - schema.Assistant: model output (generation or stream).
// - schema.Tool: tool execution result; toolName must be non-empty.
//
// For *schema.AgenticMessage events, use EventFromAgenticMessage instead.
func EventFromMessage(msg Message, msgStream *schema.StreamReader[Message],
role schema.RoleType, toolName string) *AgentEvent {
return typedEventFromMessage(msg, msgStream, role, toolName)
}
// EventFromAgenticMessage creates a TypedAgentEvent for the AgenticMessage path.
// Unlike EventFromMessage, it does not require role or toolName parameters because
// AgenticMessage carries tool results as ContentBlocks within the message itself,
// and does not support agent transfer.
//
// agenticRole identifies the role of the message (e.g. schema.AgenticRoleTypeAssistant).
// In streaming mode, the role is available on the event before consuming the stream.
func EventFromAgenticMessage(msg AgenticMessage, msgStream AgenticMessageStream, agenticRole schema.AgenticRoleType) *TypedAgentEvent[AgenticMessage] {
return &TypedAgentEvent[AgenticMessage]{
Output: &TypedAgentOutput[AgenticMessage]{
MessageOutput: &TypedMessageVariant[AgenticMessage]{
IsStreaming: msgStream != nil,
Message: msg,
MessageStream: msgStream,
AgenticRole: agenticRole,
},
},
}
}
// TransferToAgentAction represents a transfer-to-agent action.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
type TransferToAgentAction struct {
DestAgentName string
}
type TypedAgentOutput[M MessageType] struct {
MessageOutput *TypedMessageVariant[M]
CustomizedOutput any
}
type AgentOutput = TypedAgentOutput[*schema.Message]
// NewTransferToAgentAction creates an action to transfer to the specified agent.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func NewTransferToAgentAction(destAgentName string) *AgentAction {
return &AgentAction{TransferToAgent: &TransferToAgentAction{DestAgentName: destAgentName}}
}
// NewExitAction creates an action that signals the agent to exit.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func NewExitAction() *AgentAction {
return &AgentAction{Exit: true}
}
// AgentAction represents actions that an agent can emit during execution.
//
// Action Scoping in Agent Tools:
// When an agent is wrapped as an agent tool (via NewAgentTool), actions emitted by the inner agent
// are scoped to the tool boundary:
// - Interrupted: Propagated via CompositeInterrupt to allow proper interrupt/resume across boundaries
// - Exit, TransferToAgent, BreakLoop: Ignored outside the agent tool; these actions only affect
// the inner agent's execution and do not propagate to the parent agent
//
// This scoping ensures that nested agents cannot unexpectedly terminate or transfer control
// of their parent agent's execution flow.
type AgentAction struct {
Exit bool
Interrupted *InterruptInfo
TransferToAgent *TransferToAgentAction
BreakLoop *BreakLoopAction
CustomizedAction any
internalInterrupted *core.InterruptSignal
}
// RunStep represents a step in the agent execution path.
// CheckpointSchema: persisted via serialization.RunCtx (gob).
//
// NOT RECOMMENDED: RunStep is mainly relevant for agent transfer and workflow agents,
// which have not proven to be more effective empirically. Consider using ChatModelAgent
// with AgentTool or DeepAgent instead for most multi-agent scenarios.
type RunStep struct {
agentName string
}
func init() {
schema.RegisterName[[]RunStep]("eino_run_step_list")
}
func (r *RunStep) String() string {
return r.agentName
}
func (r *RunStep) Equals(r1 RunStep) bool {
return r.agentName == r1.agentName
}
func (r *RunStep) GobEncode() ([]byte, error) {
s := &runStepSerialization{AgentName: r.agentName}
buf := &bytes.Buffer{}
err := gob.NewEncoder(buf).Encode(s)
if err != nil {
return nil, fmt.Errorf("failed to gob encode RunStep: %w", err)
}
return buf.Bytes(), nil
}
func (r *RunStep) GobDecode(b []byte) error {
s := &runStepSerialization{}
err := gob.NewDecoder(bytes.NewReader(b)).Decode(s)
if err != nil {
return fmt.Errorf("failed to gob decode RunStep: %w", err)
}
r.agentName = s.AgentName
return nil
}
type runStepSerialization struct {
AgentName string
}
// TypedAgentEvent represents a single event emitted during agent execution.
// CheckpointSchema: persisted via serialization.RunCtx (gob).
type TypedAgentEvent[M MessageType] struct {
AgentName string
// RunPath represents the execution path from root agent to the current event source.
// This field is managed entirely by the framework and cannot be set by end-users.
//
// NOT RECOMMENDED: RunPath is mainly relevant for agent transfer and workflow agents,
// which have not proven to be more effective empirically. For ChatModelAgent with
// AgentTool or DeepAgent, RunPath is trivial. Consider those patterns instead.
RunPath []RunStep
Output *TypedAgentOutput[M]
Action *AgentAction
Err error
}
// AgentEvent is the default event type using *schema.Message.
type AgentEvent = TypedAgentEvent[*schema.Message]
type TypedAgentInput[M MessageType] struct {
Messages []M
EnableStreaming bool
}
type AgentInput = TypedAgentInput[*schema.Message]
// TypedAgent is the base agent interface parameterized by message type.
//
// For M = *schema.Message, the full ADK feature set is supported (multi-agent
// orchestration, cancel monitoring, retry, flowAgent).
// For M = *schema.AgenticMessage, single-agent execution works but cancel
// monitoring on the model stream and retry are not yet wired.
type TypedAgent[M MessageType] interface {
Name(ctx context.Context) string
Description(ctx context.Context) string
// Run runs the agent.
// The returned AgentEvent within the AsyncIterator must be safe to modify.
// If the returned AgentEvent within the AsyncIterator contains MessageStream,
// the MessageStream MUST be exclusive and safe to be received directly.
// NOTE: it's recommended to use SetAutomaticClose() on the MessageStream of AgentEvents emitted by AsyncIterator,
// so that even the events are not processed, the MessageStream can still be closed.
Run(ctx context.Context, input *TypedAgentInput[M], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]]
}
//go:generate mockgen -destination ../internal/mock/adk/Agent_mock.go --package adk github.com/cloudwego/eino/adk Agent,ResumableAgent
type Agent = TypedAgent[*schema.Message]
// OnSubAgents is the interface for agents that support sub-agent registration and transfer.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
type OnSubAgents interface {
OnSetSubAgents(ctx context.Context, subAgents []Agent) error
OnSetAsSubAgent(ctx context.Context, parent Agent) error
OnDisallowTransferToParent(ctx context.Context) error
}
type TypedResumableAgent[M MessageType] interface {
TypedAgent[M]
Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]]
}
type ResumableAgent = TypedResumableAgent[*schema.Message]
func concatMessageStream[M MessageType](stream *schema.StreamReader[M]) (M, error) {
var zero M
switch s := any(stream).(type) {
case *schema.StreamReader[*schema.Message]:
result, err := schema.ConcatMessageStream(s)
if err != nil {
return zero, err
}
return any(result).(M), nil
case *schema.StreamReader[*schema.AgenticMessage]:
defer s.Close()
var msgs []*schema.AgenticMessage
for {
frame, err := s.Recv()
if err == io.EOF {
break
}
if err != nil {
return zero, err
}
msgs = append(msgs, frame)
}
result, err := schema.ConcatAgenticMessages(msgs)
if err != nil {
return zero, err
}
return any(result).(M), nil
default:
panic("unreachable: unknown MessageType")
}
}
+76
View File
@@ -0,0 +1,76 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Package internal provides adk internal utils.
package internal
import (
"fmt"
"sync/atomic"
)
// Language represents the language setting for the ADK built-in prompts.
type Language uint8
const (
// LanguageEnglish represents English language.
LanguageEnglish Language = iota
// LanguageChinese represents Chinese language.
LanguageChinese
)
var language atomic.Value
// SetLanguage sets the language for the ADK built-in prompts.
// The default language is English if not explicitly set.
func SetLanguage(lang Language) error {
if lang != LanguageEnglish &&
lang != LanguageChinese {
return fmt.Errorf("invalid language: %v", lang)
}
language.Store(lang)
return nil
}
// GetLanguage returns the current language setting for the ADK built-in prompts.
// Returns LanguageEnglish if no language has been set.
func getLanguage() Language {
if l, ok := language.Load().(Language); ok {
return l
}
return LanguageEnglish
}
// I18nPrompts holds prompt strings for different languages.
type I18nPrompts struct {
English string
Chinese string
}
// SelectPrompt returns the appropriate prompt string based on the current language setting.
// Returns an error if the current language is not supported.
func SelectPrompt(prompts I18nPrompts) string {
lang := getLanguage()
switch lang {
case LanguageEnglish:
return prompts.English
case LanguageChinese:
return prompts.Chinese
default:
// unreachable
panic(fmt.Sprintf("invalid language: %v", lang))
}
}
+58
View File
@@ -0,0 +1,58 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package internal
import "github.com/google/uuid"
// EinoMsgIDKey is the Extra key used to store the eino-internal message ID.
const EinoMsgIDKey = "_eino_msg_id"
// GetMessageID returns the message ID from Extra, or "" if not set.
// Works with any map[string]any (Message.Extra or AgenticMessage.Extra).
func GetMessageID(extra map[string]any) string {
if extra == nil {
return ""
}
id, _ := extra[EinoMsgIDKey].(string)
return id
}
// SetMessageID sets the message ID in Extra and returns the resulting map.
//
// Copy-on-write: the input map is never mutated, because a message's Extra can
// be shared across fan-out stream readers and an in-place write would race with
// concurrent reads (fatal "concurrent map read and map write"). Reassigning the
// returned map to a shared Extra field is still a field-level race that cannot
// be eliminated while Extra is exported; COW only removes the map-level panic.
func SetMessageID(extra map[string]any, id string) map[string]any {
next := make(map[string]any, len(extra)+1)
for k, v := range extra {
next[k] = v
}
next[EinoMsgIDKey] = id
return next
}
// EnsureMessageID assigns a UUID v4 if no message ID is present.
// Idempotent: if ID already set, no-op.
// Returns the (possibly newly created) Extra map.
func EnsureMessageID(extra map[string]any) map[string]any {
if GetMessageID(extra) != "" {
return extra
}
return SetMessageID(extra, uuid.NewString())
}
+87
View File
@@ -0,0 +1,87 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package internal
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestGetMessageID(t *testing.T) {
t.Run("nil extra returns empty", func(t *testing.T) {
assert.Equal(t, "", GetMessageID(nil))
})
t.Run("empty extra returns empty", func(t *testing.T) {
assert.Equal(t, "", GetMessageID(map[string]any{}))
})
t.Run("wrong type returns empty", func(t *testing.T) {
extra := map[string]any{EinoMsgIDKey: 123}
assert.Equal(t, "", GetMessageID(extra))
})
t.Run("returns set ID", func(t *testing.T) {
extra := map[string]any{EinoMsgIDKey: "test-id-123"}
assert.Equal(t, "test-id-123", GetMessageID(extra))
})
}
func TestSetMessageID(t *testing.T) {
t.Run("nil extra creates map", func(t *testing.T) {
extra := SetMessageID(nil, "id-1")
assert.NotNil(t, extra)
assert.Equal(t, "id-1", extra[EinoMsgIDKey])
})
t.Run("existing extra preserved", func(t *testing.T) {
extra := map[string]any{"other_key": "other_val"}
result := SetMessageID(extra, "id-2")
assert.Equal(t, "id-2", result[EinoMsgIDKey])
assert.Equal(t, "other_val", result["other_key"])
})
}
func TestEnsureMessageID(t *testing.T) {
t.Run("nil extra gets ID", func(t *testing.T) {
extra := EnsureMessageID(nil)
id := GetMessageID(extra)
assert.NotEmpty(t, id)
assert.Len(t, id, 36) // UUID v4 format: 8-4-4-4-12 = 36 chars
})
t.Run("idempotent - does not overwrite existing ID", func(t *testing.T) {
extra := SetMessageID(nil, "existing-id")
result := EnsureMessageID(extra)
assert.Equal(t, "existing-id", GetMessageID(result))
})
t.Run("empty extra gets new ID", func(t *testing.T) {
extra := map[string]any{}
result := EnsureMessageID(extra)
id := GetMessageID(result)
assert.NotEmpty(t, id)
assert.Len(t, id, 36)
})
t.Run("generates unique IDs", func(t *testing.T) {
extra1 := EnsureMessageID(nil)
extra2 := EnsureMessageID(nil)
assert.NotEqual(t, GetMessageID(extra1), GetMessageID(extra2))
})
}
+412
View File
@@ -0,0 +1,412 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package adk
import (
"bytes"
"context"
"encoding/gob"
"errors"
"fmt"
"sync"
"github.com/cloudwego/eino/internal/core"
"github.com/cloudwego/eino/schema"
)
// ResumeInfo holds all the information necessary to resume an interrupted agent execution.
// It is created by the framework and passed to an agent's Resume method.
type ResumeInfo struct {
// EnableStreaming indicates whether the original execution was in streaming mode.
EnableStreaming bool
// Deprecated: use InterruptContexts from the embedded InterruptInfo for user-facing details,
// and GetInterruptState for internal state retrieval.
*InterruptInfo
WasInterrupted bool
InterruptState any
IsResumeTarget bool
ResumeData any
}
// InterruptInfo contains all the information about an interruption event.
// It is created by the framework when an agent returns an interrupt action.
type InterruptInfo struct {
Data any
// InterruptContexts provides a structured, user-facing view of the interrupt chain.
// Each context represents a step in the agent hierarchy that was interrupted.
InterruptContexts []*InterruptCtx
}
// TypedInterrupt creates a typed interrupt event that pauses execution to request external input.
// It is the generic counterpart of Interrupt; see Interrupt for full documentation.
func TypedInterrupt[M MessageType](ctx context.Context, info any) *TypedAgentEvent[M] {
var rp []RunStep
rCtx := getRunCtx(ctx)
if rCtx != nil {
rp = rCtx.RunPath
}
is, err := core.Interrupt(ctx, info, nil, nil,
core.WithLayerPayload(rp))
if err != nil {
return &TypedAgentEvent[M]{Err: err}
}
contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)
return &TypedAgentEvent[M]{
Action: &AgentAction{
Interrupted: &InterruptInfo{
InterruptContexts: contexts,
},
internalInterrupted: is,
},
}
}
// Interrupt creates a basic interrupt action.
// This is used when an agent needs to pause its execution to request external input or intervention,
// but does not need to save any internal state to be restored upon resumption.
// The `info` parameter is user-facing data that describes the reason for the interrupt.
func Interrupt(ctx context.Context, info any) *AgentEvent {
return TypedInterrupt[*schema.Message](ctx, info)
}
// TypedStatefulInterrupt creates a typed interrupt event that also saves the agent's internal state.
// It is the generic counterpart of StatefulInterrupt; see StatefulInterrupt for full documentation.
func TypedStatefulInterrupt[M MessageType](ctx context.Context, info any, state any) *TypedAgentEvent[M] {
var rp []RunStep
rCtx := getRunCtx(ctx)
if rCtx != nil {
rp = rCtx.RunPath
}
is, err := core.Interrupt(ctx, info, state, nil,
core.WithLayerPayload(rp))
if err != nil {
return &TypedAgentEvent[M]{Err: err}
}
contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)
return &TypedAgentEvent[M]{
Action: &AgentAction{
Interrupted: &InterruptInfo{
InterruptContexts: contexts,
},
internalInterrupted: is,
},
}
}
// StatefulInterrupt creates an interrupt action that also saves the agent's internal state.
// This is used when an agent has internal state that must be restored for it to continue correctly.
// The `info` parameter is user-facing data describing the interrupt.
// The `state` parameter is the agent's internal state object, which will be serialized and stored.
func StatefulInterrupt(ctx context.Context, info any, state any) *AgentEvent {
return TypedStatefulInterrupt[*schema.Message](ctx, info, state)
}
// TypedCompositeInterrupt creates a typed interrupt event that aggregates sub-interrupt signals.
// It is the generic counterpart of CompositeInterrupt; see CompositeInterrupt for full documentation.
func TypedCompositeInterrupt[M MessageType](ctx context.Context, info any, state any,
subInterruptSignals ...*InterruptSignal) *TypedAgentEvent[M] {
var rp []RunStep
rCtx := getRunCtx(ctx)
if rCtx != nil {
rp = rCtx.RunPath
}
is, err := core.Interrupt(ctx, info, state, subInterruptSignals,
core.WithLayerPayload(rp))
if err != nil {
return &TypedAgentEvent[M]{Err: err}
}
contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)
return &TypedAgentEvent[M]{
Action: &AgentAction{
Interrupted: &InterruptInfo{
InterruptContexts: contexts,
},
internalInterrupted: is,
},
}
}
// CompositeInterrupt creates an interrupt event that aggregates sub-interrupt signals.
func CompositeInterrupt(ctx context.Context, info any, state any,
subInterruptSignals ...*InterruptSignal) *AgentEvent {
return TypedCompositeInterrupt[*schema.Message](ctx, info, state, subInterruptSignals...)
}
// Address represents the unique, hierarchical address of a component within an execution.
// It is a slice of AddressSegments, where each segment represents one level of nesting.
// This is a type alias for core.Address. See the core package for more details.
type Address = core.Address
type AddressSegment = core.AddressSegment
type AddressSegmentType = core.AddressSegmentType
const (
AddressSegmentAgent AddressSegmentType = "agent"
AddressSegmentTool AddressSegmentType = "tool"
)
var allowedAddressSegmentTypes = []AddressSegmentType{AddressSegmentAgent, AddressSegmentTool}
// AppendAddressSegment adds an address segment for the current execution context.
func AppendAddressSegment(ctx context.Context, segType AddressSegmentType, segID string) context.Context {
return core.AppendAddressSegment(ctx, segType, segID, "")
}
// InterruptCtx provides a structured, user-facing view of a single point of interruption.
// It contains the ID and Address of the interrupted component, as well as user-defined info.
// This is a type alias for core.InterruptCtx. See the core package for more details.
type InterruptCtx = core.InterruptCtx
type InterruptSignal = core.InterruptSignal
// FromInterruptContexts converts user-facing interrupt contexts to an interrupt signal.
func FromInterruptContexts(contexts []*InterruptCtx) *InterruptSignal {
return core.FromInterruptContexts(contexts)
}
// WithCheckPointID sets the checkpoint ID used for interruption persistence.
func WithCheckPointID(id string) AgentRunOption {
return WrapImplSpecificOptFn(func(t *options) {
t.checkPointID = &id
})
}
func init() {
schema.RegisterName[*serialization]("_eino_adk_serialization")
schema.RegisterName[*WorkflowInterruptInfo]("_eino_adk_workflow_interrupt_info")
// Register []byte for gob: the cancel refactor routes bridge store checkpoint
// bytes ([]byte) through InterruptState.State (type any) inside the outer
// serialization struct. Gob requires concrete types behind interface fields
// to be registered.
gob.Register([]byte{})
}
// serialization CheckpointSchema: root checkpoint payload (gob).
// Any type tagged with `CheckpointSchema:` is persisted and must remain backward compatible.
type serialization struct {
RunCtx *runContext
// deprecated: still keep it here for backward compatibility
Info *InterruptInfo
EnableStreaming bool
InterruptID2Address map[string]Address
InterruptID2State map[string]core.InterruptState
}
func runnerLoadCheckPointImpl(store CheckPointStore, ctx context.Context, checkpointID string) (
context.Context, *runContext, *ResumeInfo, error) {
data, existed, err := store.Get(ctx, checkpointID)
if err != nil {
return nil, nil, nil, fmt.Errorf("failed to get checkpoint from store: %w", err)
}
if !existed {
return nil, nil, nil, fmt.Errorf("checkpoint[%s] not exist", checkpointID)
}
data = preprocessADKCheckpoint(data)
s := &serialization{}
err = gob.NewDecoder(bytes.NewReader(data)).Decode(s)
if err != nil {
return nil, nil, nil, fmt.Errorf("failed to decode checkpoint: %w", err)
}
ctx = core.PopulateInterruptState(ctx, s.InterruptID2Address, s.InterruptID2State)
return ctx, s.RunCtx, &ResumeInfo{
EnableStreaming: s.EnableStreaming,
InterruptInfo: s.Info,
}, nil
}
// preprocessADKCheckpoint fixes a gob incompatibility when resuming old ChatModelAgent/DeepAgents checkpoints.
//
// Background
// - ADK checkpoints are gob-encoded.
// - Some values inside checkpoints are stored as `any`, so gob includes a concrete type name
// string in the wire format and uses that name to pick the local Go type to decode into.
//
// Problem (v0.8.0-v0.8.3 checkpoints)
// - In v0.8.0-v0.8.3, *State was registered under the name "_eino_adk_react_state" AND
// implemented GobEncode/GobDecode, so the wire format for that name is "GobEncoder payload"
// (opaque bytes).
// - In v0.7.*, the same name "_eino_adk_react_state" was used but encoded as a normal struct
// (no GobEncode). Gob treats these two wire formats as incompatible.
// - Gob only allows one local Go type per name. Today we register "_eino_adk_react_state" to
// a v0.7-compatible struct decoder (stateV07). If we try to decode a v0.8.0-v0.8.3
// checkpoint under that same name, gob fails with a "want struct; got non-struct" mismatch.
//
// Solution
// - We keep "_eino_adk_react_state" mapped to the v0.7 decoder.
// - For v0.8.0-v0.8.3 checkpoints only, we rewrite the on-wire name to a same-length alias
// "_eino_adk_state_v080_", which is registered to a GobDecoder-compatible type (stateV080).
// - The alias is the same length as the original, so we can safely replace the length-prefixed
// bytes without re-encoding the whole stream.
func preprocessADKCheckpoint(data []byte) []byte {
const (
lenPrefixedReactStateName = "\x15" + stateGobNameV07
lenPrefixedCompatName = "\x15" + stateGobNameV080
lenPrefixedStateSerializationName = "\x12stateSerialization"
)
// the following line checks whether the checkpoint is persisted through v0.8.0-v0.8.3
if !bytes.Contains(data, []byte(lenPrefixedReactStateName)) || !bytes.Contains(data, []byte(lenPrefixedStateSerializationName)) {
return data
}
return bytes.ReplaceAll(data,
[]byte(lenPrefixedReactStateName),
[]byte(lenPrefixedCompatName))
}
func runnerSaveCheckPointImpl(
enableStreaming bool,
store CheckPointStore,
ctx context.Context,
key string,
info *InterruptInfo,
is *core.InterruptSignal,
) error {
if store == nil {
return nil
}
runCtx := getRunCtx(ctx)
id2Addr, id2State := core.SignalToPersistenceMaps(is)
buf := &bytes.Buffer{}
err := gob.NewEncoder(buf).Encode(&serialization{
RunCtx: runCtx,
Info: info,
InterruptID2Address: id2Addr,
InterruptID2State: id2State,
EnableStreaming: enableStreaming,
})
if err != nil {
return fmt.Errorf("failed to encode checkpoint: %w", err)
}
return store.Set(ctx, key, buf.Bytes())
}
const bridgeCheckpointID = "adk_react_mock_key"
func newBridgeStore() *bridgeStore {
return &bridgeStore{data: make(map[string][]byte)}
}
func newResumeBridgeStore(checkPointID string, data []byte) *bridgeStore {
return &bridgeStore{
data: map[string][]byte{checkPointID: data},
}
}
type bridgeStore struct {
mu sync.Mutex
data map[string][]byte
}
func (m *bridgeStore) Get(_ context.Context, key string) ([]byte, bool, error) {
m.mu.Lock()
defer m.mu.Unlock()
if v, ok := m.data[key]; ok {
return v, true, nil
}
return nil, false, nil
}
func (m *bridgeStore) Set(_ context.Context, key string, checkPoint []byte) error {
m.mu.Lock()
defer m.mu.Unlock()
if m.data == nil {
m.data = make(map[string][]byte)
}
m.data[key] = checkPoint
return nil
}
func getNextResumeAgent(ctx context.Context, _ *ResumeInfo) (string, error) {
nextAgents, err := core.GetNextResumptionPoints(ctx)
if err != nil {
return "", fmt.Errorf("failed to get next agent leading to interruption: %w", err)
}
if len(nextAgents) == 0 {
return "", errors.New("no child agents leading to interrupted agent were found")
}
if len(nextAgents) > 1 {
return "", errors.New("agent has multiple child agents leading to interruption, " +
"but concurrent transfer is not supported")
}
// get the single next agent to delegate to.
var nextAgentID string
for id := range nextAgents {
nextAgentID = id
break
}
return nextAgentID, nil
}
func getNextResumeAgents(ctx context.Context, _ *ResumeInfo) (map[string]bool, error) {
nextAgents, err := core.GetNextResumptionPoints(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get next agents leading to interruption: %w", err)
}
if len(nextAgents) == 0 {
return nil, errors.New("no child agents leading to interrupted agent were found")
}
return nextAgents, nil
}
func buildResumeInfo(ctx context.Context, nextAgentID string, info *ResumeInfo) (
context.Context, *ResumeInfo) {
ctx = AppendAddressSegment(ctx, AddressSegmentAgent, nextAgentID)
nextResumeInfo := &ResumeInfo{
EnableStreaming: info.EnableStreaming,
InterruptInfo: info.InterruptInfo,
}
wasInterrupted, hasState, state := core.GetInterruptState[any](ctx)
nextResumeInfo.WasInterrupted = wasInterrupted
if hasState {
nextResumeInfo.InterruptState = state
}
if wasInterrupted {
isResumeTarget, hasData, data := core.GetResumeContext[any](ctx)
nextResumeInfo.IsResumeTarget = isResumeTarget
if hasData {
nextResumeInfo.ResumeData = data
}
}
ctx = updateRunPathOnly(ctx, nextAgentID)
return ctx, nextResumeInfo
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+213
View File
@@ -0,0 +1,213 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Package agentsmd provides a middleware that automatically injects Agents.md
// file contents into model input messages. The injection is transient — content
// is prepended at model call time and never persisted to conversation state,
// so it is naturally excluded from summarization / compression.
package agentsmd
import (
"context"
"fmt"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
// Config defines the configuration for the agentsmd middleware.
type Config struct {
// Backend provides file access for loading Agents.md files.
// Implementations can use local filesystem, remote storage, or any other backend.
// Required.
Backend Backend
// AgentsMDFiles specifies the ordered list of Agents.md file paths to load.
// Files are loaded and injected in the given order.
// Supports @import syntax inside files for recursive inclusion (max depth 5).
AgentsMDFiles []string
// AllAgentsMDMaxBytes limits the total byte size of all loaded Agents.md content.
// Files are loaded in order; once the cumulative size exceeds this limit,
// remaining files are skipped. Each individual file is always loaded in full.
// 0 means no limit.
AllAgentsMDMaxBytes int
// OnLoadWarning is an optional callback invoked when a non-fatal error occurs
// during Agents.md file loading (e.g. file not found, circular @import, depth
// exceeded). If nil, warnings are logged via log.Printf.
//
// Note: Backend.Read errors other than os.ErrNotExist (e.g. permission denied,
// I/O errors) are NOT treated as warnings and will abort the loading process.
OnLoadWarning func(filePath string, err error)
}
// NewTyped creates a generic agentsmd middleware that injects Agents.md content into every
// model call. The content is loaded from the configured file paths via Backend
// on each model invocation.
//
// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage.
//
// Recommended: place this middleware AFTER the summarization middleware, so that
// Agents.md content is excluded from summarization/compression.
func NewTyped[M adk.MessageType](_ context.Context, cfg *Config) (adk.TypedChatModelAgentMiddleware[M], error) {
if err := cfg.validate(); err != nil {
return nil, err
}
return &typedMiddleware[M]{
loader: newLoaderConfig(cfg.Backend, cfg.AgentsMDFiles, cfg.AllAgentsMDMaxBytes, cfg.OnLoadWarning),
}, nil
}
// New creates an agentsmd middleware that injects Agents.md content into every
// model call. The content is loaded from the configured file paths via Backend
// on each model invocation.
//
// Recommended: place this middleware AFTER the summarization middleware, so that
// Agents.md content is excluded from summarization/compression.
func New(ctx context.Context, cfg *Config) (adk.ChatModelAgentMiddleware, error) {
return NewTyped[*schema.Message](ctx, cfg)
}
type typedMiddleware[M adk.MessageType] struct {
*adk.TypedBaseChatModelAgentMiddleware[M]
loader *loaderConfig
}
const agentsMDCacheKey = "__agentsmd_content_cache__"
const agentsMDExtraKey = "__agentsmd_content__"
// BeforeModelRewriteState injects Agents.md content as a User message before
// the first User message in the conversation. The injected message is tagged
// with an Extra key so that repeated invocations are idempotent.
func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) {
// Idempotent: if we already injected, return early.
for _, msg := range state.Messages {
if hasAgentsMDExtra(msg) {
return ctx, state, nil
}
}
content, err := m.loadContent(ctx)
if err != nil {
return ctx, nil, err
}
if content == "" {
return ctx, state, nil
}
nState := *state
nState.Messages = typedInsertBeforeFirstUser(state.Messages, content)
return ctx, &nState, nil
}
// hasAgentsMDExtra checks whether a message has the agentsmd extra key set.
func hasAgentsMDExtra[M adk.MessageType](msg M) bool {
switch v := any(msg).(type) {
case *schema.Message:
if v.Extra != nil {
if _, ok := v.Extra[agentsMDExtraKey]; ok {
return true
}
}
case *schema.AgenticMessage:
if v.Extra != nil {
if _, ok := v.Extra[agentsMDExtraKey]; ok {
return true
}
}
}
return false
}
// typedInsertBeforeFirstUser inserts a user message with agentsmd content before the first User message.
func typedInsertBeforeFirstUser[M adk.MessageType](msgs []M, content string) []M {
newMsg := makeUserMsgWithExtra[M](content)
result := make([]M, 0, len(msgs)+1)
for i, msg := range msgs {
if isUserRole(msg) {
result = append(result, newMsg)
result = append(result, msgs[i:]...)
return result
}
result = append(result, msg)
}
result = append(result, newMsg)
return result
}
func isUserRole[M adk.MessageType](msg M) bool {
switch m := any(msg).(type) {
case *schema.Message:
return m.Role == schema.User
case *schema.AgenticMessage:
return m.Role == schema.AgenticRoleTypeUser
}
return false
}
func makeUserMsgWithExtra[M adk.MessageType](content string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
msg := schema.UserMessage(content)
msg.Extra = map[string]any{agentsMDExtraKey: true}
return any(msg).(M)
case *schema.AgenticMessage:
msg := schema.UserAgenticMessage(content)
msg.Extra = map[string]any{agentsMDExtraKey: true}
return any(msg).(M)
}
panic("unreachable")
}
// loadContent retrieves the Agents.md content, using a per-Run cache to avoid
// reloading on every model call within the same Run().
func (m *typedMiddleware[M]) loadContent(ctx context.Context) (string, error) {
if cached, found, err := adk.GetRunLocalValue(ctx, agentsMDCacheKey); err == nil && found {
if s, ok := cached.(string); ok {
return s, nil
}
}
content, err := m.loader.load(ctx)
if err != nil {
return "", fmt.Errorf("[agentsmd]: failed to load agent files: %w", err)
}
if content != "" {
_ = adk.SetRunLocalValue(ctx, agentsMDCacheKey, content)
}
return content, nil
}
func (c *Config) validate() error {
if c == nil {
return fmt.Errorf("[agentsmd]: config is required")
}
if c.Backend == nil {
return fmt.Errorf("[agentsmd]: backend is required")
}
if len(c.AgentsMDFiles) == 0 {
return fmt.Errorf("[agentsmd]: at least one agent file path is required")
}
if c.AllAgentsMDMaxBytes < 0 {
return fmt.Errorf("[agentsmd]: AllAgentMDDocsMaxBytes must be non-negative")
}
return nil
}
@@ -0,0 +1,351 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package agentsmd
import (
"context"
"strings"
"testing"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
// --- generic table-driven test helpers ---
func makeUserMsg[M adk.MessageType](content string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any(&schema.Message{Role: schema.User, Content: content}).(M)
case *schema.AgenticMessage:
return any(schema.UserAgenticMessage(content)).(M)
default:
panic("unreachable")
}
}
func makeSystemMsg[M adk.MessageType](content string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any(&schema.Message{Role: schema.System, Content: content}).(M)
case *schema.AgenticMessage:
return any(schema.SystemAgenticMessage(content)).(M)
default:
panic("unreachable")
}
}
func makeAssistantMsg[M adk.MessageType](content string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any(&schema.Message{Role: schema.Assistant, Content: content}).(M)
case *schema.AgenticMessage:
return any(&schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: content})},
}).(M)
default:
panic("unreachable")
}
}
func getMsgRole[M adk.MessageType](msg M) string {
switch v := any(msg).(type) {
case *schema.Message:
return string(v.Role)
case *schema.AgenticMessage:
return string(v.Role)
default:
panic("unreachable")
}
}
func getMsgContent[M adk.MessageType](msg M) string {
switch v := any(msg).(type) {
case *schema.Message:
return v.Content
case *schema.AgenticMessage:
for _, block := range v.ContentBlocks {
if block == nil {
continue
}
if block.UserInputText != nil {
return block.UserInputText.Text
}
if block.AssistantGenText != nil {
return block.AssistantGenText.Text
}
}
return ""
default:
panic("unreachable")
}
}
func getMsgExtra[M adk.MessageType](msg M) map[string]any {
switch v := any(msg).(type) {
case *schema.Message:
return v.Extra
case *schema.AgenticMessage:
return v.Extra
default:
panic("unreachable")
}
}
// --- generic table-driven test ---
type agentsMDTestCase struct {
name string
run func(t *testing.T)
}
func testAgentsMDGeneric[M adk.MessageType](t *testing.T) {
tests := []agentsMDTestCase{
{
name: "BasicInjection",
run: func(t *testing.T) {
b := newMemBackend()
b.set("/agent.md", "You are a helpful assistant.")
ctx := context.Background()
mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}})
if err != nil {
t.Fatal(err)
}
state := &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}}
_, state, err = mw.BeforeModelRewriteState(ctx, state, nil)
if err != nil {
t.Fatal(err)
}
if len(state.Messages) != 2 {
t.Fatalf("expected 2 messages, got %d", len(state.Messages))
}
if getMsgRole(state.Messages[0]) != "user" {
t.Fatalf("expected first message role user, got %s", getMsgRole(state.Messages[0]))
}
if !strings.Contains(getMsgContent(state.Messages[0]), "You are a helpful assistant.") {
t.Fatalf("expected agent.md content in first message, got %q", getMsgContent(state.Messages[0]))
}
if !strings.Contains(getMsgContent(state.Messages[0]), "<system-reminder>") {
t.Fatalf("expected system-reminder tag, got %q", getMsgContent(state.Messages[0]))
}
if count := strings.Count(getMsgContent(state.Messages[0]), "<system-reminder>"); count != 1 {
t.Fatalf("expected exactly one opening system-reminder tag, got %d in %q", count, getMsgContent(state.Messages[0]))
}
if count := strings.Count(getMsgContent(state.Messages[0]), "</system-reminder>"); count != 1 {
t.Fatalf("expected exactly one closing system-reminder tag, got %d in %q", count, getMsgContent(state.Messages[0]))
}
if getMsgContent(state.Messages[1]) != "hello" {
t.Fatalf("expected original message preserved, got %q", getMsgContent(state.Messages[1]))
}
},
},
{
name: "InsertBeforeFirstUserMessage",
run: func(t *testing.T) {
b := newMemBackend()
b.set("/agent.md", "agent instructions")
ctx := context.Background()
mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}})
if err != nil {
t.Fatal(err)
}
input := []M{
makeSystemMsg[M]("system prompt"),
makeUserMsg[M]("hello"),
}
state := &adk.TypedChatModelAgentState[M]{Messages: input}
_, state, err = mw.BeforeModelRewriteState(ctx, state, nil)
if err != nil {
t.Fatal(err)
}
if len(state.Messages) != 3 {
t.Fatalf("expected 3 messages, got %d", len(state.Messages))
}
if getMsgRole(state.Messages[0]) != "system" {
t.Fatalf("expected first message role system, got %s", getMsgRole(state.Messages[0]))
}
if getMsgContent(state.Messages[0]) != "system prompt" {
t.Fatalf("expected system prompt preserved, got %q", getMsgContent(state.Messages[0]))
}
if getMsgRole(state.Messages[1]) != "user" || !strings.Contains(getMsgContent(state.Messages[1]), "agent instructions") {
t.Fatalf("expected agentmd message at index 1, got role=%s content=%q", getMsgRole(state.Messages[1]), getMsgContent(state.Messages[1]))
}
if getMsgRole(state.Messages[2]) != "user" || getMsgContent(state.Messages[2]) != "hello" {
t.Fatalf("expected original user message at index 2, got role=%s content=%q", getMsgRole(state.Messages[2]), getMsgContent(state.Messages[2]))
}
},
},
{
name: "InsertWithNoUserMessage",
run: func(t *testing.T) {
b := newMemBackend()
b.set("/agent.md", "agent instructions")
ctx := context.Background()
mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}})
if err != nil {
t.Fatal(err)
}
input := []M{
makeSystemMsg[M]("system prompt"),
makeAssistantMsg[M]("assistant reply"),
}
state := &adk.TypedChatModelAgentState[M]{Messages: input}
_, state, err = mw.BeforeModelRewriteState(ctx, state, nil)
if err != nil {
t.Fatal(err)
}
if len(state.Messages) != 3 {
t.Fatalf("expected 3 messages, got %d", len(state.Messages))
}
if getMsgRole(state.Messages[0]) != "system" {
t.Fatalf("expected System at index 0, got %s", getMsgRole(state.Messages[0]))
}
if getMsgRole(state.Messages[1]) != "assistant" {
t.Fatalf("expected Assistant at index 1, got %s", getMsgRole(state.Messages[1]))
}
if getMsgRole(state.Messages[2]) != "user" || !strings.Contains(getMsgContent(state.Messages[2]), "agent instructions") {
t.Fatalf("expected agentmd appended at end, got role=%s content=%q", getMsgRole(state.Messages[2]), getMsgContent(state.Messages[2]))
}
},
},
{
name: "AllFilesEmpty",
run: func(t *testing.T) {
b := newMemBackend()
b.set("/agent.md", "")
ctx := context.Background()
mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}})
if err != nil {
t.Fatal(err)
}
state := &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}}
_, state, err = mw.BeforeModelRewriteState(ctx, state, nil)
if err != nil {
t.Fatal(err)
}
if len(state.Messages) != 1 {
t.Fatalf("expected 1 message (no agentmd prepended), got %d", len(state.Messages))
}
if getMsgContent(state.Messages[0]) != "hello" {
t.Fatalf("expected original message unchanged, got %q", getMsgContent(state.Messages[0]))
}
},
},
{
name: "Idempotency",
run: func(t *testing.T) {
b := newMemBackend()
b.set("/agent.md", "agent instructions")
ctx := context.Background()
mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}})
if err != nil {
t.Fatal(err)
}
state := &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}}
_, state, err = mw.BeforeModelRewriteState(ctx, state, nil)
if err != nil {
t.Fatal(err)
}
if len(state.Messages) != 2 {
t.Fatalf("expected 2 messages after first call, got %d", len(state.Messages))
}
// Verify the marker is set in Extra.
extra := getMsgExtra(state.Messages[0])
if extra == nil {
t.Fatal("expected Extra to be set on injected message")
}
if _, ok := extra[agentsMDExtraKey]; !ok {
t.Fatalf("expected agentsMDExtraKey in Extra, got %v", extra)
}
// Call again with the same state (which now contains the marker message).
_, state, err = mw.BeforeModelRewriteState(ctx, state, nil)
if err != nil {
t.Fatal(err)
}
if len(state.Messages) != 2 {
t.Fatalf("expected 2 messages after second call (idempotent), got %d", len(state.Messages))
}
if !strings.Contains(getMsgContent(state.Messages[0]), "agent instructions") {
t.Fatalf("expected agentmd content preserved, got %q", getMsgContent(state.Messages[0]))
}
},
},
{
name: "ReinsertAfterRemoval",
run: func(t *testing.T) {
b := newMemBackend()
b.set("/agent.md", "agent instructions")
ctx := context.Background()
mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}})
if err != nil {
t.Fatal(err)
}
state := &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}}
_, state, err = mw.BeforeModelRewriteState(ctx, state, nil)
if err != nil {
t.Fatal(err)
}
if len(state.Messages) != 2 {
t.Fatalf("expected 2 messages after first call, got %d", len(state.Messages))
}
// Simulate removal of the marker message (e.g., by summarization).
state = &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}}
_, state, err = mw.BeforeModelRewriteState(ctx, state, nil)
if err != nil {
t.Fatal(err)
}
if len(state.Messages) != 2 {
t.Fatalf("expected 2 messages after re-insert, got %d", len(state.Messages))
}
if !strings.Contains(getMsgContent(state.Messages[0]), "agent instructions") {
t.Fatalf("expected agentmd content re-inserted, got %q", getMsgContent(state.Messages[0]))
}
},
},
}
for _, tc := range tests {
t.Run(tc.name, tc.run)
}
}
func TestAgentsMDGeneric(t *testing.T) {
t.Run("Message", testAgentsMDGeneric[*schema.Message])
t.Run("AgenticMessage", testAgentsMDGeneric[*schema.AgenticMessage])
}
File diff suppressed because it is too large Load Diff
+299
View File
@@ -0,0 +1,299 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package agentsmd
import (
"context"
"errors"
"fmt"
"log"
"os"
"path/filepath"
"regexp"
"strings"
"github.com/cloudwego/eino/adk/filesystem"
"github.com/cloudwego/eino/adk/internal"
)
// importRegex matches @path/to/file anywhere in text.
// The path must start with a letter, digit, dot, underscore, slash, or tilde, followed by
// path characters (letters, digits, dots, slashes, hyphens, underscores).
// A post-match filter further requires the path to contain "/" or end with
// an allowed extension (see allowedImportExts), so bare words like @someone
// and email-like patterns like @example.com are ignored.
var importRegex = regexp.MustCompile(`@([a-zA-Z0-9_.~/][a-zA-Z0-9_.~/\-]*)`)
// allowedImportExts is the set of file extensions recognised as @import targets.
// Paths without "/" must end with one of these extensions to be treated as imports;
// this avoids false positives on email addresses (@example.com) and mentions (@foo.bar).
var allowedImportExts = map[string]bool{
".md": true,
".txt": true,
".mdx": true,
".yaml": true,
".yml": true,
".json": true,
".toml": true,
}
const maxImportDepth = 5
// ReadRequest is an alias for filesystem.ReadRequest.
type ReadRequest = filesystem.ReadRequest
type FileContent = filesystem.FileContent
// Backend defines the file access interface for loading Agents.md files.
// Implementations can use local filesystem, remote storage, or any other backend.
type Backend interface {
// Read reads the content of a file.
// If the file does not exist, implementations should return an error wrapping
// os.ErrNotExist (so that errors.Is(err, os.ErrNotExist) returns true). This allows the loader
// to silently skip missing files and notify via OnLoadWarning callback.
// Other errors (e.g. permission denied, I/O errors) will abort the loading process.
Read(ctx context.Context, req *ReadRequest) (*FileContent, error)
}
// loaderConfig holds the immutable configuration for creating loaders.
// It is safe for concurrent use by multiple goroutines.
type loaderConfig struct {
backend Backend
files []string // ordered file paths from config
maxBytes int // cumulative read budget; 0 means unlimited
onWarning func(filePath string, err error) // callback for non-fatal loading warnings
}
func newLoaderConfig(backend Backend, files []string, maxBytes int, onWarning func(filePath string, err error)) *loaderConfig {
if onWarning == nil {
onWarning = func(filePath string, err error) {
log.Printf("[agentsmd] warning: %s: %v", filePath, err)
}
}
return &loaderConfig{
backend: backend,
files: files,
maxBytes: maxBytes,
onWarning: onWarning,
}
}
// loader handles loading and @import resolution for agents.md files.
// A new loader is created for each load() call to avoid sharing mutable state
// (totalBytes) across concurrent invocations.
type loader struct {
*loaderConfig
totalBytes int // accumulated bytes during this load call
}
func (cfg *loaderConfig) newLoader() *loader {
return &loader{loaderConfig: cfg}
}
// load reads all agents.md files and returns the formatted content.
// Each top-level file and its @imported files appear as separate sections.
func (cfg *loaderConfig) load(ctx context.Context) (string, error) {
l := cfg.newLoader()
var parts []loadedFile
seen := make(map[string]bool) // dedup across all files and imports
for i, filePath := range l.files {
files, err := l.loadFile(ctx, filePath, 0, make(map[string]bool), seen)
if err != nil {
return "", fmt.Errorf("failed to load %q: %w", filePath, err)
}
// If loading this file caused the budget to be exceeded, skip it
// (but always include the first file).
if i > 0 && l.maxBytes > 0 && l.totalBytes > l.maxBytes {
l.onWarning(filePath, fmt.Errorf("skipped: cumulative size %d exceeds max bytes %d", l.totalBytes, l.maxBytes))
break
}
parts = append(parts, files...)
}
return formatContent(parts), nil
}
// loadFile reads a file via Backend and collects @imported files as separate entries.
// Returns a slice where the first element is this file itself, followed by all
// transitively imported files (in encounter order, preserving @path in original text).
// visited tracks the current ancestor chain to detect circular imports.
// seen tracks globally loaded files to avoid duplicate reads and byte counting.
func (l *loader) loadFile(ctx context.Context, filePath string, depth int, visited map[string]bool, seen map[string]bool) ([]loadedFile, error) {
filePath = filepath.Clean(filePath)
if depth > maxImportDepth {
l.onWarning(filePath, fmt.Errorf("@import depth exceeds maximum of %d", maxImportDepth))
return nil, nil
}
if visited[filePath] {
l.onWarning(filePath, fmt.Errorf("circular @import detected"))
return nil, nil
}
if seen[filePath] {
return nil, nil
}
visited[filePath] = true
defer delete(visited, filePath)
fileContent, err := l.backend.Read(ctx, &ReadRequest{FilePath: filePath, Offset: 1})
if err != nil {
if errors.Is(err, os.ErrNotExist) {
l.onWarning(filePath, fmt.Errorf("file not found, skipping"))
return nil, nil
}
return nil, err
}
content := ""
if fileContent != nil {
content = fileContent.Content
}
l.totalBytes += len(content)
seen[filePath] = true
if content == "" {
return nil, nil
}
// Collect imported files as separate sections (content stays untouched).
imports, err := l.collectImports(ctx, filePath, content, depth, visited, seen)
if err != nil {
return nil, err
}
// This file first, then its imports.
result := make([]loadedFile, 0, 1+len(imports))
result = append(result, loadedFile{path: filePath, content: content})
result = append(result, imports...)
return result, nil
}
// collectImports scans content for @path/to/file references and loads each
// imported file (plus its transitive imports). The original content is NOT modified.
// Returns the list of imported loadedFile entries in encounter order.
// seen is shared across the entire load call to avoid duplicate reads.
// Non-fatal errors (file not found, depth exceeded, circular import) are reported
// via onWarning and skipped. Fatal errors (e.g. I/O) are returned.
func (l *loader) collectImports(ctx context.Context, hostPath, content string, depth int, visited map[string]bool, seen map[string]bool) ([]loadedFile, error) {
dir := filepath.Dir(hostPath)
var imports []loadedFile
matches := importRegex.FindAllStringSubmatch(content, -1)
for _, match := range matches {
rawPath := match[1]
// Only treat as import if path contains "/" or ends with an allowed extension.
// This avoids false positives on email addresses and social mentions.
if !strings.Contains(rawPath, "/") && !allowedImportExts[filepath.Ext(rawPath)] {
continue
}
// If budget is exhausted, skip further imports.
if l.maxBytes > 0 && l.totalBytes > l.maxBytes {
break
}
importPath := rawPath
if !filepath.IsAbs(importPath) {
importPath = filepath.Join(dir, importPath)
}
if seen[importPath] {
continue
}
files, err := l.loadFile(ctx, importPath, depth+1, visited, seen)
if err != nil {
return nil, fmt.Errorf("failed to import %q from %q: %w", rawPath, hostPath, err)
}
imports = append(imports, files...)
}
return imports, nil
}
type loadedFile struct {
path string
content string
}
const formatHeaderEn = `<system-reminder>
As you answer the user's questions, you can use the following context:
Codebase and user instructions are shown below. Be sure to adhere to these instructions. IMPORTANT: These instructions OVERRIDE any default behavior and you MUST follow them exactly as written.
`
const formatHeaderCn = `<system-reminder>
在回答用户问题时,你可以使用以下上下文:
代码库和用户指令如下。请务必遵守这些指令。重要提示:这些指令会覆盖任何默认行为,你必须严格按照要求执行。
`
const formatFileHeaderEn = "\nContents of "
const formatFileHeaderCn = "\n文件内容:"
const formatFileLabelEn = " (instructions):\n\n"
const formatFileLabelCn = "(指令):\n\n"
const formatFooterEn = `IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.
</system-reminder>`
const formatFooterCn = `重要提示:此上下文可能与你的任务相关,也可能不相关。除非此上下文与你的任务高度相关,否则不要响应此上下文。
</system-reminder>`
func formatContent(files []loadedFile) string {
if len(files) == 0 {
return ""
}
header := internal.SelectPrompt(internal.I18nPrompts{
English: formatHeaderEn,
Chinese: formatHeaderCn,
})
fileHeader := internal.SelectPrompt(internal.I18nPrompts{
English: formatFileHeaderEn,
Chinese: formatFileHeaderCn,
})
fileLabel := internal.SelectPrompt(internal.I18nPrompts{
English: formatFileLabelEn,
Chinese: formatFileLabelCn,
})
footer := internal.SelectPrompt(internal.I18nPrompts{
English: formatFooterEn,
Chinese: formatFooterCn,
})
var sb strings.Builder
sb.WriteString(header)
for _, f := range files {
sb.WriteString(fileHeader)
sb.WriteString(f.path)
sb.WriteString(fileLabel)
sb.WriteString(f.content)
sb.WriteString("\n")
}
sb.WriteString(footer)
return sb.String()
}
@@ -0,0 +1,162 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package toolsearch
const (
toolDescription = `Search for or select deferred tools to make them available for use.
MANDATORY PREREQUISITE - THIS IS A HARD REQUIREMENT
You MUST use this tool to load deferred tools BEFORE calling them directly.
This is a BLOCKING REQUIREMENT - deferred tools are NOT available until you load them using this tool. Look for <available-deferred-tools> messages in the conversation for the list of tools you can discover. Both query modes (keyword search and direct selection) load the returned tools — once a tool appears in the results, it is immediately available to call.
Why this is non-negotiable:
- Deferred tools are not loaded until discovered via this tool
- Calling a deferred tool without first loading it will fail
Query modes:
1. Keyword search - Use keywords when you're unsure which tool to use or need to discover multiple tools at once:
- "list directory" - find tools for listing directories
- "notebook jupyter" - find notebook editing tools
- "slack message" - find slack messaging tools
- Returns up to 5 matching tools ranked by relevance
- All returned tools are immediately available to call — no further selection step needed
2. Direct selection - Use select:<tool_name> when you know the exact tool name:
- "select:mcp__slack__read_channel"
- "select:NotebookEdit"
- "select:Read,Edit,Grep" - load multiple tools at once with comma separation
- Returns the named tool(s) if they exist
IMPORTANT: Both modes load tools equally. Do NOT follow up a keyword search with select: calls for tools already returned — they are already loaded.
3. Required keyword - Prefix with + to require a match:
- "+linear create issue" - only tools from "linear", ranked by "create"/"issue"
- "+slack send" - only "slack" tools, ranked by "send"
- Useful when you know the service name but not the exact tool
CORRECT Usage Patterns:
<example>
User: I need to work with slack somehow
Assistant: Let me search for slack tools.
[Calls tool_search with query: "slack"]
Assistant: Found several options including mcp__slack__read_channel.
[Calls mcp__slack__read_channel directly — it was loaded by the keyword search]
</example>
<example>
User: Edit the Jupyter notebook
Assistant: Let me load the notebook editing tool.
[Calls tool_search with query: "select:NotebookEdit"]
[Calls NotebookEdit]
</example>
<example>
User: List files in the src directory
Assistant: I can see mcp__filesystem__list_directory in the available tools. Let me select it.
[Calls tool_search with query: "select:mcp__filesystem__list_directory"]
[Calls the tool]
</example>
INCORRECT Usage Patterns - NEVER DO THESE:
<bad-example>
User: Read my slack messages
Assistant: [Directly calls mcp__slack__read_channel without loading it first]
WRONG - You must load the tool FIRST using this tool
</bad-example>
<bad-example>
Assistant: [Calls tool_search with query: "slack", gets back mcp__slack__read_channel]
Assistant: [Calls tool_search with query: "select:mcp__slack__read_channel"]
WRONG - The keyword search already loaded the tool. The select call is redundant.
</bad-example>`
toolDescriptionChinese = `搜索或选择延迟加载(deferred)的工具,使其可供调用。
强制前提条件(MANDATORY PREREQUISITE)— 硬性要求
在直接调用任何 延迟加载工具(deferred tools) 之前,你 必须先使用此工具将其加载。
这是一个 阻塞性要求(BLOCKING REQUIREMENT) — 延迟加载工具在被加载之前是 不可用的。你需要在对话中查找 <available-deferred-tools> 消息,以获取可以发现的工具列表。无论使用哪种查询方式(关键字搜索 或 直接选择),只要工具出现在返回结果中,它们就会自动被加载并立即可调用。
为什么这是不可协商的规则:
- 延迟加载工具在被发现之前不会被加载
- 如果你在加载之前直接调用延迟工具,调用将会失败
查询模式:
1. 关键字搜索(Keyword search)- 当你不确定具体需要哪个工具,或希望一次发现多个工具时使用关键字搜索:
- "list directory" — 查找用于列出目录的工具
- "notebook jupyter" — 查找 Jupyter Notebook 编辑工具
- "slack message" — 查找 Slack 消息相关工具
- 返回最多 5 个最相关的工具
- 所有返回的工具都会立即加载并可直接调用 — 不需要额外执行 select 步骤
2. 直接选择(Direct selection)— 当你已经知道工具的确切名称时使用 select:<tool_name>
- "select:mcp__slack__read_channel"
- "select:NotebookEdit"
- "select:Read,Edit,Grep" — 一次加载多个工具
- 如果工具存在,将被加载并返回
重要说明:两种模式的加载效果完全相同。不要在关键词搜索之后,对返回的工具再次进行 select: 选择 — 它们已经加载好了。
3. 必须匹配关键字(Required keyword)— 在关键字前添加 + 可以 强制匹配特定服务或来源。
- "+linear create issue" — 仅返回名字中包含 "linear" 的工具,按 "create" / "issue" 排序
- "+slack send" — 仅返回名字中包含 "slack" 的工具,按 "send" 排序
- 适用于你知道服务名称但不知道具体工具名称
正确使用示例:
<example>
User: 我需要处理 Slack 相关的事情
Assistant: 让我搜索 Slack 工具。
[调用 tool_searchquery: "slack"]
Assistant: 找到多个选项,包括 mcp__slack__read_channel。
[直接调用 mcp__slack__read_channel — 关键字搜索已经加载了该工具]
</example>
<example>
User: 编辑这个 Jupyter Notebook
Assistant: 让我加载 Notebook 编辑工具。
[调用 tool_searchquery: "select:NotebookEdit"]
[调用 NotebookEdit]
</example>
<example>
User: 列出 src 目录中的文件
Assistant: 我看到可用工具中有 mcp__filesystem__list_directory,让我加载它。
[调用 tool_searchquery: "select:mcp__filesystem__list_directory"]
[调用该工具]
</example>
错误用法(严禁)
<bad-example>
User: 读取我的 Slack 消息
Assistant: [不调用 tool_search 工具加载,直接调用 mcp__slack__read_channel]
错误 — 在调用工具之前没有先使用 tool_search 加载该工具。
</bad-example>
<bad-example>
Assistant:[调用 tool_searchquery: "slack",返回 mcp__slack__read_channel]
Assistant:[再次调用 tool_searchquery: "select:mcp__slack__read_channel"]
错误 — 关键字搜索 已经加载了该工具,再次 select 是冗余操作。`
systemReminderTpl = `<available-deferred-tools>
{{- range .Tools }}
{{ . }}
{{- end }}
</available-deferred-tools>`
)
@@ -0,0 +1,647 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Package toolsearch provides tool search middleware.
package toolsearch
import (
"bytes"
"context"
"encoding/json"
"fmt"
"sort"
"strings"
"text/template"
"unicode"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/adk/internal"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema"
)
// Config is the configuration for the tool search middleware.
type Config struct {
// DynamicTools is a list of tools that can be dynamically searched and loaded by the agent.
DynamicTools []tool.BaseTool
// UseModelToolSearch indicates whether the ChatModel natively supports tool search.
//
// When true, the middleware delegates tool search to the model's native capability.
//
// When false (default), the middleware manages tool visibility by filtering the tool list
// based on tool_search results before each model call. Note that this approach may
// invalidate the model's KV-cache (as the tool list changes between calls), and effectiveness
// depends on the model's ability to work with a dynamically changing tool set.
UseModelToolSearch bool
}
// NewTyped constructs and returns the generic tool search middleware.
//
// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage.
func NewTyped[M adk.MessageType](ctx context.Context, config *Config) (adk.TypedChatModelAgentMiddleware[M], error) {
if config == nil {
return nil, fmt.Errorf("config is required")
}
if len(config.DynamicTools) == 0 {
return nil, fmt.Errorf("tools is required")
}
tpl, err := template.New("").Parse(systemReminderTpl)
if err != nil {
return nil, err
}
dynamicToolInfos := make([]*schema.ToolInfo, 0, len(config.DynamicTools))
mapOfDynamicTools := make(map[string]*schema.ToolInfo, len(config.DynamicTools))
toolNames := make([]string, 0, len(config.DynamicTools))
for _, t := range config.DynamicTools {
info, infoErr := t.Info(ctx)
if infoErr != nil {
return nil, fmt.Errorf("failed to get dynamic tool info: %w", infoErr)
}
if _, ok := mapOfDynamicTools[info.Name]; ok {
return nil, fmt.Errorf("duplicate dynamic tool name: %s", info.Name)
}
toolNames = append(toolNames, info.Name)
mapOfDynamicTools[info.Name] = info
dynamicToolInfos = append(dynamicToolInfos, info)
}
buf := &bytes.Buffer{}
err = tpl.Execute(buf, systemReminder{Tools: toolNames})
if err != nil {
return nil, fmt.Errorf("failed to format system reminder template: %w", err)
}
return &typedMiddleware[M]{
dynamicTools: config.DynamicTools,
mapOfDynamicTools: mapOfDynamicTools,
dynamicToolInfos: dynamicToolInfos,
useModelToolSearch: config.UseModelToolSearch,
sr: buf.String(),
}, nil
}
// New constructs and returns the tool search middleware.
//
// The tool search middleware enables dynamic tool selection for agents with large tool libraries.
// Instead of passing all tools to the model at once (which can overwhelm context limits),
// this middleware:
//
// 1. Adds a "tool_search" meta-tool that accepts keyword queries to search tools
// 2. Initially hides all dynamic tools from the model's tool list
// 3. When the model calls tool_search, matching tools become available for subsequent calls
//
// Example usage:
//
// middleware, _ := toolsearch.New(ctx, &toolsearch.Config{
// DynamicTools: []tool.BaseTool{weatherTool, stockTool, currencyTool, ...},
// })
// agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
// // ...
// Handlers: []adk.ChatModelAgentMiddleware{middleware},
// })
func New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) {
return NewTyped[*schema.Message](ctx, config)
}
type systemReminder struct {
Tools []string
}
type typedMiddleware[M adk.MessageType] struct {
*adk.TypedBaseChatModelAgentMiddleware[M]
dynamicTools []tool.BaseTool
mapOfDynamicTools map[string]*schema.ToolInfo
dynamicToolInfos []*schema.ToolInfo
useModelToolSearch bool
sr string
}
func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) {
if runCtx == nil {
return ctx, runCtx, nil
}
nRunCtx := *runCtx
nRunCtx.Tools = make([]tool.BaseTool, len(runCtx.Tools), len(runCtx.Tools)+1+len(m.dynamicTools))
copy(nRunCtx.Tools, runCtx.Tools)
nRunCtx.Tools = append(nRunCtx.Tools, newToolSearchTool(m.mapOfDynamicTools, m.useModelToolSearch))
nRunCtx.Tools = append(nRunCtx.Tools, m.dynamicTools...)
if m.useModelToolSearch {
nRunCtx.ToolSearchTool = getToolSearchToolInfo()
}
return ctx, &nRunCtx, nil
}
const toolSearchInitializedKey = "__toolsearch_initialized__"
const toolSearchReminderExtraKey = "__toolsearch_reminder__"
func (m *typedMiddleware[M]) isInitialized(ctx context.Context) bool {
val, ok, err := adk.GetRunLocalValue(ctx, toolSearchInitializedKey)
if err != nil || !ok {
return false
}
b, _ := val.(bool)
return b
}
func (m *typedMiddleware[M]) markInitialized(ctx context.Context) {
_ = adk.SetRunLocalValue(ctx, toolSearchInitializedKey, true)
}
func (m *typedMiddleware[M]) ensureReminder(msgs []M) []M {
for _, msg := range msgs {
if hasToolSearchReminderExtra(msg) {
return msgs
}
}
reminder := makeReminderMsg[M](m.sr)
result := make([]M, 0, len(msgs)+1)
inserted := false
for _, msg := range msgs {
if !inserted && !isSystemRoleTS(msg) {
inserted = true
result = append(result, reminder)
}
result = append(result, msg)
}
if !inserted {
result = append(result, reminder)
}
return result
}
func isSystemRoleTS[M adk.MessageType](msg M) bool {
switch m := any(msg).(type) {
case *schema.Message:
return m.Role == schema.System
case *schema.AgenticMessage:
return m.Role == schema.AgenticRoleTypeSystem
}
return false
}
func makeReminderMsg[M adk.MessageType](content string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
msg := schema.UserMessage(content)
msg.Extra = map[string]any{toolSearchReminderExtraKey: true}
return any(msg).(M)
case *schema.AgenticMessage:
msg := schema.UserAgenticMessage(content)
msg.Extra = map[string]any{toolSearchReminderExtraKey: true}
return any(msg).(M)
}
panic("unreachable")
}
func hasToolSearchReminderExtra[M adk.MessageType](msg M) bool {
switch v := any(msg).(type) {
case *schema.Message:
if v.Extra != nil {
if b, ok := v.Extra[toolSearchReminderExtraKey]; ok {
if bVal, _ := b.(bool); bVal {
return true
}
}
}
case *schema.AgenticMessage:
if v.Extra != nil {
if b, ok := v.Extra[toolSearchReminderExtraKey]; ok {
if bVal, _ := b.(bool); bVal {
return true
}
}
}
}
return false
}
func (m *typedMiddleware[M]) extractDynamicTools(tools []*schema.ToolInfo) []*schema.ToolInfo {
var result []*schema.ToolInfo
for _, t := range tools {
if _, ok := m.mapOfDynamicTools[t.Name]; ok {
result = append(result, t)
}
}
return result
}
func (m *typedMiddleware[M]) stripDynamicTools(tools []*schema.ToolInfo) []*schema.ToolInfo {
var result []*schema.ToolInfo
for _, t := range tools {
if _, ok := m.mapOfDynamicTools[t.Name]; !ok {
result = append(result, t)
}
}
return result
}
func removeTool(tools []*schema.ToolInfo, name string) []*schema.ToolInfo {
var result []*schema.ToolInfo
for _, t := range tools {
if t.Name != name {
result = append(result, t)
}
}
return result
}
func toolNameSet(tools []*schema.ToolInfo) map[string]bool {
m := make(map[string]bool, len(tools))
for _, t := range tools {
m[t.Name] = true
}
return m
}
func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) {
state.Messages = m.ensureReminder(state.Messages)
if !m.isInitialized(ctx) {
m.markInitialized(ctx)
if m.useModelToolSearch {
// Model-native search: move dynamic tools to DeferredToolInfos for server-side retrieval,
// keep only static tools in ToolInfos, and remove the tool_search tool (the model handles search itself).
state.DeferredToolInfos = m.extractDynamicTools(state.ToolInfos)
state.ToolInfos = m.stripDynamicTools(state.ToolInfos)
state.ToolInfos = removeTool(state.ToolInfos, toolSearchToolName)
} else {
// Client-side search: hide dynamic tools initially; they become visible
// only after the model calls tool_search and forward selection adds them back.
state.ToolInfos = m.stripDynamicTools(state.ToolInfos)
}
}
// Forward selection (client-side search only): scan tool_search results in the
// conversation history and add the selected dynamic tools back to ToolInfos.
if !m.useModelToolSearch {
existing := toolNameSet(state.ToolInfos)
for _, msg := range state.Messages {
content, ok := extractToolSearchResult(msg, toolSearchToolName)
if !ok {
continue
}
var result toolSearchResult
if err := json.Unmarshal([]byte(content), &result); err != nil {
continue
}
for _, name := range result.Matches {
if existing[name] {
continue
}
if info, ok := m.mapOfDynamicTools[name]; ok {
state.ToolInfos = append(state.ToolInfos, info)
existing[name] = true
}
}
}
}
return ctx, state, nil
}
// extractToolSearchResult checks if the given message is a tool result from the tool_search tool,
// and if so returns the content string. Returns ("", false) if not a matching tool result.
func extractToolSearchResult[M adk.MessageType](msg M, toolName string) (string, bool) {
switch v := any(msg).(type) {
case *schema.Message:
if v.Role == schema.Tool && v.ToolName == toolName {
return v.Content, true
}
case *schema.AgenticMessage:
for _, block := range v.ContentBlocks {
if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult &&
block.FunctionToolResult != nil && block.FunctionToolResult.Name == toolName {
for _, b := range block.FunctionToolResult.Content {
if b != nil && b.Text != nil {
return b.Text.Text, true
}
}
}
}
}
return "", false
}
func newToolSearchTool(tools map[string]*schema.ToolInfo, useModelToolSearch bool) tool.BaseTool {
if useModelToolSearch {
return &modelToolSearchTool{tools: tools}
}
return &toolSearchTool{tools: tools}
}
type toolSearchArgs struct {
Query string `json:"query"`
MaxResults *int `json:"max_results,omitempty"`
}
type toolSearchResult struct {
Matches []string `json:"matches"`
}
type toolSearchTool struct {
tools map[string]*schema.ToolInfo
}
func (t *toolSearchTool) Info(_ context.Context) (*schema.ToolInfo, error) {
return getToolSearchToolInfo(), nil
}
func (t *toolSearchTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
matches, err := search(argumentsInJSON, t.tools)
if err != nil {
return "", err
}
result := &toolSearchResult{}
for _, m := range matches {
result.Matches = append(result.Matches, m.Name)
}
b, err := json.Marshal(result)
if err != nil {
return "", fmt.Errorf("failed to marshal tool search result: %w", err)
}
return string(b), nil
}
type modelToolSearchTool struct {
tools map[string]*schema.ToolInfo
}
func (t *modelToolSearchTool) Info(_ context.Context) (*schema.ToolInfo, error) {
return getToolSearchToolInfo(), nil
}
func (t *modelToolSearchTool) InvokableRun(_ context.Context, argumentsInJSON *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) {
ret, err := search(argumentsInJSON.Text, t.tools)
if err != nil {
return nil, err
}
return &schema.ToolResult{Parts: []schema.ToolOutputPart{
{
Type: schema.ToolPartTypeToolSearchResult,
ToolSearchResult: &schema.ToolSearchResult{
Tools: ret,
},
},
}}, nil
}
const (
toolSearchToolName = "tool_search"
defaultMaxResults = 5
)
func getToolSearchToolInfo() *schema.ToolInfo {
return &schema.ToolInfo{
Name: toolSearchToolName,
Desc: internal.SelectPrompt(internal.I18nPrompts{
English: toolDescription,
Chinese: toolDescriptionChinese,
}),
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
"query": {
Type: schema.String,
Desc: "Query to find deferred tools. Use \"select:<tool_name>\" for direct selection, or keywords to search.",
Required: true,
},
"max_results": {
Type: schema.Integer,
Desc: "Maximum number of results to return (default: 5)",
Required: false,
},
}),
}
}
func search(argumentsInJSON string, tools map[string]*schema.ToolInfo) ([]*schema.ToolInfo, error) {
var args toolSearchArgs
if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil {
return nil, fmt.Errorf("failed to unmarshal tool search arguments: %w", err)
}
query := strings.TrimSpace(args.Query)
if query == "" {
return nil, fmt.Errorf("query is required")
}
maxResults := defaultMaxResults
if args.MaxResults != nil && *args.MaxResults > 0 {
maxResults = *args.MaxResults
}
var matches []string
// Direct selection mode: select:tool1,tool2
// max_results is intentionally not applied here because the model has
// already specified the exact tools it wants by name.
if strings.HasPrefix(query, "select:") {
names := strings.Split(strings.TrimPrefix(query, "select:"), ",")
toolSet := make(map[string]bool, len(tools))
for name := range tools {
toolSet[name] = true
}
for _, name := range names {
name = strings.TrimSpace(name)
if name != "" && toolSet[name] {
matches = append(matches, name)
}
}
} else {
matches = keywordSearch(query, maxResults, tools)
}
ret := make([]*schema.ToolInfo, 0, len(matches))
for _, name := range matches {
ti, ok := tools[name]
if !ok {
continue
}
ret = append(ret, ti)
}
return ret, nil
}
func intMax(a, b int) int {
if a > b {
return a
}
return b
}
func intMin(a, b int) int {
if a < b {
return a
}
return b
}
// scoredTool pairs a tool name with its search score.
type scoredTool struct {
name string
score int
}
// keywordSearch scores all tools against the query keywords and returns the top N.
func keywordSearch(query string, maxResults int, tools map[string]*schema.ToolInfo) []string {
keywords := parseKeywords(query)
if len(keywords) == 0 {
return nil
}
var scored []scoredTool
for name, tm := range tools {
nameParts := splitToolName(name)
nameLower := strings.ToLower(name)
descLower := strings.ToLower(tm.Desc)
totalScore := 0
allRequiredFound := true
for _, kw := range keywords {
kwLower := strings.ToLower(kw.word)
kwScore := 0
// Score against name parts
for _, part := range nameParts {
partLower := strings.ToLower(part)
if partLower == kwLower {
kwScore = intMax(kwScore, 10)
} else if strings.Contains(partLower, kwLower) {
kwScore = intMax(kwScore, 5)
}
}
// Score against full name
if strings.Contains(nameLower, kwLower) {
kwScore = intMax(kwScore, 3)
}
// Score against description (substring match)
if descLower != "" && strings.Contains(descLower, kwLower) {
kwScore = intMax(kwScore, 2)
}
if kw.required && kwScore == 0 {
allRequiredFound = false
break
}
totalScore += kwScore
}
if !allRequiredFound {
continue
}
if totalScore > 0 {
scored = append(scored, scoredTool{name: name, score: totalScore})
}
}
// Sort by score descending, then by name for stability
sort.Slice(scored, func(i, j int) bool {
if scored[i].score != scored[j].score {
return scored[i].score > scored[j].score
}
return scored[i].name < scored[j].name
})
results := make([]string, 0, intMin(maxResults, len(scored)))
for i := 0; i < len(scored) && i < maxResults; i++ {
results = append(results, scored[i].name)
}
return results
}
// keyword represents a parsed search keyword.
type keyword struct {
word string
required bool
}
// parseKeywords splits a query string into keywords, handling the '+' required prefix.
func parseKeywords(query string) (keywords []keyword) {
parts := strings.Fields(query)
for _, p := range parts {
if strings.HasPrefix(p, "+") {
word := strings.TrimPrefix(p, "+")
if word != "" {
keywords = append(keywords, keyword{word: word, required: true})
}
} else if p != "" {
keywords = append(keywords, keyword{word: p, required: false})
}
}
return
}
// splitToolName splits a tool name into parts by underscores, double underscores (MCP separator),
// and camelCase boundaries.
func splitToolName(name string) []string {
// First split by double underscore (MCP server__tool separator)
segments := strings.Split(name, "__")
var parts []string
for _, seg := range segments {
// Split each segment by single underscore
underscoreParts := strings.Split(seg, "_")
for _, up := range underscoreParts {
if up == "" {
continue
}
// Further split by camelCase
camelParts := splitCamelCase(up)
parts = append(parts, camelParts...)
}
}
return parts
}
// splitCamelCase splits a camelCase or PascalCase string into its constituent words.
func splitCamelCase(s string) []string {
if s == "" {
return nil
}
var parts []string
runes := []rune(s)
start := 0
for i := 1; i < len(runes); i++ {
if unicode.IsUpper(runes[i]) {
if unicode.IsLower(runes[i-1]) {
parts = append(parts, string(runes[start:i]))
start = i
} else if i+1 < len(runes) && unicode.IsLower(runes[i+1]) {
parts = append(parts, string(runes[start:i]))
start = i
}
}
}
parts = append(parts, string(runes[start:]))
return parts
}
@@ -0,0 +1,404 @@
/*
* Copyright 2026 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package toolsearch
import (
"context"
"encoding/json"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema"
)
// ---------------------------------------------------------------------------
// Generic table-driven tests covering both *schema.Message and *schema.AgenticMessage
// ---------------------------------------------------------------------------
// --- Generic message construction helpers ---
func makeUserMsg[M adk.MessageType](content string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any(schema.UserMessage(content)).(M)
case *schema.AgenticMessage:
return any(schema.UserAgenticMessage(content)).(M)
default:
panic("unreachable")
}
}
func makeSystemMsg[M adk.MessageType](content string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any(&schema.Message{Role: schema.System, Content: content}).(M)
case *schema.AgenticMessage:
return any(schema.SystemAgenticMessage(content)).(M)
default:
panic("unreachable")
}
}
type testToolCall struct {
ID string
Name string
Arguments string
}
func makeAssistantMsgWithToolCalls[M adk.MessageType](toolCalls []testToolCall) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
tcs := make([]schema.ToolCall, len(toolCalls))
for i, tc := range toolCalls {
tcs[i] = schema.ToolCall{
ID: tc.ID,
Function: schema.FunctionCall{Name: tc.Name, Arguments: tc.Arguments},
}
}
return any(schema.AssistantMessage("", tcs)).(M)
case *schema.AgenticMessage:
blocks := make([]*schema.ContentBlock, len(toolCalls))
for i, tc := range toolCalls {
blocks[i] = schema.NewContentBlock(&schema.FunctionToolCall{
CallID: tc.ID,
Name: tc.Name,
Arguments: tc.Arguments,
})
}
return any(&schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: blocks,
}).(M)
default:
panic("unreachable")
}
}
func makeToolResultMsg[M adk.MessageType](content string, callID string, toolName string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any(&schema.Message{
Role: schema.Tool,
ToolName: toolName,
ToolCallID: callID,
Content: content,
}).(M)
case *schema.AgenticMessage:
return any(&schema.AgenticMessage{
Role: schema.AgenticRoleTypeUser,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.FunctionToolResult{
CallID: callID,
Name: toolName,
Content: []*schema.FunctionToolResultContentBlock{
{Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: content}},
},
}),
},
}).(M)
default:
panic("unreachable")
}
}
func getMsgRole[M adk.MessageType](msg M) string {
switch v := any(msg).(type) {
case *schema.Message:
return string(v.Role)
case *schema.AgenticMessage:
return string(v.Role)
default:
panic("unreachable")
}
}
func getMsgContent[M adk.MessageType](msg M) string {
switch v := any(msg).(type) {
case *schema.Message:
return v.Content
case *schema.AgenticMessage:
for _, block := range v.ContentBlocks {
if block != nil && block.Type == schema.ContentBlockTypeUserInputText && block.UserInputText != nil {
return block.UserInputText.Text
}
}
return ""
default:
panic("unreachable")
}
}
func getMsgExtra[M adk.MessageType](msg M) map[string]any {
switch v := any(msg).(type) {
case *schema.Message:
return v.Extra
case *schema.AgenticMessage:
return v.Extra
default:
panic("unreachable")
}
}
func setMsgExtra[M adk.MessageType](msg M, key string, val any) {
switch v := any(msg).(type) {
case *schema.Message:
if v.Extra == nil {
v.Extra = make(map[string]any)
}
v.Extra[key] = val
case *schema.AgenticMessage:
if v.Extra == nil {
v.Extra = make(map[string]any)
}
v.Extra[key] = val
default:
panic("unreachable")
}
}
func newTestMiddlewareTyped[M adk.MessageType](t *testing.T, tools []tool.BaseTool) *typedMiddleware[M] {
t.Helper()
ctx := context.Background()
mw, err := NewTyped[M](ctx, &Config{
DynamicTools: tools,
UseModelToolSearch: false,
})
require.NoError(t, err)
return mw.(*typedMiddleware[M])
}
func countRemindersGeneric[M adk.MessageType](msgs []M) int {
count := 0
for _, msg := range msgs {
extra := getMsgExtra(msg)
if extra != nil {
if v, _ := extra[toolSearchReminderExtraKey].(bool); v {
count++
}
}
}
return count
}
// --- Generic test functions ---
func testEnsureReminderGeneric[M adk.MessageType](t *testing.T) {
dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"}
m := newTestMiddlewareTyped[M](t, []tool.BaseTool{dynamicA})
t.Run("normal: system then user", func(t *testing.T) {
input := []M{
makeSystemMsg[M]("sys"),
makeUserMsg[M]("hi"),
}
got := m.ensureReminder(input)
require.Len(t, got, 3)
assert.Equal(t, "system", getMsgRole(got[0]))
// Reminder inserted after system
extra := getMsgExtra(got[1])
require.NotNil(t, extra)
assert.Equal(t, true, extra[toolSearchReminderExtraKey])
assert.Equal(t, "hi", getMsgContent(got[2]))
})
t.Run("all system messages", func(t *testing.T) {
input := []M{
makeSystemMsg[M]("sys1"),
makeSystemMsg[M]("sys2"),
}
got := m.ensureReminder(input)
require.Len(t, got, 3)
assert.Equal(t, "system", getMsgRole(got[0]))
assert.Equal(t, "system", getMsgRole(got[1]))
// Reminder appended at end
extra := getMsgExtra(got[2])
require.NotNil(t, extra)
assert.Equal(t, true, extra[toolSearchReminderExtraKey])
})
t.Run("empty input", func(t *testing.T) {
got := m.ensureReminder(nil)
require.Len(t, got, 1)
extra := getMsgExtra(got[0])
require.NotNil(t, extra)
assert.Equal(t, true, extra[toolSearchReminderExtraKey])
})
t.Run("no system messages", func(t *testing.T) {
input := []M{
makeUserMsg[M]("hi"),
}
got := m.ensureReminder(input)
require.Len(t, got, 2)
// Reminder inserted at position 0
extra := getMsgExtra(got[0])
require.NotNil(t, extra)
assert.Equal(t, true, extra[toolSearchReminderExtraKey])
assert.Equal(t, "hi", getMsgContent(got[1]))
})
t.Run("idempotent: does not insert twice", func(t *testing.T) {
reminder := makeUserMsg[M]("<reminder>")
setMsgExtra(reminder, toolSearchReminderExtraKey, true)
input := []M{
reminder,
makeUserMsg[M]("hi"),
}
got := m.ensureReminder(input)
require.Len(t, got, 2)
assert.Equal(t, "hi", getMsgContent(got[1]))
})
}
func testMode1InitializationGeneric[M adk.MessageType](t *testing.T) {
dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"}
dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"}
m := newTestMiddlewareTyped[M](t, []tool.BaseTool{dynamicA, dynamicB})
ctx := context.Background()
state := &adk.TypedChatModelAgentState[M]{
Messages: []M{
makeSystemMsg[M]("sys"),
makeUserMsg[M]("hello"),
},
ToolInfos: []*schema.ToolInfo{
ti("static_tool", "Static tool"),
getToolSearchToolInfo(),
ti("dynamic_tool_a", "Dynamic tool A"),
ti("dynamic_tool_b", "Dynamic tool B"),
},
}
// Initialization strips dynamic tools, keeps tool_search and static tools.
_, state, err := m.BeforeModelRewriteState(ctx, state, nil)
require.NoError(t, err)
names := toolNames(state.ToolInfos)
assert.Equal(t, []string{"static_tool", "tool_search"}, names)
assert.Nil(t, state.DeferredToolInfos, "Mode 1 should not populate DeferredToolInfos")
// Verify reminder was inserted.
assert.Equal(t, 1, countRemindersGeneric(state.Messages), "reminder should be inserted")
}
func testMode1ForwardSelectionGeneric[M adk.MessageType](t *testing.T) {
dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"}
dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"}
m := newTestMiddlewareTyped[M](t, []tool.BaseTool{dynamicA, dynamicB})
ctx := context.Background()
// Simulate state AFTER initialization (dynamic tools already stripped).
// Include a tool_search result message that selected dynamic_tool_a.
toolSearchResultJSON, _ := json.Marshal(toolSearchResult{Matches: []string{"dynamic_tool_a"}})
// Build the reminder message with the extra marker
reminderMsg := makeUserMsg[M]("hello")
setMsgExtra(reminderMsg, toolSearchReminderExtraKey, true)
state := &adk.TypedChatModelAgentState[M]{
Messages: []M{
makeSystemMsg[M]("sys"),
reminderMsg,
makeAssistantMsgWithToolCalls[M]([]testToolCall{
{ID: "tc1", Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`},
}),
makeToolResultMsg[M](string(toolSearchResultJSON), "tc1", toolSearchToolName),
},
ToolInfos: []*schema.ToolInfo{
ti("static_tool", "Static tool"),
getToolSearchToolInfo(),
},
}
// Forward selection should add dynamic_tool_a from the tool_search result.
_, state, err := m.BeforeModelRewriteState(ctx, state, nil)
require.NoError(t, err)
names := toolNames(state.ToolInfos)
assert.Equal(t, []string{"dynamic_tool_a", "static_tool", "tool_search"}, names)
// Call again: forward selection should be idempotent (dynamic_tool_a already present).
_, state, err = m.BeforeModelRewriteState(ctx, state, nil)
require.NoError(t, err)
names = toolNames(state.ToolInfos)
assert.Equal(t, []string{"dynamic_tool_a", "static_tool", "tool_search"}, names)
}
func testMalformedJSONGeneric[M adk.MessageType](t *testing.T) {
dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"}
m := newTestMiddlewareTyped[M](t, []tool.BaseTool{dynamicA})
ctx := context.Background()
// Build the reminder message with the extra marker
reminderMsg := makeUserMsg[M]("reminder")
setMsgExtra(reminderMsg, toolSearchReminderExtraKey, true)
state := &adk.TypedChatModelAgentState[M]{
Messages: []M{
makeSystemMsg[M]("sys"),
reminderMsg,
makeAssistantMsgWithToolCalls[M]([]testToolCall{
{ID: "tc1", Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`},
}),
makeToolResultMsg[M](`{invalid json!!!`, "tc1", toolSearchToolName),
},
ToolInfos: []*schema.ToolInfo{
ti("static_tool", "Static tool"),
getToolSearchToolInfo(),
},
}
_, state, err := m.BeforeModelRewriteState(ctx, state, nil)
require.NoError(t, err, "malformed JSON in tool_search result should not cause an error")
names := toolNames(state.ToolInfos)
assert.NotContains(t, names, "dynamic_tool_a", "malformed JSON result should be skipped")
assert.Contains(t, names, "static_tool")
assert.Contains(t, names, "tool_search")
}
// --- Top-level generic test runner ---
func TestToolSearchGeneric(t *testing.T) {
t.Run("Message", func(t *testing.T) {
t.Run("EnsureReminder", testEnsureReminderGeneric[*schema.Message])
t.Run("Mode1Init", testMode1InitializationGeneric[*schema.Message])
t.Run("Mode1ForwardSelection", testMode1ForwardSelectionGeneric[*schema.Message])
t.Run("MalformedJSON", testMalformedJSONGeneric[*schema.Message])
})
t.Run("AgenticMessage", func(t *testing.T) {
t.Run("EnsureReminder", testEnsureReminderGeneric[*schema.AgenticMessage])
t.Run("Mode1Init", testMode1InitializationGeneric[*schema.AgenticMessage])
t.Run("Mode1ForwardSelection", testMode1ForwardSelectionGeneric[*schema.AgenticMessage])
t.Run("MalformedJSON", testMalformedJSONGeneric[*schema.AgenticMessage])
})
}
File diff suppressed because it is too large Load Diff
+33
View File
@@ -0,0 +1,33 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Package filesystem provides middlewares.
package filesystem
import (
"github.com/cloudwego/eino/adk/filesystem"
)
type FileInfo = filesystem.FileInfo
type GrepMatch = filesystem.GrepMatch
type LsInfoRequest = filesystem.LsInfoRequest
type ReadRequest = filesystem.ReadRequest
type MultiModalReadRequest = filesystem.MultiModalReadRequest
type GrepRequest = filesystem.GrepRequest
type GlobInfoRequest = filesystem.GlobInfoRequest
type WriteRequest = filesystem.WriteRequest
type EditRequest = filesystem.EditRequest
type FileContent = filesystem.FileContent
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,177 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package filesystem
import (
"bufio"
"context"
"errors"
"fmt"
"io"
"strings"
"unicode/utf8"
"github.com/slongfield/pyfmt"
"github.com/cloudwego/eino/adk/filesystem"
"github.com/cloudwego/eino/adk/internal"
"github.com/cloudwego/eino/compose"
"github.com/cloudwego/eino/schema"
)
type toolResultOffloadingConfig struct {
Backend filesystem.Backend
TokenLimit int
PathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error)
}
func newToolResultOffloading(ctx context.Context, config *toolResultOffloadingConfig) compose.ToolMiddleware {
offloading := &toolResultOffloading{
backend: config.Backend,
tokenLimit: config.TokenLimit,
pathGenerator: config.PathGenerator,
}
if offloading.tokenLimit == 0 {
offloading.tokenLimit = 20000
}
if offloading.pathGenerator == nil {
offloading.pathGenerator = func(ctx context.Context, input *compose.ToolInput) (string, error) {
return fmt.Sprintf("/large_tool_result/%s", input.CallID), nil
}
}
return compose.ToolMiddleware{
Invokable: offloading.invoke,
Streamable: offloading.stream,
}
}
type toolResultOffloading struct {
backend filesystem.Backend
tokenLimit int
pathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error)
}
func (t *toolResultOffloading) invoke(endpoint compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
output, err := endpoint(ctx, input)
if err != nil {
return nil, err
}
result, err := t.handleResult(ctx, output.Result, input)
if err != nil {
return nil, err
}
return &compose.ToolOutput{Result: result}, nil
}
}
func (t *toolResultOffloading) stream(endpoint compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
output, err := endpoint(ctx, input)
if err != nil {
return nil, err
}
result, err := concatString(output.Result)
if err != nil {
return nil, err
}
result, err = t.handleResult(ctx, result, input)
if err != nil {
return nil, err
}
return &compose.StreamToolOutput{Result: schema.StreamReaderFromArray([]string{result})}, nil
}
}
func (t *toolResultOffloading) handleResult(ctx context.Context, result string, input *compose.ToolInput) (string, error) {
if len(result) > t.tokenLimit*4 {
path, err := t.pathGenerator(ctx, input)
if err != nil {
return "", err
}
nResult := formatToolMessage(result)
msgTemplate := internal.SelectPrompt(internal.I18nPrompts{
English: tooLargeToolMessage,
Chinese: tooLargeToolMessageChinese,
})
nResult, err = pyfmt.Fmt(msgTemplate, map[string]any{
"tool_call_id": input.CallID,
"file_path": path,
"content_sample": nResult,
})
if err != nil {
return "", err
}
err = t.backend.Write(ctx, &WriteRequest{
FilePath: path,
Content: result,
})
if err != nil {
return "", err
}
return nResult, nil
}
return result, nil
}
func concatString(sr *schema.StreamReader[string]) (string, error) {
if sr == nil {
return "", errors.New("stream is nil")
}
sb := strings.Builder{}
for {
str, err := sr.Recv()
if errors.Is(err, io.EOF) {
return sb.String(), nil
}
if err != nil {
return "", err
}
sb.WriteString(str)
}
}
func formatToolMessage(s string) string {
reader := bufio.NewScanner(strings.NewReader(s))
var b strings.Builder
lineNum := 1
for reader.Scan() {
if lineNum > 10 {
break
}
line := reader.Text()
if utf8.RuneCountInString(line) > 1000 {
runes := []rune(line)
line = string(runes[:1000])
}
b.WriteString(fmt.Sprintf("%d: %s\n", lineNum, line))
lineNum++
}
return b.String()
}
@@ -0,0 +1,611 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package filesystem
import (
"context"
"errors"
"fmt"
"io"
"strings"
"testing"
"github.com/cloudwego/eino/compose"
"github.com/cloudwego/eino/schema"
)
// mockBackend is a simple in-memory backend for testing
type mockBackend struct {
files map[string]string
}
func newMockBackend() *mockBackend {
return &mockBackend{
files: make(map[string]string),
}
}
func (m *mockBackend) Write(ctx context.Context, req *WriteRequest) error {
m.files[req.FilePath] = req.Content
return nil
}
func (m *mockBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) {
content, ok := m.files[req.FilePath]
if !ok {
return nil, errors.New("file not found")
}
return &FileContent{Content: content}, nil
}
func (m *mockBackend) LsInfo(ctx context.Context, _ *LsInfoRequest) ([]FileInfo, error) {
return nil, nil
}
func (m *mockBackend) GrepRaw(ctx context.Context, _ *GrepRequest) ([]GrepMatch, error) {
return nil, nil
}
func (m *mockBackend) GlobInfo(ctx context.Context, _ *GlobInfoRequest) ([]FileInfo, error) {
return nil, nil
}
func (m *mockBackend) Edit(ctx context.Context, _ *EditRequest) error {
return nil
}
func TestToolResultOffloading_SmallResult(t *testing.T) {
ctx := context.Background()
backend := newMockBackend()
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 100, // Small limit for testing
}
middleware := newToolResultOffloading(ctx, config)
// Create a mock endpoint that returns a small result
smallResult := "This is a small result"
mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
return &compose.ToolOutput{Result: smallResult}, nil
}
// Wrap the endpoint with the middleware
wrappedEndpoint := middleware.Invokable(mockEndpoint)
// Execute
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_123",
}
output, err := wrappedEndpoint(ctx, input)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Small result should pass through unchanged
if output.Result != smallResult {
t.Errorf("expected result %q, got %q", smallResult, output.Result)
}
// No file should be written
if len(backend.files) != 0 {
t.Errorf("expected no files to be written, got %d files", len(backend.files))
}
}
func TestToolResultOffloading_LargeResult(t *testing.T) {
ctx := context.Background()
backend := newMockBackend()
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 10, // Very small limit to trigger offloading
}
middleware := newToolResultOffloading(ctx, config)
// Create a large result (more than 10 * 4 = 40 bytes)
largeResult := strings.Repeat("This is a long line of text that will exceed the token limit.\n", 10)
mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
return &compose.ToolOutput{Result: largeResult}, nil
}
wrappedEndpoint := middleware.Invokable(mockEndpoint)
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_456",
}
output, err := wrappedEndpoint(ctx, input)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Result should be replaced with a message
if !strings.Contains(output.Result, "Tool result too large") {
t.Errorf("expected result to contain 'Tool result too large', got %q", output.Result)
}
if !strings.Contains(output.Result, "call_456") {
t.Errorf("expected result to contain call ID 'call_456', got %q", output.Result)
}
if !strings.Contains(output.Result, "/large_tool_result/call_456") {
t.Errorf("expected result to contain file path, got %q", output.Result)
}
// File should be written
if len(backend.files) != 1 {
t.Fatalf("expected 1 file to be written, got %d files", len(backend.files))
}
savedContent, ok := backend.files["/large_tool_result/call_456"]
if !ok {
t.Fatalf("expected file at /large_tool_result/call_456, got files: %v", backend.files)
}
if savedContent != largeResult {
t.Errorf("saved content doesn't match original result")
}
}
func TestToolResultOffloading_CustomPathGenerator(t *testing.T) {
ctx := context.Background()
backend := newMockBackend()
customPath := "/custom/path/result.txt"
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 10,
PathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) {
return customPath, nil
},
}
middleware := newToolResultOffloading(ctx, config)
largeResult := strings.Repeat("Large content ", 100)
mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
return &compose.ToolOutput{Result: largeResult}, nil
}
wrappedEndpoint := middleware.Invokable(mockEndpoint)
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_789",
}
output, err := wrappedEndpoint(ctx, input)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Check custom path is used
if !strings.Contains(output.Result, customPath) {
t.Errorf("expected result to contain custom path %q, got %q", customPath, output.Result)
}
// File should be written to custom path
savedContent, ok := backend.files[customPath]
if !ok {
t.Fatalf("expected file at %q, got files: %v", customPath, backend.files)
}
if savedContent != largeResult {
t.Errorf("saved content doesn't match original result")
}
}
func TestToolResultOffloading_PathGeneratorError(t *testing.T) {
ctx := context.Background()
backend := newMockBackend()
expectedErr := errors.New("path generation failed")
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 10,
PathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) {
return "", expectedErr
},
}
middleware := newToolResultOffloading(ctx, config)
largeResult := strings.Repeat("Large content ", 100)
mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
return &compose.ToolOutput{Result: largeResult}, nil
}
wrappedEndpoint := middleware.Invokable(mockEndpoint)
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_error",
}
_, err := wrappedEndpoint(ctx, input)
if err == nil {
t.Fatal("expected error, got nil")
}
if !errors.Is(err, expectedErr) {
t.Errorf("expected error %v, got %v", expectedErr, err)
}
}
func TestToolResultOffloading_EndpointError(t *testing.T) {
ctx := context.Background()
backend := newMockBackend()
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 100,
}
middleware := newToolResultOffloading(ctx, config)
expectedErr := errors.New("endpoint execution failed")
mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
return nil, expectedErr
}
wrappedEndpoint := middleware.Invokable(mockEndpoint)
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_endpoint_error",
}
_, err := wrappedEndpoint(ctx, input)
if err == nil {
t.Fatal("expected error, got nil")
}
if !errors.Is(err, expectedErr) {
t.Errorf("expected error %v, got %v", expectedErr, err)
}
}
func TestToolResultOffloading_DefaultTokenLimit(t *testing.T) {
ctx := context.Background()
backend := newMockBackend()
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 0, // Should default to 20000
}
middleware := newToolResultOffloading(ctx, config)
// Create a result smaller than 20000 * 4 = 80000 bytes
smallResult := strings.Repeat("x", 1000)
mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
return &compose.ToolOutput{Result: smallResult}, nil
}
wrappedEndpoint := middleware.Invokable(mockEndpoint)
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_default",
}
output, err := wrappedEndpoint(ctx, input)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Should pass through unchanged
if output.Result != smallResult {
t.Errorf("expected result to pass through unchanged")
}
// No file should be written
if len(backend.files) != 0 {
t.Errorf("expected no files to be written, got %d files", len(backend.files))
}
}
func TestToolResultOffloading_Stream(t *testing.T) {
ctx := context.Background()
backend := newMockBackend()
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 10,
}
middleware := newToolResultOffloading(ctx, config)
// Create a streaming endpoint that returns large content
largeResult := strings.Repeat("Large streaming content ", 100)
mockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
// Split the result into chunks
chunks := []string{largeResult[:len(largeResult)/2], largeResult[len(largeResult)/2:]}
return &compose.StreamToolOutput{
Result: schema.StreamReaderFromArray(chunks),
}, nil
}
wrappedEndpoint := middleware.Streamable(mockStreamEndpoint)
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_stream",
}
output, err := wrappedEndpoint(ctx, input)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Read the stream
var result strings.Builder
for {
chunk, err := output.Result.Recv()
if errors.Is(err, io.EOF) {
break
}
if err != nil {
t.Fatalf("error reading stream: %v", err)
}
result.WriteString(chunk)
}
resultStr := result.String()
// Result should be replaced with a message
if !strings.Contains(resultStr, "Tool result too large") {
t.Errorf("expected result to contain 'Tool result too large', got %q", resultStr)
}
if !strings.Contains(resultStr, "call_stream") {
t.Errorf("expected result to contain call ID 'call_stream', got %q", resultStr)
}
// File should be written
if len(backend.files) != 1 {
t.Fatalf("expected 1 file to be written, got %d files", len(backend.files))
}
savedContent, ok := backend.files["/large_tool_result/call_stream"]
if !ok {
t.Fatalf("expected file at /large_tool_result/call_stream, got files: %v", backend.files)
}
if savedContent != largeResult {
t.Errorf("saved content doesn't match original result")
}
}
func TestToolResultOffloading_StreamError(t *testing.T) {
ctx := context.Background()
backend := newMockBackend()
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 10,
}
middleware := newToolResultOffloading(ctx, config)
expectedErr := errors.New("stream endpoint failed")
mockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
return nil, expectedErr
}
wrappedEndpoint := middleware.Streamable(mockStreamEndpoint)
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_stream_error",
}
_, err := wrappedEndpoint(ctx, input)
if err == nil {
t.Fatal("expected error, got nil")
}
if !errors.Is(err, expectedErr) {
t.Errorf("expected error %v, got %v", expectedErr, err)
}
}
func TestFormatToolMessage(t *testing.T) {
tests := []struct {
name string
input string
expected string
}{
{
name: "single line",
input: "single line",
expected: "1: single line\n",
},
{
name: "multiple lines",
input: "line1\nline2\nline3",
expected: "1: line1\n2: line2\n3: line3\n",
},
{
name: "more than 10 lines",
input: "1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12",
expected: "1: 1\n2: 2\n3: 3\n4: 4\n5: 5\n6: 6\n7: 7\n8: 8\n9: 9\n10: 10\n",
},
{
name: "long line truncation",
input: strings.Repeat("a", 1500),
expected: fmt.Sprintf("1: %s\n", strings.Repeat("a", 1000)),
},
{
name: "unicode characters",
input: "你好世界\n测试",
expected: "1: 你好世界\n2: 测试\n",
},
{
name: "long unicode line",
input: strings.Repeat("你", 1500),
expected: fmt.Sprintf("1: %s\n", strings.Repeat("你", 1000)),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := formatToolMessage(tt.input)
if result != tt.expected {
t.Errorf("formatToolMessage() = %q, want %q", result, tt.expected)
}
})
}
}
func TestConcatString(t *testing.T) {
tests := []struct {
name string
chunks []string
expected string
expectError bool
}{
{
name: "single chunk",
chunks: []string{"hello"},
expected: "hello",
},
{
name: "multiple chunks",
chunks: []string{"hello", " ", "world"},
expected: "hello world",
},
{
name: "empty chunks",
chunks: []string{"", "", ""},
expected: "",
},
{
name: "mixed chunks",
chunks: []string{"a", "", "b", "c"},
expected: "abc",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
sr := schema.StreamReaderFromArray(tt.chunks)
result, err := concatString(sr)
if tt.expectError {
if err == nil {
t.Error("expected error, got nil")
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if result != tt.expected {
t.Errorf("concatString() = %q, want %q", result, tt.expected)
}
})
}
// Test nil stream
t.Run("nil stream", func(t *testing.T) {
_, err := concatString(nil)
if err == nil {
t.Error("expected error for nil stream, got nil")
}
if !strings.Contains(err.Error(), "stream is nil") {
t.Errorf("expected 'stream is nil' error, got %v", err)
}
})
}
func TestToolResultOffloading_BackendWriteError(t *testing.T) {
ctx := context.Background()
// Create a backend that fails on write
backend := &failingBackend{
writeErr: errors.New("write failed"),
}
config := &toolResultOffloadingConfig{
Backend: backend,
TokenLimit: 10,
}
middleware := newToolResultOffloading(ctx, config)
largeResult := strings.Repeat("Large content ", 100)
mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
return &compose.ToolOutput{Result: largeResult}, nil
}
wrappedEndpoint := middleware.Invokable(mockEndpoint)
input := &compose.ToolInput{
Name: "test_tool",
CallID: "call_write_error",
}
_, err := wrappedEndpoint(ctx, input)
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), "write failed") {
t.Errorf("expected 'write failed' error, got %v", err)
}
}
// failingBackend is a mock backend that can be configured to fail
type failingBackend struct {
writeErr error
}
func (f *failingBackend) Write(ctx context.Context, req *WriteRequest) error {
if f.writeErr != nil {
return f.writeErr
}
return nil
}
func (f *failingBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) {
return &FileContent{}, nil
}
func (f *failingBackend) LsInfo(ctx context.Context, _ *LsInfoRequest) ([]FileInfo, error) {
return nil, nil
}
func (f *failingBackend) GrepRaw(ctx context.Context, _ *GrepRequest) ([]GrepMatch, error) {
return nil, nil
}
func (f *failingBackend) GlobInfo(ctx context.Context, _ *GlobInfoRequest) ([]FileInfo, error) {
return nil, nil
}
func (f *failingBackend) Edit(ctx context.Context, _ *EditRequest) error {
return nil
}
+270
View File
@@ -0,0 +1,270 @@
/*
* Copyright (c) 2025 Harrison Chase
* Copyright (c) 2025 CloudWeGo Authors
* SPDX-License-Identifier: MIT
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package filesystem
// This file contains prompt templates and tool descriptions adapted from the DeepAgents project.
// Original source: https://github.com/langchain-ai/deepagents
//
// These prompts are used under the terms of the original project's open source license.
// When using this code in your own open source project, ensure compliance with the original license requirements.
const (
tooLargeToolMessage = `Tool result too large, the result of this tool call {tool_call_id} was saved in the filesystem at this path: {file_path}
You 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.
You can do this by specifying an offset and limit in the read_file tool call.
For example, to read the first 100 lines, you can use the read_file tool with offset=0 and limit=100.
Here are the first 10 lines of the result:
{content_sample}`
tooLargeToolMessageChinese = `工具结果过大,此工具调用 {tool_call_id} 的结果已保存到文件系统的以下路径:{file_path}
你可以使用 read_file 工具从文件系统读取结果,但请确保每次只读取部分结果。
你可以通过在 read_file 工具调用中指定 offset 和 limit 来实现。
例如,要读取前 100 行,你可以使用 read_file 工具,设置 offset=0 和 limit=100。
以下是结果的前 10 行:
{content_sample}`
ListFilesToolDesc = `Lists all files in the filesystem, filtering by directory.
Usage:
- The path parameter must be an absolute path, not a relative path
- The ls tool will return a list of all files in the specified directory.
- This is very useful for exploring the file system and finding the right file to read or edit.
- You should almost ALWAYS use this tool before using the read_file or edit_file tools.`
ListFilesToolDescChinese = `列出文件系统中的所有文件,按目录过滤。
使用方法:
- path 参数必须是绝对路径,不能是相对路径
- ls 工具将返回指定目录中所有文件的列表
- 这对于探索文件系统和找到要读取或编辑的正确文件非常有用
- 在使用 read_file 或 edit_file 工具之前,你几乎总是应该先使用此工具`
ReadFileToolDesc = `Reads a file from the filesystem. You can access any file directly by using this tool.
Assume 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.
Usage:
- The file_path parameter must be an absolute path, not a relative path
- By default, it reads up to 2000 lines starting from the beginning of the file
- **IMPORTANT for large files and codebase exploration**: Use pagination with offset and limit parameters to avoid context overflow
- First scan: read_file(path, limit=100) to see file structure
- Read more sections: read_file(path, offset=100, limit=200) for next 200 lines
- Only omit limit (read full file) when necessary for editing
- Specify offset and limit: read_file(path, offset=0, limit=100) reads first 100 lines
- Results are returned using cat -n format, with line numbers starting at 1
- 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.
- If you read a file that exists but has empty contents you will receive a system reminder warning in place of file contents.
- You should ALWAYS make sure a file has been read before editing it.`
ReadFileToolDescChinese = `从文件系统读取文件。你可以使用此工具直接访问任何文件。
假设此工具能够读取机器上的所有文件。如果用户提供了文件路径,假设该路径是有效的。读取不存在的文件是可以的;将返回错误。
使用方法:
- file_path 参数必须是绝对路径,不能是相对路径
- 默认情况下,从文件开头读取最多 2000 行
- **大文件和代码库探索的重要提示**:使用 offset 和 limit 参数进行分页,以避免上下文溢出
- 首次扫描:read_file(path, limit=100) 查看文件结构
- 读取更多部分:read_file(path, offset=100, limit=200) 读取接下来的 200 行
- 仅在编辑必要时才省略 limit(读取完整文件)
- 指定 offset 和 limitread_file(path, offset=0, limit=100) 读取前 100 行
- 结果以 cat -n 格式返回,行号从 1 开始
- 你可以在单个响应中调用多个工具。最好同时推测性地批量读取多个可能有用的文件
- 如果你读取的文件存在但内容为空,你将收到系统提醒警告而不是文件内容
- 在编辑文件之前,你应该始终确保已读取该文件`
// EnhancedReadFileDescSuffix is appended to ReadFileToolDesc when using MultiModalReadFileTool.
EnhancedReadFileDescSuffix = `
- This tool supports reading image files (e.g., PNG, JPG, etc.). When reading an image file, the contents are presented visually, as the underlying model is a multimodal LLM.
- This tool can read PDF files (.pdf). For large PDFs (more than 10 pages), you MUST provide the pages parameter to read specific page ranges (e.g., pages: "1-5"). Reading a large PDF without the pages parameter will fail. Maximum 20 pages per request.`
EnhancedReadFileDescSuffixChinese = `
- 此工具支持读取图片文件(如 PNG、JPG 等)。读取图片文件时,内容将以视觉方式呈现,因为底层模型是多模态 LLM。
- 此工具可以读取 PDF 文件(.pdf)。对于大型 PDF(超过 10 页),你必须提供 pages 参数来指定页面范围(例如 pages: "1-5")。不提供 pages 参数读取大型 PDF 将会失败。每次请求最多 20 页。`
EditFileToolDesc = `Performs exact string replacements in files.
Usage:
- 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.
- 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.
- ALWAYS prefer editing existing files. NEVER write new files unless explicitly required.
- Only use emojis if the user explicitly requests it. Avoid adding emojis to files unless asked.
- 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'.
- Use 'replace_all' for replacing and renaming strings across the file. This parameter is useful if you want to rename a variable for instance.`
EditFileToolDescChinese = `在文件中执行精确的字符串替换。
使用方法:
- 在编辑之前,你必须在对话中至少使用一次 'read_file' 工具。如果你在未读取文件的情况下尝试编辑,此工具将报错
- 当从 Read 工具输出编辑文本时,请确保保留行号前缀之后的确切缩进(制表符/空格)。行号前缀格式为:空格 + 行号 + 制表符。制表符之后的所有内容都是要匹配的实际文件内容。永远不要在 old_string 或 new_string 中包含行号前缀的任何部分
- 始终优先编辑现有文件。除非明确要求,否则不要创建新文件
- 仅在用户明确要求时使用表情符号。除非被要求,否则避免在文件中添加表情符号
- 如果 'old_string' 在文件中不唯一,编辑将失败。要么提供包含更多上下文的更长字符串使其唯一,要么使用 'replace_all' 更改 'old_string' 的每个实例
- 使用 'replace_all' 在整个文件中替换和重命名字符串。例如,如果你想重命名变量,此参数很有用`
WriteFileToolDesc = `Writes a file to the local filesystem.
Usage:
- This tool will overwrite the existing file if there is one at the provided path.
- 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.
- ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required.
- NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User.
- Only use emojis if the user explicitly requests it. Avoid writing emojis to files unless asked.`
WriteFileToolDescChinese = `将文件写入本地文件系统。
使用方法:
- 如果提供的路径已存在文件,此工具将覆盖现有文件
- 如果这是一个现有文件,你必须先使用 Read 工具读取文件内容。如果你没有先读取文件,此工具将失败
- 始终优先编辑代码库中的现有文件。除非明确要求,否则不要创建新文件
- 不要主动创建文档文件(*.md)或 README 文件。仅在用户明确要求时才创建文档文件
- 仅在用户明确要求时使用表情符号。除非被要求,否则避免在文件中写入表情符号`
GlobToolDesc = `Fast file pattern matching tool that works with any codebase size
- Supports glob patterns like "**/*.js" or "src/**/*.ts"
- Returns matching file paths sorted by modification time
- Use this tool when you need to find files by name patterns
- 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.
Examples:
- '**/*.py' - Find all Python files
- '*.txt' - Find all text files in root
- '/subdir/**/*.md' - Find all markdown files under /subdir`
GlobToolDescChinese = `适用于任何代码库大小的快速文件模式匹配工具
- 支持 glob 模式,如 "**/*.js" 或 "src/**/*.ts"
- 返回按修改时间排序的匹配文件路径
- 当你需要按名称模式查找文件时使用此工具
- 你可以在单个响应中调用多个工具。最好同时并行执行多个可能有用的搜索
示例:
- '**/*.py' - 查找所有 Python 文件
- '*.txt' - 查找根目录中的所有文本文件
- '/subdir/**/*.md' - 查找 /subdir 下的所有 markdown 文件`
GrepToolDesc = `
A powerful search tool built on ripgrep
Usage:
- 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.
- Supports full regex syntax (e.g., "log.*Error", "function\s+\w+")
- Filter files with glob parameter (e.g., "*.js", "**/*.tsx") or type parameter (e.g., "js", "py", "rust")
- Output modes: "content" shows matching lines, "files_with_matches" shows only file paths (default), "count" shows match counts
- Use Task tool for open-ended searches requiring multiple rounds
- Pattern syntax: Uses ripgrep (not grep) - literal braces need escaping (use 'interface\{\}' to find 'interface{}' in Go code)
- Multiline matching: By default patterns match within single lines only. For cross-line patterns like 'struct \{[\s\S]*?field', use 'multiline: true'`
GrepToolDescChinese = `
基于 ripgrep 的强大搜索工具
使用方法:
- 始终使用 Grep 进行搜索任务。不要将 'grep' 或 'rg' 作为 Bash 命令调用。Grep 工具已针对正确的权限和访问进行了优化
- 支持完整的正则表达式语法(例如,"log.*Error""function\s+\w+"
- 使用 glob 参数(例如,"*.js""**/*.tsx")或 type 参数(例如,"js""py""rust")过滤文件
- 输出模式:"content" 显示匹配行,"files_with_matches" 仅显示文件路径(默认),"count" 显示匹配计数
- 对于需要多轮的开放式搜索,使用 Task 工具
- 模式语法:使用 ripgrep(不是 grep)- 字面大括号需要转义(使用 'interface\{\}' 在 Go 代码中查找 'interface{}'
- 多行匹配:默认情况下,模式仅在单行内匹配。对于跨行模式如 'struct \{[\s\S]*?field',使用 'multiline: true'`
ExecuteToolDesc = `
Executes a given command in the sandbox environment with proper handling and security measures.
Before executing the command, please follow these steps:
1. Directory Verification:
- 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
- For example, before running "mkdir foo/bar", first use ls to check that "foo" exists and is the intended parent directory
2. Command Execution:
- Always quote file paths that contain spaces with double quotes (e.g., cd "path with spaces/file.txt")
- Examples of proper quoting:
- cd "/Users/name/My Documents" (correct)
- cd /Users/name/My Documents (incorrect - will fail)
- python "/path/with spaces/script.py" (correct)
- python /path/with spaces/script.py (incorrect - will fail)
- After ensuring proper quoting, execute the command
- Capture the output of the command
Usage notes:
- The command parameter is required
- Commands run in an isolated sandbox environment
- Returns combined stdout/stderr output with exit code
- If the output is very large, it may be truncated
- 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.
- When issuing multiple commands, use the ';' or '&&' operator to separate them. DO NOT use newlines (newlines are ok in quoted strings)
- Use '&&' when commands depend on each other (e.g., "mkdir dir && cd dir")
- Use ';' only when you need to run commands sequentially but don't care if earlier commands fail
- Try to maintain your current working directory throughout the session by using absolute paths and avoiding usage of cd
Examples:
Good examples:
- execute(command="pytest /foo/bar/tests")
- execute(command="python /path/to/script.py")
- execute(command="npm install && npm test")
Bad examples (avoid these):
- execute(command="cd /foo/bar && pytest tests") # Use absolute path instead
- execute(command="cat file.txt") # Use read_file tool instead
- execute(command="find . -name '*.py'") # Use glob tool instead
- execute(command="grep -r 'pattern' .") # Use grep tool instead
`
ExecuteToolDescChinese = `
在沙箱环境中执行给定命令,具有适当的处理和安全措施。
执行命令前,请按照以下步骤操作:
1. 目录验证:
- 如果命令将创建新目录或文件,首先使用 ls 工具验证父目录是否存在且是正确的位置
- 例如,在运行 "mkdir foo/bar" 之前,首先使用 ls 检查 "foo" 是否存在且是预期的父目录
2. 命令执行:
- 始终用双引号引用包含空格的文件路径(例如,cd "path with spaces/file.txt"
- 正确引用的示例:
- cd "/Users/name/My Documents"(正确)
- cd /Users/name/My Documents(错误 - 将失败)
- python "/path/with spaces/script.py"(正确)
- python /path/with spaces/script.py(错误 - 将失败)
- 确保正确引用后,执行命令
- 捕获命令的输出
使用说明:
- command 参数是必需的
- 命令在隔离的沙箱环境中运行
- 返回合并的 stdout/stderr 输出和退出代码
- 如果输出非常大,可能会被截断
- 非常重要:你必须避免使用 find 和 grep 等搜索命令。请改用 grep、glob 工具进行搜索。你必须避免使用 cat、head、tail 等读取工具,请使用 read_file 读取文件
- 发出多个命令时,使用 ';' 或 '&&' 运算符分隔它们。不要使用换行符(引号字符串中的换行符是可以的)
- 当命令相互依赖时使用 '&&'(例如,"mkdir dir && cd dir"
- 仅当你需要按顺序运行命令但不关心早期命令是否失败时使用 ';'
- 尝试通过使用绝对路径并避免使用 cd 来在整个会话中保持当前工作目录
示例:
好的示例:
- execute(command="pytest /foo/bar/tests")
- execute(command="python /path/to/script.py")
- execute(command="npm install && npm test")
不好的示例(避免这些):
- execute(command="cd /foo/bar && pytest tests") # 改用绝对路径
- execute(command="cat file.txt") # 改用 read_file 工具
- execute(command="find . -name '*.py'") # 改用 glob 工具
- execute(command="grep -r 'pattern' .") # 改用 grep 工具
`
)
@@ -0,0 +1,263 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Package patchtoolcalls provides a middleware that patches dangling tool calls in the message history.
package patchtoolcalls
import (
"context"
"fmt"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/adk/internal"
"github.com/cloudwego/eino/schema"
)
// Config defines the configuration options for the patch tool calls middleware.
type Config struct {
// PatchedContentGenerator is an optional custom function to generate the content
// of patched tool messages. If not provided, a default message will be used.
//
// Parameters:
// - ctx: the context for the operation
// - toolName: the name of the tool that was called
// - toolCallID: the id of the tool call
//
// Returns:
// - string: the content to use for the patched tool message
// - error: any error that occurred during generation
PatchedContentGenerator func(ctx context.Context, toolName, toolCallID string) (string, error)
}
// NewTyped creates a new generic patch tool calls middleware.
//
// The middleware scans the message history before each model invocation and inserts
// placeholder tool messages for any tool calls that don't have corresponding responses.
func NewTyped[M adk.MessageType](_ context.Context, cfg *Config) (adk.TypedChatModelAgentMiddleware[M], error) {
if cfg == nil {
cfg = &Config{}
}
return &typedMiddleware[M]{
gen: cfg.PatchedContentGenerator,
}, nil
}
// New creates a new patch tool calls middleware with the given configuration.
//
// The middleware scans the message history before each model invocation and inserts
// placeholder tool messages for any tool calls that don't have corresponding responses.
func New(ctx context.Context, cfg *Config) (adk.ChatModelAgentMiddleware, error) {
return NewTyped[*schema.Message](ctx, cfg)
}
type typedMiddleware[M adk.MessageType] struct {
*adk.TypedBaseChatModelAgentMiddleware[M]
gen func(ctx context.Context, toolName, toolCallID string) (string, error)
}
func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M],
mc *adk.TypedModelContext[M],
) (context.Context, *adk.TypedChatModelAgentState[M], error) {
if len(state.Messages) == 0 {
return ctx, state, nil
}
var zero M
switch any(zero).(type) {
case *schema.Message:
return patchToolCallsForMessage(ctx, m.gen, any(state).(*adk.TypedChatModelAgentState[*schema.Message]), mc)
case *schema.AgenticMessage:
return patchToolCallsForAgenticMessage(ctx, m.gen, any(state).(*adk.TypedChatModelAgentState[*schema.AgenticMessage]), mc)
default:
panic("unreachable: unknown MessageType")
}
}
func patchToolCallsForMessage[M adk.MessageType](ctx context.Context,
gen func(ctx context.Context, toolName, toolCallID string) (string, error),
state *adk.TypedChatModelAgentState[*schema.Message],
_ *adk.TypedModelContext[M],
) (context.Context, *adk.TypedChatModelAgentState[M], error) {
patched := make([]*schema.Message, 0, len(state.Messages))
for i, msg := range state.Messages {
patched = append(patched, msg)
if msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 {
continue
}
for _, tc := range msg.ToolCalls {
if hasCorrespondingToolMessage(state.Messages[i+1:], tc.ID) {
continue
}
toolMsg, err := createPatchedToolMessage(ctx, gen, tc)
if err != nil {
return ctx, nil, err
}
patched = append(patched, toolMsg)
}
}
nState := *state
nState.Messages = patched
return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil
}
func patchToolCallsForAgenticMessage[M adk.MessageType](ctx context.Context,
gen func(ctx context.Context, toolName, toolCallID string) (string, error),
state *adk.TypedChatModelAgentState[*schema.AgenticMessage],
_ *adk.TypedModelContext[M],
) (context.Context, *adk.TypedChatModelAgentState[M], error) {
patched := make([]*schema.AgenticMessage, 0, len(state.Messages))
for i, msg := range state.Messages {
patched = append(patched, msg)
if msg.Role != schema.AgenticRoleTypeAssistant {
continue
}
// Collect tool call IDs from this assistant message.
var toolCalls []struct {
callID string
name string
}
for _, block := range msg.ContentBlocks {
if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil {
toolCalls = append(toolCalls, struct {
callID string
name string
}{callID: block.FunctionToolCall.CallID, name: block.FunctionToolCall.Name})
}
}
if len(toolCalls) == 0 {
continue
}
for _, tc := range toolCalls {
if hasCorrespondingAgenticToolResult(state.Messages[i+1:], tc.callID) {
continue
}
toolMsg, err := createPatchedAgenticToolMessage(ctx, gen, tc.name, tc.callID)
if err != nil {
return ctx, nil, err
}
patched = append(patched, toolMsg)
}
}
nState := *state
nState.Messages = patched
return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil
}
func hasCorrespondingToolMessage(messages []*schema.Message, toolCallID string) bool {
for _, msg := range messages {
// Only consider successive tool messages after the tool call message
if msg.Role != schema.Tool {
return false
}
if msg.ToolCallID == toolCallID {
return true
}
}
return false
}
func hasCorrespondingAgenticToolResult(messages []*schema.AgenticMessage, toolCallID string) bool {
for _, msg := range messages {
// Only consider successive tool messages after the tool call message
if msg.Role != schema.AgenticRoleTypeUser {
return false
}
hasToolResult := false
for _, block := range msg.ContentBlocks {
if block == nil {
continue
}
if block.Type == schema.ContentBlockTypeFunctionToolResult {
hasToolResult = true
if block.FunctionToolResult != nil && block.FunctionToolResult.CallID == toolCallID {
return true
}
}
if block.Type == schema.ContentBlockTypeToolSearchResult {
hasToolResult = true
if block.ToolSearchFunctionToolResult != nil && block.ToolSearchFunctionToolResult.CallID == toolCallID {
return true
}
}
}
if !hasToolResult {
return false
}
}
return false
}
func createPatchedToolMessage(ctx context.Context, gen func(ctx context.Context, toolName, toolCallID string) (string, error), tc schema.ToolCall) (*schema.Message, error) {
if gen != nil {
content, err := gen(ctx, tc.Function.Name, tc.ID)
if err != nil {
return nil, err
}
return schema.ToolMessage(content, tc.ID, schema.WithToolName(tc.Function.Name)), nil
}
tpl := internal.SelectPrompt(internal.I18nPrompts{
English: defaultPatchedToolMessageTemplate,
Chinese: defaultPatchedToolMessageTemplateChinese,
})
return schema.ToolMessage(fmt.Sprintf(tpl, tc.Function.Name, tc.ID), tc.ID, schema.WithToolName(tc.Function.Name)), nil
}
func createPatchedAgenticToolMessage(ctx context.Context, gen func(ctx context.Context, toolName, toolCallID string) (string, error), toolName, callID string) (*schema.AgenticMessage, error) {
var content string
if gen != nil {
var err error
content, err = gen(ctx, toolName, callID)
if err != nil {
return nil, err
}
} else {
tpl := internal.SelectPrompt(internal.I18nPrompts{
English: defaultPatchedToolMessageTemplate,
Chinese: defaultPatchedToolMessageTemplateChinese,
})
content = fmt.Sprintf(tpl, toolName, callID)
}
return &schema.AgenticMessage{
Role: schema.AgenticRoleTypeUser,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.FunctionToolResult{
CallID: callID,
Name: toolName,
Content: []*schema.FunctionToolResultContentBlock{
{Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: content}},
},
}),
},
}, nil
}
const (
defaultPatchedToolMessageTemplate = "Tool call %s with id %s was canceled - another message came in before it could be completed."
defaultPatchedToolMessageTemplateChinese = "工具调用 %s(ID 为 %s)已被取消——在其完成之前收到了另一条消息。"
)
@@ -0,0 +1,394 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package patchtoolcalls
import (
"context"
"fmt"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func TestNewTypedAgenticMessage(t *testing.T) {
ctx := context.Background()
mw, err := NewTyped[*schema.AgenticMessage](ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, mw)
var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw
}
type testToolCall struct {
ID string
Name string
Arguments string
}
func makeUserMsg[M adk.MessageType](content string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any(schema.UserMessage(content)).(M)
case *schema.AgenticMessage:
return any(schema.UserAgenticMessage(content)).(M)
}
panic("unreachable")
}
func makeAssistantMsgWithToolCalls[M adk.MessageType](content string, toolCalls []testToolCall) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
tcs := make([]schema.ToolCall, len(toolCalls))
for i, tc := range toolCalls {
tcs[i] = schema.ToolCall{ID: tc.ID, Function: schema.FunctionCall{Name: tc.Name, Arguments: tc.Arguments}}
}
return any(schema.AssistantMessage(content, tcs)).(M)
case *schema.AgenticMessage:
blocks := make([]*schema.ContentBlock, 0, len(toolCalls)+1)
if content != "" {
blocks = append(blocks, schema.NewContentBlock(&schema.AssistantGenText{Text: content}))
}
for _, tc := range toolCalls {
blocks = append(blocks, schema.NewContentBlock(&schema.FunctionToolCall{CallID: tc.ID, Name: tc.Name, Arguments: tc.Arguments}))
}
return any(&schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: blocks,
}).(M)
}
panic("unreachable")
}
func makeToolResultMsg[M adk.MessageType](content string, callID string, toolName string) M {
var zero M
switch any(zero).(type) {
case *schema.Message:
return any(schema.ToolMessage(content, callID, schema.WithToolName(toolName))).(M)
case *schema.AgenticMessage:
return any(&schema.AgenticMessage{
Role: schema.AgenticRoleTypeUser,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.FunctionToolResult{
CallID: callID,
Name: toolName,
Content: []*schema.FunctionToolResultContentBlock{
{Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: content}},
},
}),
},
}).(M)
}
panic("unreachable")
}
func assertMsgContent[M adk.MessageType](t *testing.T, msg M, expectedContent string) {
t.Helper()
switch m := any(msg).(type) {
case *schema.Message:
assert.Equal(t, expectedContent, m.Content)
case *schema.AgenticMessage:
for _, block := range m.ContentBlocks {
if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil {
for _, b := range block.FunctionToolResult.Content {
if b.Text != nil {
assert.Equal(t, expectedContent, b.Text.Text)
return
}
}
}
}
t.Errorf("no text content found in agentic message, expected %q", expectedContent)
}
}
func assertToolResultID[M adk.MessageType](t *testing.T, msg M, expectedID string) {
t.Helper()
switch m := any(msg).(type) {
case *schema.Message:
assert.Equal(t, expectedID, m.ToolCallID)
case *schema.AgenticMessage:
for _, block := range m.ContentBlocks {
if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil {
assert.Equal(t, expectedID, block.FunctionToolResult.CallID)
return
}
}
t.Errorf("no tool result found in agentic message, expected call ID %q", expectedID)
}
}
func assertToolResultName[M adk.MessageType](t *testing.T, msg M, expectedName string) {
t.Helper()
switch m := any(msg).(type) {
case *schema.Message:
assert.Equal(t, expectedName, m.ToolName)
case *schema.AgenticMessage:
for _, block := range m.ContentBlocks {
if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil {
assert.Equal(t, expectedName, block.FunctionToolResult.Name)
return
}
}
t.Errorf("no tool result found in agentic message, expected tool name %q", expectedName)
}
}
func testPatchToolCallsGeneric[M adk.MessageType](t *testing.T) {
ctx := context.Background()
tests := []struct {
name string
config *Config
messages []M
wantLen int
checkPatchedAt int // index of the patched message to check (-1 if no check needed)
wantCallID string
wantToolName string
wantContent string
}{
{
name: "empty messages",
config: nil,
messages: nil,
wantLen: 0,
checkPatchedAt: -1,
},
{
name: "no tool calls to patch",
config: nil,
messages: []M{
makeUserMsg[M]("hello"),
makeAssistantMsgWithToolCalls[M]("hi there", nil),
},
wantLen: 2,
checkPatchedAt: -1,
},
{
name: "missing tool result",
config: nil,
messages: []M{
makeUserMsg[M]("hello"),
makeAssistantMsgWithToolCalls[M]("", []testToolCall{
{ID: "call_1", Name: "tool_a", Arguments: "{}"},
{ID: "call_2", Name: "tool_b", Arguments: "{}"},
}),
makeToolResultMsg[M]("result_a", "call_1", "tool_a"),
},
wantLen: 4,
checkPatchedAt: 2,
wantCallID: "call_2",
wantToolName: "tool_b",
wantContent: fmt.Sprintf(defaultPatchedToolMessageTemplate, "tool_b", "call_2"),
},
{
name: "custom content generator",
config: &Config{
PatchedContentGenerator: func(ctx context.Context, toolName, toolCallID string) (string, error) {
return fmt.Sprintf("123 %s %s", toolName, toolCallID), nil
},
},
messages: []M{
makeUserMsg[M]("hello"),
makeAssistantMsgWithToolCalls[M]("", []testToolCall{
{ID: "call_1", Name: "tool_a", Arguments: "{}"},
{ID: "call_2", Name: "tool_b", Arguments: "{}"},
}),
makeToolResultMsg[M]("result_a", "call_1", "tool_a"),
},
wantLen: 4,
checkPatchedAt: 2,
wantCallID: "call_2",
wantToolName: "tool_b",
wantContent: "123 tool_b call_2",
},
{
name: "two consecutive assistant messages with tool calls",
config: nil,
messages: []M{
makeUserMsg[M]("hello"),
makeAssistantMsgWithToolCalls[M]("", []testToolCall{
{ID: "call_1", Name: "tool_a", Arguments: "{}"},
}),
makeAssistantMsgWithToolCalls[M]("continued...", nil),
},
wantLen: 4,
checkPatchedAt: 2,
wantCallID: "call_1",
wantToolName: "tool_a",
wantContent: fmt.Sprintf(defaultPatchedToolMessageTemplate, "tool_a", "call_1"),
},
{
name: "assistant message followed by user message without tool result",
config: nil,
messages: []M{
makeUserMsg[M]("hello"),
makeAssistantMsgWithToolCalls[M]("", []testToolCall{
{ID: "call_1", Name: "tool_a", Arguments: "{}"},
}),
makeUserMsg[M]("continued..."),
},
wantLen: 4,
checkPatchedAt: 2,
wantCallID: "call_1",
wantToolName: "tool_a",
wantContent: fmt.Sprintf(defaultPatchedToolMessageTemplate, "tool_a", "call_1"),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mw, err := NewTyped[M](ctx, tt.config)
assert.NoError(t, err)
state := &adk.TypedChatModelAgentState[M]{
Messages: tt.messages,
}
_, newState, err := mw.BeforeModelRewriteState(ctx, state, nil)
assert.NoError(t, err)
assert.Len(t, newState.Messages, tt.wantLen)
if tt.checkPatchedAt >= 0 && tt.checkPatchedAt < len(newState.Messages) {
patched := newState.Messages[tt.checkPatchedAt]
assertToolResultID(t, patched, tt.wantCallID)
assertToolResultName(t, patched, tt.wantToolName)
assertMsgContent(t, patched, tt.wantContent)
}
})
}
}
func TestPatchToolCallsGeneric(t *testing.T) {
t.Run("Message", testPatchToolCallsGeneric[*schema.Message])
t.Run("AgenticMessage", testPatchToolCallsGeneric[*schema.AgenticMessage])
}
func TestPatchToolCallsAgenticToolSearchResult(t *testing.T) {
ctx := context.Background()
mw, err := NewTyped[*schema.AgenticMessage](ctx, nil)
require.NoError(t, err)
messages := []*schema.AgenticMessage{
makeAssistantMsgWithToolCalls[*schema.AgenticMessage]("", []testToolCall{
{ID: "call_1", Name: "tool_search", Arguments: `{"query":"dynamic"}`},
}),
{
Role: schema.AgenticRoleTypeUser,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.ToolSearchFunctionToolResult{
CallID: "call_1",
Name: "tool_search",
Result: &schema.ToolSearchResult{Tools: []*schema.ToolInfo{
{Name: "dynamic_tool", Desc: "dynamic tool"},
}},
}),
},
},
}
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: messages}
_, newState, err := mw.BeforeModelRewriteState(ctx, state, nil)
require.NoError(t, err)
assert.Len(t, newState.Messages, 2)
assert.Equal(t, schema.ContentBlockTypeToolSearchResult, newState.Messages[1].ContentBlocks[0].Type)
}
// TestPatchToolCalls_NilFunctionToolCallInBlock verifies the middleware handles
// a ContentBlock with Type=FunctionToolCall but FunctionToolCall=nil without panicking.
func TestPatchToolCalls_NilFunctionToolCallInBlock(t *testing.T) {
ctx := context.Background()
mw, err := NewTyped[*schema.AgenticMessage](ctx, nil)
require.NoError(t, err)
msgs := []*schema.AgenticMessage{
schema.UserAgenticMessage("hello"),
{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
{
Type: schema.ContentBlockTypeFunctionToolCall,
FunctionToolCall: nil, // nil despite type indicating tool call
},
schema.NewContentBlock(&schema.FunctionToolCall{
CallID: "call_1",
Name: "real_tool",
}),
},
},
}
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: msgs}
_, newState, err := mw.BeforeModelRewriteState(ctx, state, nil)
assert.NoError(t, err)
assert.Len(t, newState.Messages, 3, "should patch call_1 but skip nil FunctionToolCall block")
patchMsg := newState.Messages[2]
assert.Equal(t, schema.AgenticRoleTypeUser, patchMsg.Role)
foundResult := false
for _, block := range patchMsg.ContentBlocks {
if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult &&
block.FunctionToolResult != nil && block.FunctionToolResult.CallID == "call_1" {
foundResult = true
}
}
assert.True(t, foundResult, "patched message should contain tool result for call_1")
}
// TestPatchToolCalls_AgenticMessage_NilBlockInUserMessage verifies the middleware handles
// a User Agentic Message with nil ContentBlock without panicking.
func TestPatchToolCalls_AgenticMessage_NilBlockInUserMessage(t *testing.T) {
ctx := context.Background()
mw, err := NewTyped[*schema.AgenticMessage](ctx, nil)
require.NoError(t, err)
msgs := []*schema.AgenticMessage{
schema.UserAgenticMessage("hello"),
makeAssistantMsgWithToolCalls[*schema.AgenticMessage]("", []testToolCall{
{ID: "call_1", Name: "tool_a", Arguments: "{}"},
}),
{
Role: schema.AgenticRoleTypeUser,
ContentBlocks: []*schema.ContentBlock{
nil, // nil block to test robustness
},
},
}
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: msgs}
_, newState, err := mw.BeforeModelRewriteState(ctx, state, nil)
assert.NoError(t, err, "should not panic when encountering nil block in user message")
assert.Len(t, newState.Messages, 4, "should patch call_1 and insert tool response")
// Verify the patched message is inserted at index 2
patchMsg := newState.Messages[2]
assert.Equal(t, schema.AgenticRoleTypeUser, patchMsg.Role)
foundResult := false
for _, block := range patchMsg.ContentBlocks {
if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult &&
block.FunctionToolResult != nil && block.FunctionToolResult.CallID == "call_1" {
foundResult = true
break
}
}
assert.True(t, foundResult, "patched message should contain tool result for call_1")
}
+80
View File
@@ -0,0 +1,80 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"errors"
"path/filepath"
"strings"
"sync"
fspkg "github.com/cloudwego/eino/adk/filesystem"
)
type inMemoryBackend struct {
files map[string]string
mu sync.RWMutex
}
func newInMemoryBackend() *inMemoryBackend {
return &inMemoryBackend{
files: make(map[string]string),
}
}
func (b *inMemoryBackend) LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error) {
b.mu.RLock()
defer b.mu.RUnlock()
reqPath := strings.TrimSuffix(req.Path, "/")
var result []FileInfo
for path := range b.files {
dir := filepath.Dir(path)
if dir == reqPath {
result = append(result, FileInfo{Path: path})
}
}
return result, nil
}
func (b *inMemoryBackend) Read(ctx context.Context, req *ReadRequest) (*fspkg.FileContent, error) {
b.mu.RLock()
defer b.mu.RUnlock()
content, ok := b.files[req.FilePath]
if !ok {
return nil, errors.New("file not found")
}
return &fspkg.FileContent{Content: content}, nil
}
func (b *inMemoryBackend) Write(ctx context.Context, req *WriteRequest) error {
b.mu.Lock()
defer b.mu.Unlock()
b.files[req.FilePath] = req.Content
return nil
}
func (b *inMemoryBackend) Delete(ctx context.Context, req *DeleteRequest) error {
b.mu.Lock()
defer b.mu.Unlock()
delete(b.files, req.FilePath)
return nil
}
+81
View File
@@ -0,0 +1,81 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"fmt"
"sync"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
// Config is the configuration for the tool search middleware.
type Config struct {
Backend Backend
BaseDir string
}
// NewTyped creates a new plantask middleware that provides task management tools for agents.
// It adds TaskCreate, TaskGet, TaskUpdate, and TaskList tools to the agent's tool set,
// allowing agents to create and manage structured task lists during coding sessions.
//
// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage.
func NewTyped[M adk.MessageType](_ context.Context, config *Config) (adk.TypedChatModelAgentMiddleware[M], error) {
if config == nil {
return nil, fmt.Errorf("config is required")
}
if config.Backend == nil {
return nil, fmt.Errorf("backend is required")
}
if config.BaseDir == "" {
return nil, fmt.Errorf("baseDir is required")
}
return &typedMiddleware[M]{backend: config.Backend, baseDir: config.BaseDir}, nil
}
// New creates a new plantask middleware that provides task management tools for agents.
// It adds TaskCreate, TaskGet, TaskUpdate, and TaskList tools to the agent's tool set,
// allowing agents to create and manage structured task lists during coding sessions.
func New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) {
return NewTyped[*schema.Message](ctx, config)
}
type typedMiddleware[M adk.MessageType] struct {
*adk.TypedBaseChatModelAgentMiddleware[M]
backend Backend
baseDir string
}
func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) {
if runCtx == nil {
return ctx, runCtx, nil
}
nRunCtx := *runCtx
lock := sync.Mutex{}
nRunCtx.Tools = append(nRunCtx.Tools,
newTaskCreateTool(m.backend, m.baseDir, &lock),
newTaskGetTool(m.backend, m.baseDir, &lock),
newTaskUpdateTool(m.backend, m.baseDir, &lock),
newTaskListTool(m.backend, m.baseDir, &lock),
)
return ctx, &nRunCtx, nil
}
+137
View File
@@ -0,0 +1,137 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema"
)
func TestNew(t *testing.T) {
ctx := context.Background()
_, err := New(ctx, nil)
assert.Error(t, err)
assert.Contains(t, err.Error(), "config is required")
_, err = New(ctx, &Config{})
assert.Error(t, err)
assert.Contains(t, err.Error(), "backend is required")
_, err = New(ctx, &Config{Backend: newInMemoryBackend()})
assert.Error(t, err)
assert.Contains(t, err.Error(), "baseDir is required")
m, err := New(ctx, &Config{Backend: newInMemoryBackend(), BaseDir: "/tmp/tasks"})
assert.NoError(t, err)
assert.NotNil(t, m)
}
func TestMiddlewareBeforeAgent(t *testing.T) {
ctx := context.Background()
backend := newInMemoryBackend()
baseDir := "/tmp/tasks"
m, err := New(ctx, &Config{Backend: backend, BaseDir: baseDir})
assert.NoError(t, err)
mw := m.(*typedMiddleware[*schema.Message])
ctx, runCtx, err := mw.BeforeAgent(ctx, nil)
assert.NoError(t, err)
assert.Nil(t, runCtx)
runCtx = &adk.ChatModelAgentContext{
Tools: []tool.BaseTool{},
}
ctx, newRunCtx, err := mw.BeforeAgent(ctx, runCtx)
assert.NoError(t, err)
assert.NotNil(t, newRunCtx)
assert.Len(t, newRunCtx.Tools, 4)
toolNames := make([]string, 0, 4)
for _, t := range newRunCtx.Tools {
info, _ := t.Info(ctx)
toolNames = append(toolNames, info.Name)
}
assert.Contains(t, toolNames, "TaskCreate")
assert.Contains(t, toolNames, "TaskGet")
assert.Contains(t, toolNames, "TaskUpdate")
assert.Contains(t, toolNames, "TaskList")
}
func TestIntegration(t *testing.T) {
ctx := context.Background()
backend := newInMemoryBackend()
baseDir := "/tmp/tasks"
lock := &sync.Mutex{}
createTool := newTaskCreateTool(backend, baseDir, lock)
getTool := newTaskGetTool(backend, baseDir, lock)
updateTool := newTaskUpdateTool(backend, baseDir, lock)
listTool := newTaskListTool(backend, baseDir, lock)
result, err := createTool.InvokableRun(ctx, `{"subject": "Task 1", "description": "First task"}`)
assert.NoError(t, err)
assert.Contains(t, result, "Task #1")
result, err = createTool.InvokableRun(ctx, `{"subject": "Task 2", "description": "Second task"}`)
assert.NoError(t, err)
assert.Contains(t, result, "Task #2")
_, err = updateTool.InvokableRun(ctx, `{"taskId": "2", "addBlockedBy": ["1"]}`)
assert.NoError(t, err)
result, err = listTool.InvokableRun(ctx, `{}`)
assert.NoError(t, err)
assert.Contains(t, result, "#1 [pending] Task 1")
assert.Contains(t, result, "#2 [pending] Task 2")
assert.Contains(t, result, "[blocked by #1]")
_, err = updateTool.InvokableRun(ctx, `{"taskId": "1", "status": "in_progress"}`)
assert.NoError(t, err)
result, err = getTool.InvokableRun(ctx, `{"taskId": "1"}`)
assert.NoError(t, err)
assert.Contains(t, result, "Status: in_progress")
_, err = updateTool.InvokableRun(ctx, `{"taskId": "1", "status": "completed"}`)
assert.NoError(t, err)
result, err = listTool.InvokableRun(ctx, `{}`)
assert.NoError(t, err)
assert.Contains(t, result, "#1 [completed] Task 1")
}
func TestNewTypedAgenticMessage(t *testing.T) {
ctx := context.Background()
mw, err := NewTyped[*schema.AgenticMessage](ctx, &Config{
Backend: newInMemoryBackend(),
BaseDir: "/tmp/tasks",
})
assert.NoError(t, err)
assert.NotNil(t, mw)
var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw
}
+123
View File
@@ -0,0 +1,123 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"regexp"
"github.com/cloudwego/eino/adk/middlewares/filesystem"
)
var validTaskIDRegex = regexp.MustCompile(`^\d+$`)
const highWatermarkFileName = ".highwatermark"
type task struct {
ID string `json:"id"`
Subject string `json:"subject"`
Description string `json:"description"`
Status string `json:"status"`
Blocks []string `json:"blocks"`
BlockedBy []string `json:"blockedBy"`
ActiveForm string `json:"activeForm,omitempty"`
Owner string `json:"owner,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type taskOut struct {
Result string `json:"result"`
}
const (
taskStatusPending = "pending"
taskStatusInProgress = "in_progress"
taskStatusCompleted = "completed"
taskStatusDeleted = "deleted"
)
type FileInfo = filesystem.FileInfo
type LsInfoRequest = filesystem.LsInfoRequest
type ReadRequest = filesystem.ReadRequest
type WriteRequest = filesystem.WriteRequest
type DeleteRequest struct {
FilePath string
}
// Backend defines the storage interface for task persistence.
// Implementations can use local filesystem, remote storage, or any other storage backend.
type Backend interface {
// LsInfo lists file information in the specified directory.
LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error)
// Read reads the content of a file.
Read(ctx context.Context, req *ReadRequest) (*filesystem.FileContent, error)
// Write writes content to a file, creating it if it doesn't exist.
Write(ctx context.Context, req *WriteRequest) error
// Delete removes a file from storage.
Delete(ctx context.Context, req *DeleteRequest) error
}
func isValidTaskID(taskID string) bool {
return validTaskIDRegex.MatchString(taskID)
}
func appendUnique(slice []string, items ...string) []string {
seen := make(map[string]struct{}, len(slice))
for _, s := range slice {
seen[s] = struct{}{}
}
for _, item := range items {
if _, exists := seen[item]; !exists {
slice = append(slice, item)
seen[item] = struct{}{}
}
}
return slice
}
func hasCyclicDependency(taskMap map[string]*task, blockerID, blockedID string) bool {
if blockerID == blockedID {
return true
}
visited := make(map[string]bool)
return canReach(taskMap, blockedID, blockerID, visited)
}
func canReach(taskMap map[string]*task, fromID, toID string, visited map[string]bool) bool {
if fromID == toID {
return true
}
if visited[fromID] {
return false
}
visited[fromID] = true
fromTask, exists := taskMap[fromID]
if !exists {
return false
}
for _, blockedID := range fromTask.Blocks {
if canReach(taskMap, blockedID, toID, visited) {
return true
}
}
return false
}
+265
View File
@@ -0,0 +1,265 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"fmt"
"path/filepath"
"sync"
"github.com/bytedance/sonic"
"github.com/cloudwego/eino/adk/internal"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema"
)
func newTaskCreateTool(backend Backend, baseDir string, lock *sync.Mutex) *taskCreateTool {
return &taskCreateTool{Backend: backend, BaseDir: baseDir, lock: lock}
}
type taskCreateTool struct {
Backend Backend
BaseDir string
lock *sync.Mutex
}
type taskCreateArgs struct {
Subject string `json:"subject"`
Description string `json:"description"`
ActiveForm string `json:"activeForm,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
}
func (t *taskCreateTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
desc := internal.SelectPrompt(internal.I18nPrompts{
English: taskCreateToolDesc,
Chinese: taskCreateToolDescChinese,
})
return &schema.ToolInfo{
Name: TaskCreateToolName,
Desc: desc,
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
"subject": {
Type: schema.String,
Desc: "A brief title for the task",
Required: true,
},
"description": {
Type: schema.String,
Desc: "A detailed description of what needs to be done",
Required: true,
},
"activeForm": {
Type: schema.String,
Desc: "Present continuous form shown in spinner when in_progress (e.g., \"Running tests\")",
Required: false,
},
"metadata": {
Type: schema.Object,
Desc: "Arbitrary metadata to attach to the task",
SubParams: map[string]*schema.ParameterInfo{
"propertyNames": {
Type: schema.String,
},
},
Required: false,
},
}),
}, nil
}
func (t *taskCreateTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
t.lock.Lock()
defer t.lock.Unlock()
params := &taskCreateArgs{}
err := sonic.UnmarshalString(argumentsInJSON, params)
if err != nil {
return "", err
}
files, err := t.Backend.LsInfo(ctx, &LsInfoRequest{
Path: t.BaseDir,
})
if err != nil {
return "", fmt.Errorf("%s list files in %s failed, err: %w", TaskCreateToolName, t.BaseDir, err)
}
highwatermark := int64(0)
for _, file := range files {
fileName := filepath.Base(file.Path)
if fileName == highWatermarkFileName {
content, readErr := t.Backend.Read(ctx, &ReadRequest{
FilePath: file.Path,
})
if readErr != nil {
return "", fmt.Errorf("%s read highwatermark file %s failed, err: %w", TaskCreateToolName, file.Path, readErr)
}
if content.Content != "" {
var val int64
if _, scanErr := fmt.Sscanf(content.Content, "%d", &val); scanErr == nil {
highwatermark = val
}
}
break
}
}
taskID := highwatermark + 1
taskFileName := fmt.Sprintf("%d.json", taskID)
for _, file := range files {
fileName := filepath.Base(file.Path)
if fileName == taskFileName {
return "", fmt.Errorf("task #%d already exists", taskID)
}
}
newTask := &task{
ID: fmt.Sprintf("%d", taskID),
Subject: params.Subject,
Description: params.Description,
Status: taskStatusPending,
Blocks: []string{},
BlockedBy: []string{},
ActiveForm: params.ActiveForm,
Metadata: params.Metadata,
}
taskData, err := sonic.MarshalString(newTask)
if err != nil {
return "", fmt.Errorf("%s marshal task #%d failed, err: %w", TaskCreateToolName, taskID, err)
}
// Write highwatermark file first
highwatermarkPath := filepath.Join(t.BaseDir, highWatermarkFileName)
err = t.Backend.Write(ctx, &WriteRequest{
FilePath: highwatermarkPath,
Content: fmt.Sprintf("%d", taskID),
})
if err != nil {
return "", fmt.Errorf("%s update highwatermark file %s failed, err: %w", TaskCreateToolName, highwatermarkPath, err)
}
taskFilePath := filepath.Join(t.BaseDir, taskFileName)
err = t.Backend.Write(ctx, &WriteRequest{
FilePath: taskFilePath,
Content: taskData,
})
if err != nil {
return "", fmt.Errorf("%s create Task #%d failed, err: %w", TaskCreateToolName, taskID, err)
}
resp := &taskOut{
Result: fmt.Sprintf("Task #%d created successfully: %s", taskID, params.Subject),
}
jsonResp, err := sonic.MarshalString(resp)
if err != nil {
return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskCreateToolName, err)
}
return jsonResp, nil
}
const TaskCreateToolName = "TaskCreate"
const 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.
It also helps the user understand the progress of the task and overall progress of their requests.
## When to Use This Tool
Use this tool proactively in these scenarios:
- Complex multi-step tasks - When a task requires 3 or more distinct steps or actions
- Non-trivial and complex tasks - Tasks that require careful planning or multiple operations
- Plan mode - When using plan mode, create a task list to track the work
- User explicitly requests todo list - When the user directly asks you to use the todo list
- User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated)
- After receiving new instructions - Immediately capture user requirements as tasks
- When you start working on a task - Mark it as in_progress BEFORE beginning work
- After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation
## When NOT to Use This Tool
Skip using this tool when:
- There is only a single, straightforward task
- The task is trivial and tracking it provides no organizational benefit
- The task can be completed in less than 3 trivial steps
- The task is purely conversational or informational
NOTE 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.
## Task Fields
- **subject**: A brief, actionable title in imperative form (e.g., "Fix authentication bug in login flow")
- **description**: Detailed description of what needs to be done, including context and acceptance criteria
- **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.
**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".
## Tips
- Create tasks with clear, specific subjects that describe the outcome
- Include enough detail in the description for another agent to understand and complete the task
- After creating tasks, use TaskUpdate to set up dependencies (blocks/blockedBy) if needed
- Check TaskList first to avoid creating duplicate tasks
`
const taskCreateToolDescChinese = `使用此工具为当前编码会话创建结构化的任务列表。这有助于跟踪进度、组织复杂任务,并向用户展示工作的完整性。
它还帮助用户了解任务的进度和请求的整体进展。
## 何时使用此工具
在以下场景中主动使用此工具:
- 复杂的多步骤任务 - 当任务需要 3 个或更多不同的步骤或操作时
- 非简单的复杂任务 - 需要仔细规划或多个操作的任务
- 计划模式 - 使用计划模式时,创建任务列表来跟踪工作
- 用户明确要求待办列表 - 当用户直接要求使用待办列表时
- 用户提供多个任务 - 当用户提供待办事项列表时(编号或逗号分隔)
- 收到新指令后 - 立即将用户需求记录为任务
- 开始处理任务时 - 在开始工作之前将其标记为 in_progress
- 完成任务后 - 将其标记为已完成,并添加实施过程中发现的任何后续任务
## 何时不使用此工具
在以下情况下跳过使用此工具:
- 只有一个简单直接的任务
- 任务很简单,跟踪它没有组织上的好处
- 任务可以在少于 3 个简单步骤内完成
- 任务纯粹是对话性或信息性的
注意:如果只有一个简单任务要做,不应该使用此工具。在这种情况下,直接完成任务更好。
## 任务字段
- **subject**:简短的、可操作的标题,使用祈使句形式(例如,"修复登录流程中的认证错误")
- **description**:需要完成的工作的详细描述,包括上下文和验收标准
- **activeForm**:任务处于 in_progress 状态时在加载动画中显示的现在进行时形式(例如,"正在修复认证错误")。这会在你处理任务时显示给用户。
**重要**:创建任务时始终提供 activeForm。subject 应该是祈使句("运行测试"),而 activeForm 应该是现在进行时("正在运行测试")。所有任务创建时状态为 "pending"。
## 提示
- 创建具有清晰、具体主题的任务,描述预期结果
- 在描述中包含足够的细节,以便其他代理能够理解并完成任务
- 创建任务后,如果需要,使用 TaskUpdate 设置依赖关系(blocks/blockedBy
- 先检查 TaskList 以避免创建重复任务
`
@@ -0,0 +1,91 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"path/filepath"
"sync"
"testing"
"github.com/bytedance/sonic"
"github.com/stretchr/testify/assert"
)
func TestTaskCreateTool(t *testing.T) {
ctx := context.Background()
backend := newInMemoryBackend()
baseDir := "/tmp/tasks"
lock := &sync.Mutex{}
tool := newTaskCreateTool(backend, baseDir, lock)
info, err := tool.Info(ctx)
assert.NoError(t, err)
assert.Equal(t, TaskCreateToolName, info.Name)
assert.Equal(t, taskCreateToolDesc, info.Desc)
result, err := tool.InvokableRun(ctx, `{"subject": "Test Task", "description": "Test description", "activeForm": "Testing"}`)
assert.NoError(t, err)
assert.Equal(t, `{"result":"Task #1 created successfully: Test Task"}`, result)
content, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")})
assert.NoError(t, err)
var taskData task
err = sonic.UnmarshalString(content.Content, &taskData)
assert.NoError(t, err)
assert.Equal(t, "1", taskData.ID)
assert.Equal(t, "Test Task", taskData.Subject)
assert.Equal(t, "Test description", taskData.Description)
assert.Equal(t, taskStatusPending, taskData.Status)
assert.Equal(t, "Testing", taskData.ActiveForm)
hwContent, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, highWatermarkFileName)})
assert.NoError(t, err)
assert.Equal(t, "1", hwContent.Content)
result, err = tool.InvokableRun(ctx, `{"subject": "Second Task", "description": "Second description"}`)
assert.NoError(t, err)
assert.Equal(t, `{"result":"Task #2 created successfully: Second Task"}`, result)
hwContent, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, highWatermarkFileName)})
assert.NoError(t, err)
assert.Equal(t, "2", hwContent.Content)
}
func TestTaskCreateToolWithMetadata(t *testing.T) {
ctx := context.Background()
backend := newInMemoryBackend()
baseDir := "/tmp/tasks"
lock := &sync.Mutex{}
tool := newTaskCreateTool(backend, baseDir, lock)
result, err := tool.InvokableRun(ctx, `{"subject": "Task with metadata", "description": "Has metadata", "metadata": {"key1": "value1", "key2": "value2"}}`)
assert.NoError(t, err)
assert.Contains(t, result, "Task #1 created successfully")
content, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")})
assert.NoError(t, err)
var taskData task
err = sonic.UnmarshalString(content.Content, &taskData)
assert.NoError(t, err)
assert.Equal(t, "value1", taskData.Metadata["key1"])
assert.Equal(t, "value2", taskData.Metadata["key2"])
}
+176
View File
@@ -0,0 +1,176 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"fmt"
"path/filepath"
"strings"
"sync"
"github.com/bytedance/sonic"
"github.com/cloudwego/eino/adk/internal"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema"
)
func newTaskGetTool(backend Backend, baseDir string, lock *sync.Mutex) *taskGetTool {
return &taskGetTool{Backend: backend, BaseDir: baseDir, lock: lock}
}
type taskGetTool struct {
Backend Backend
BaseDir string
lock *sync.Mutex
}
func (t *taskGetTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
desc := internal.SelectPrompt(internal.I18nPrompts{
English: taskGetToolDesc,
Chinese: taskGetToolDescChinese,
})
return &schema.ToolInfo{
Name: TaskGetToolName,
Desc: desc,
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
"taskId": {
Type: schema.String,
Desc: "The ID of the task to retrieve",
Required: true,
},
}),
}, nil
}
type taskGetArgs struct {
TaskID string `json:"taskId"`
}
func (t *taskGetTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
t.lock.Lock()
defer t.lock.Unlock()
params := &taskGetArgs{}
err := sonic.UnmarshalString(argumentsInJSON, params)
if err != nil {
return "", err
}
if !isValidTaskID(params.TaskID) {
return "", fmt.Errorf("%s validate task ID failed, err: invalid format: %s", TaskGetToolName, params.TaskID)
}
taskFileName := fmt.Sprintf("%s.json", params.TaskID)
taskFilePath := filepath.Join(t.BaseDir, taskFileName)
content, err := t.Backend.Read(ctx, &ReadRequest{
FilePath: taskFilePath,
})
if err != nil {
return "", fmt.Errorf("%s get Task #%s failed, err: %w", TaskGetToolName, params.TaskID, err)
}
taskData := &task{}
err = sonic.UnmarshalString(content.Content, taskData)
if err != nil {
return "", fmt.Errorf("%s get Task #%s failed, err: %w", TaskGetToolName, params.TaskID, err)
}
var result strings.Builder
result.WriteString(fmt.Sprintf("Task #%s: %s\n", taskData.ID, taskData.Subject))
result.WriteString(fmt.Sprintf("Status: %s\n", taskData.Status))
result.WriteString(fmt.Sprintf("Description: %s\n", taskData.Description))
if len(taskData.BlockedBy) > 0 {
blockedByIDs := make([]string, len(taskData.BlockedBy))
for i, id := range taskData.BlockedBy {
blockedByIDs[i] = "#" + id
}
result.WriteString(fmt.Sprintf("Blocked by: %s\n", strings.Join(blockedByIDs, ", ")))
}
if len(taskData.Blocks) > 0 {
blocksIDs := make([]string, len(taskData.Blocks))
for i, id := range taskData.Blocks {
blocksIDs[i] = "#" + id
}
result.WriteString(fmt.Sprintf("Blocks: %s\n", strings.Join(blocksIDs, ", ")))
}
if taskData.Owner != "" {
result.WriteString(fmt.Sprintf("Owner: %s\n", taskData.Owner))
}
resp := &taskOut{
Result: result.String(),
}
jsonResp, err := sonic.MarshalString(resp)
if err != nil {
return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskGetToolName, err)
}
return jsonResp, nil
}
const TaskGetToolName = "TaskGet"
const taskGetToolDesc = `Use this tool to retrieve a task by its ID from the task list.
## When to Use This Tool
- When you need the full description and context before starting work on a task
- To understand task dependencies (what it blocks, what blocks it)
- After being assigned a task, to get complete requirements
## Output
Returns full task details:
- **subject**: Task title
- **description**: Detailed requirements and context
- **status**: 'pending', 'in_progress', or 'completed'
- **blocks**: Tasks waiting on this one to complete
- **blockedBy**: Tasks that must complete before this one can start
## Tips
- After fetching a task, verify its blockedBy list is empty before beginning work.
- Use TaskList to see all tasks in summary form.
`
const taskGetToolDescChinese = `使用此工具通过任务 ID 从任务列表中获取任务。
## 何时使用此工具
- 当你需要在开始处理任务之前获取完整的描述和上下文时
- 了解任务依赖关系(它阻塞什么,什么阻塞它)
- 被分配任务后,获取完整的需求
## 输出
返回完整的任务详情:
- **subject**:任务标题
- **description**:详细的需求和上下文
- **status**'pending'、'in_progress' 或 'completed'
- **blocks**:等待此任务完成的任务
- **blockedBy**:必须在此任务开始之前完成的任务
## 提示
- 获取任务后,在开始工作之前验证其 blockedBy 列表是否为空。
- 使用 TaskList 查看所有任务的摘要形式。
`
+80
View File
@@ -0,0 +1,80 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"path/filepath"
"sync"
"testing"
"github.com/bytedance/sonic"
"github.com/stretchr/testify/assert"
)
func TestTaskGetTool(t *testing.T) {
ctx := context.Background()
backend := newInMemoryBackend()
baseDir := "/tmp/tasks"
lock := &sync.Mutex{}
taskData := &task{
ID: "1",
Subject: "Test Task",
Description: "Test description",
Status: taskStatusPending,
Blocks: []string{"2", "3"},
BlockedBy: []string{"4"},
}
taskJSON, _ := sonic.MarshalString(taskData)
_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: taskJSON})
tool := newTaskGetTool(backend, baseDir, lock)
info, err := tool.Info(ctx)
assert.NoError(t, err)
assert.Equal(t, TaskGetToolName, info.Name)
assert.Equal(t, taskGetToolDesc, info.Desc)
result, err := tool.InvokableRun(ctx, `{"taskId": "1"}`)
assert.NoError(t, err)
assert.Contains(t, result, "Task #1: Test Task")
assert.Contains(t, result, "Status: "+taskStatusPending)
assert.Contains(t, result, "Description: Test description")
assert.Contains(t, result, "Blocked by: #4")
assert.Contains(t, result, "Blocks: #2, #3")
_, err = tool.InvokableRun(ctx, `{"taskId": "999"}`)
assert.Error(t, err)
}
func TestTaskGetToolInvalidTaskID(t *testing.T) {
ctx := context.Background()
backend := newInMemoryBackend()
baseDir := "/tmp/tasks"
lock := &sync.Mutex{}
tool := newTaskGetTool(backend, baseDir, lock)
_, err := tool.InvokableRun(ctx, `{"taskId": "../../../etc/passwd"}`)
assert.Error(t, err)
assert.Contains(t, err.Error(), "validate task ID failed")
_, err = tool.InvokableRun(ctx, `{"taskId": "abc"}`)
assert.Error(t, err)
assert.Contains(t, err.Error(), "validate task ID failed")
}
+194
View File
@@ -0,0 +1,194 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"fmt"
"path/filepath"
"sort"
"strings"
"sync"
"github.com/bytedance/sonic"
"github.com/cloudwego/eino/adk/internal"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema"
)
func newTaskListTool(backend Backend, baseDir string, lock *sync.Mutex) *taskListTool {
return &taskListTool{Backend: backend, BaseDir: baseDir, lock: lock}
}
type taskListTool struct {
Backend Backend
BaseDir string
lock *sync.Mutex
}
func (t *taskListTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
desc := internal.SelectPrompt(internal.I18nPrompts{
English: taskListToolDesc,
Chinese: taskListToolDescChinese,
})
return &schema.ToolInfo{
Name: TaskListToolName,
Desc: desc,
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{}),
}, nil
}
func listTasks(ctx context.Context, backend Backend, baseDir string) ([]*task, error) {
files, err := backend.LsInfo(ctx, &LsInfoRequest{
Path: baseDir,
})
if err != nil {
return nil, fmt.Errorf("%s list files in %s failed, err: %w", TaskListToolName, baseDir, err)
}
var tasks []*task
for _, file := range files {
fileName := filepath.Base(file.Path)
if !strings.HasSuffix(fileName, ".json") {
continue
}
taskID := strings.TrimSuffix(fileName, ".json")
if !isValidTaskID(taskID) {
continue
}
content, err := backend.Read(ctx, &ReadRequest{
FilePath: file.Path,
})
if err != nil {
return nil, fmt.Errorf("%s read task file %s failed, err: %w", TaskListToolName, file.Path, err)
}
taskData := &task{}
err = sonic.UnmarshalString(content.Content, taskData)
if err != nil {
return nil, fmt.Errorf("%s parse task file %s failed, err: %w", TaskListToolName, file.Path, err)
}
tasks = append(tasks, taskData)
}
// sort tasks by ID
sort.Slice(tasks, func(i, j int) bool {
return tasks[i].ID < tasks[j].ID
})
return tasks, nil
}
func (t *taskListTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
t.lock.Lock()
defer t.lock.Unlock()
tasks, err := listTasks(ctx, t.Backend, t.BaseDir)
if err != nil {
return "", err
}
if len(tasks) == 0 {
resp := &taskOut{
Result: "No tasks found.",
}
jsonResp, marshalErr := sonic.MarshalString(resp)
if marshalErr != nil {
return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskListToolName, marshalErr)
}
return jsonResp, nil
}
var result strings.Builder
for i, taskData := range tasks {
if i > 0 {
result.WriteString("\n")
}
result.WriteString(fmt.Sprintf("#%s [%s] %s", taskData.ID, taskData.Status, taskData.Subject))
if taskData.Owner != "" {
result.WriteString(fmt.Sprintf(" [owner: %s]", taskData.Owner))
}
if len(taskData.BlockedBy) > 0 {
blockedByIDs := make([]string, len(taskData.BlockedBy))
for j, id := range taskData.BlockedBy {
blockedByIDs[j] = "#" + id
}
result.WriteString(fmt.Sprintf(" [blocked by %s]", strings.Join(blockedByIDs, ", ")))
}
}
resp := &taskOut{
Result: result.String(),
}
jsonResp, err := sonic.MarshalString(resp)
if err != nil {
return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskListToolName, err)
}
return jsonResp, nil
}
const TaskListToolName = "TaskList"
const taskListToolDesc = `Use this tool to list all tasks in the task list.
## When to Use This Tool
- To see what tasks are available to work on (status: 'pending', no owner, not blocked)
- To check overall progress on the project
- To find tasks that are blocked and need dependencies resolved
- After completing a task, to check for newly unblocked work or claim the next available task
- **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
## Output
Returns a summary of each task:
- **id**: Task identifier (use with TaskGet, TaskUpdate)
- **subject**: Brief description of the task
- **status**: 'pending', 'in_progress', or 'completed'
- **owner**: Agent ID if assigned, empty if available
- **blockedBy**: List of open task IDs that must be resolved first (tasks with blockedBy cannot be claimed until dependencies resolve)
Use TaskGet with a specific task ID to view full details including description and comments.
`
const taskListToolDescChinese = `使用此工具列出任务列表中的所有任务。
## 何时使用此工具
- 查看可以处理的任务(状态:'pending',无所有者,未被阻塞)
- 检查项目的整体进度
- 查找被阻塞且需要解决依赖关系的任务
- 完成任务后,检查新解除阻塞的工作或认领下一个可用任务
- **优先按 ID 顺序处理任务**(最小 ID 优先),当有多个任务可用时,因为较早的任务通常为后续任务建立上下文
## 输出
返回每个任务的摘要:
- **id**:任务标识符(与 TaskGet、TaskUpdate 一起使用)
- **subject**:任务的简要描述
- **status**'pending'、'in_progress' 或 'completed'
- **owner**:如果已分配则为代理 ID,如果可用则为空
- **blockedBy**:必须首先解决的开放任务 ID 列表(具有 blockedBy 的任务在依赖关系解决之前无法被认领)
使用 TaskGet 配合特定任务 ID 查看完整详情,包括描述和评论。
`
@@ -0,0 +1,60 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"path/filepath"
"sync"
"testing"
"github.com/bytedance/sonic"
"github.com/stretchr/testify/assert"
)
func TestTaskListTool(t *testing.T) {
ctx := context.Background()
backend := newInMemoryBackend()
baseDir := "/tmp/tasks"
lock := &sync.Mutex{}
tool := newTaskListTool(backend, baseDir, lock)
info, err := tool.Info(ctx)
assert.NoError(t, err)
assert.Equal(t, TaskListToolName, info.Name)
assert.Equal(t, taskListToolDesc, info.Desc)
result, err := tool.InvokableRun(ctx, `{}`)
assert.NoError(t, err)
assert.Equal(t, `{"result":"No tasks found."}`, result)
task1 := &task{ID: "1", Subject: "Task 1", Status: taskStatusPending, BlockedBy: []string{"2"}}
task1JSON, _ := sonic.MarshalString(task1)
_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON})
task2 := &task{ID: "2", Subject: "Task 2", Status: taskStatusInProgress, Owner: "agent1"}
task2JSON, _ := sonic.MarshalString(task2)
_ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON})
result, err = tool.InvokableRun(ctx, `{}`)
assert.NoError(t, err)
assert.Contains(t, result, "#1 ["+taskStatusPending+"] Task 1")
assert.Contains(t, result, "[blocked by #2]")
assert.Contains(t, result, "#2 ["+taskStatusInProgress+"] Task 2")
assert.Contains(t, result, "[owner: agent1]")
}
+565
View File
@@ -0,0 +1,565 @@
/*
* Copyright 2025 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package plantask
import (
"context"
"fmt"
"path/filepath"
"strings"
"sync"
"github.com/bytedance/sonic"
"github.com/cloudwego/eino/adk/internal"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema"
)
func newTaskUpdateTool(backend Backend, baseDir string, lock *sync.Mutex) *taskUpdateTool {
return &taskUpdateTool{Backend: backend, BaseDir: baseDir, lock: lock}
}
type taskUpdateTool struct {
Backend Backend
BaseDir string
lock *sync.Mutex
}
type taskUpdateArgs struct {
TaskID string `json:"taskId"`
Subject string `json:"subject,omitempty"`
Description string `json:"description,omitempty"`
ActiveForm string `json:"activeForm,omitempty"`
Status string `json:"status,omitempty"`
AddBlocks []string `json:"addBlocks,omitempty"`
AddBlockedBy []string `json:"addBlockedBy,omitempty"`
Owner string `json:"owner,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
}
func (t *taskUpdateTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
desc := internal.SelectPrompt(internal.I18nPrompts{
English: taskUpdateToolDesc,
Chinese: taskUpdateToolDescChinese,
})
return &schema.ToolInfo{
Name: TaskUpdateToolName,
Desc: desc,
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
"taskId": {
Type: schema.String,
Desc: "The ID of the task to update",
Required: true,
},
"subject": {
Type: schema.String,
Desc: "New subject for the task",
Required: false,
},
"description": {
Type: schema.String,
Desc: "New description for the task",
Required: false,
},
"activeForm": {
Type: schema.String,
Desc: "Present continuous form shown in spinner when in_progress (e.g., \"Running tests\")",
Required: false,
},
"status": {
Type: schema.String,
Desc: "New status for the task: 'pending', 'in_progress', 'completed', or 'deleted'",
Enum: []string{"pending", "in_progress", "completed", "deleted"},
Required: false,
},
"addBlocks": {
Type: schema.Array,
Desc: "Task IDs that this task blocks",
ElemInfo: &schema.ParameterInfo{Type: schema.String},
Required: false,
},
"addBlockedBy": {
Type: schema.Array,
Desc: "Task IDs that block this task",
ElemInfo: &schema.ParameterInfo{Type: schema.String},
Required: false,
},
"owner": {
Type: schema.String,
Desc: "New owner for the task",
Required: false,
},
"metadata": {
Type: schema.Object,
Desc: "Metadata keys to merge into the task. Set a key to null to delete it.",
SubParams: map[string]*schema.ParameterInfo{
"propertyNames": {
Type: schema.String,
},
},
Required: false,
},
}),
}, nil
}
func (t *taskUpdateTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
t.lock.Lock()
defer t.lock.Unlock()
params := &taskUpdateArgs{}
err := sonic.UnmarshalString(argumentsInJSON, params)
if err != nil {
return "", err
}
if !isValidTaskID(params.TaskID) {
return "", fmt.Errorf("%s validate task ID failed, err: invalid format: %s", TaskUpdateToolName, params.TaskID)
}
taskFileName := fmt.Sprintf("%s.json", params.TaskID)
taskFilePath := filepath.Join(t.BaseDir, taskFileName)
if params.Status == taskStatusDeleted {
if removeErr := t.removeTaskFromDependencies(ctx, params.TaskID); removeErr != nil {
return "", fmt.Errorf("%s remove Task #%s from dependencies failed, err: %w", TaskUpdateToolName, params.TaskID, removeErr)
}
err = t.Backend.Delete(ctx, &DeleteRequest{
FilePath: taskFilePath,
})
if err != nil {
return "", fmt.Errorf("%s delete Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err)
}
resp := &taskOut{
Result: fmt.Sprintf("Updated task #%s deleted", params.TaskID),
}
jsonResp, marshalErr := sonic.MarshalString(resp)
if marshalErr != nil {
return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskUpdateToolName, marshalErr)
}
return jsonResp, nil
}
content, err := t.Backend.Read(ctx, &ReadRequest{
FilePath: taskFilePath,
})
if err != nil {
return "", fmt.Errorf("%s read Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err)
}
taskData := &task{}
err = sonic.UnmarshalString(content.Content, taskData)
if err != nil {
return "", fmt.Errorf("%s parse Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err)
}
var updatedFields []string
if params.Subject != "" {
taskData.Subject = params.Subject
updatedFields = append(updatedFields, "subject")
}
if params.Description != "" {
taskData.Description = params.Description
updatedFields = append(updatedFields, "description")
}
if params.ActiveForm != "" {
taskData.ActiveForm = params.ActiveForm
updatedFields = append(updatedFields, "activeForm")
}
if params.Status != "" {
taskData.Status = params.Status
updatedFields = append(updatedFields, "status")
}
if len(params.AddBlocks) > 0 || len(params.AddBlockedBy) > 0 {
tasks, listErr := listTasks(ctx, t.Backend, t.BaseDir)
if listErr != nil {
return "", fmt.Errorf("%s list tasks failed, err: %w", TaskUpdateToolName, listErr)
}
taskMap := make(map[string]*task, len(tasks))
for _, tk := range tasks {
taskMap[tk.ID] = tk
}
if len(params.AddBlocks) > 0 {
for _, blockedTaskID := range params.AddBlocks {
if !isValidTaskID(blockedTaskID) {
return "", fmt.Errorf("%s validate blocked task ID failed, err: invalid format: %s", TaskUpdateToolName, blockedTaskID)
}
if hasCyclicDependency(taskMap, params.TaskID, blockedTaskID) {
return "", fmt.Errorf("%s adding Task #%s to blocks of Task #%s would create a cyclic dependency", TaskUpdateToolName, blockedTaskID, params.TaskID)
}
}
for _, blockedTaskID := range params.AddBlocks {
if addErr := t.addBlockedByToTask(ctx, blockedTaskID, params.TaskID); addErr != nil {
return "", fmt.Errorf("%s update Task #%s blocks failed, err: %w", TaskUpdateToolName, params.TaskID, addErr)
}
}
taskData.Blocks = appendUnique(taskData.Blocks, params.AddBlocks...)
updatedFields = append(updatedFields, "blocks")
}
if len(params.AddBlockedBy) > 0 {
for _, blockingTaskID := range params.AddBlockedBy {
if !isValidTaskID(blockingTaskID) {
return "", fmt.Errorf("%s validate blocking task ID failed, err: invalid format: %s", TaskUpdateToolName, blockingTaskID)
}
if hasCyclicDependency(taskMap, blockingTaskID, params.TaskID) {
return "", fmt.Errorf("%s adding Task #%s to blockedBy of Task #%s would create a cyclic dependency", TaskUpdateToolName, blockingTaskID, params.TaskID)
}
}
for _, blockingTaskID := range params.AddBlockedBy {
if addErr := t.addBlocksToTask(ctx, blockingTaskID, params.TaskID); addErr != nil {
return "", fmt.Errorf("%s update Task #%s blockedBy failed, err: %w", TaskUpdateToolName, params.TaskID, addErr)
}
}
taskData.BlockedBy = appendUnique(taskData.BlockedBy, params.AddBlockedBy...)
updatedFields = append(updatedFields, "blockedBy")
}
}
if params.Owner != "" {
taskData.Owner = params.Owner
updatedFields = append(updatedFields, "owner")
}
if params.Metadata != nil {
if taskData.Metadata == nil {
taskData.Metadata = make(map[string]any)
}
for k, v := range params.Metadata {
if v == nil {
delete(taskData.Metadata, k)
} else {
taskData.Metadata[k] = v
}
}
updatedFields = append(updatedFields, "metadata")
}
updatedContent, err := sonic.MarshalString(taskData)
if err != nil {
return "", fmt.Errorf("%s marshal Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err)
}
err = t.Backend.Write(ctx, &WriteRequest{
FilePath: taskFilePath,
Content: updatedContent,
})
if err != nil {
return "", fmt.Errorf("%s write Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err)
}
if params.Status == taskStatusCompleted {
if checkErr := t.checkIfNeedDeleteAllTasks(ctx); checkErr != nil {
return "", fmt.Errorf("%s check and delete all tasks failed, err: %w", TaskUpdateToolName, checkErr)
}
}
resp := &taskOut{
Result: fmt.Sprintf("Updated task #%s %s", params.TaskID, strings.Join(updatedFields, ", ")),
}
jsonResp, err := sonic.MarshalString(resp)
if err != nil {
return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskUpdateToolName, err)
}
return jsonResp, nil
}
func (t *taskUpdateTool) removeTaskFromDependencies(ctx context.Context, deletedTaskID string) error {
tasks, err := listTasks(ctx, t.Backend, t.BaseDir)
if err != nil {
return err
}
for _, taskData := range tasks {
if taskData.ID == deletedTaskID {
continue
}
modified := false
newBlocks := make([]string, 0, len(taskData.Blocks))
for _, id := range taskData.Blocks {
if id != deletedTaskID {
newBlocks = append(newBlocks, id)
} else {
modified = true
}
}
newBlockedBy := make([]string, 0, len(taskData.BlockedBy))
for _, id := range taskData.BlockedBy {
if id != deletedTaskID {
newBlockedBy = append(newBlockedBy, id)
} else {
modified = true
}
}
if modified {
taskData.Blocks = newBlocks
taskData.BlockedBy = newBlockedBy
updatedContent, err := sonic.MarshalString(taskData)
if err != nil {
return fmt.Errorf("failed to marshal task #%s: %w", taskData.ID, err)
}
taskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf("%s.json", taskData.ID))
if err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil {
return fmt.Errorf("failed to write task #%s: %w", taskData.ID, err)
}
}
}
return nil
}
func (t *taskUpdateTool) addBlockedByToTask(ctx context.Context, targetTaskID, blockerTaskID string) error {
taskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf("%s.json", targetTaskID))
content, err := t.Backend.Read(ctx, &ReadRequest{FilePath: taskFilePath})
if err != nil {
return fmt.Errorf("failed to read task #%s for updating blockedBy: %w", targetTaskID, err)
}
targetTask := &task{}
if unmarshalErr := sonic.UnmarshalString(content.Content, targetTask); unmarshalErr != nil {
return fmt.Errorf("failed to parse task #%s: %w", targetTaskID, unmarshalErr)
}
targetTask.BlockedBy = appendUnique(targetTask.BlockedBy, blockerTaskID)
updatedContent, err := sonic.MarshalString(targetTask)
if err != nil {
return fmt.Errorf("failed to marshal task #%s: %w", targetTaskID, err)
}
if err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil {
return fmt.Errorf("failed to write task #%s: %w", targetTaskID, err)
}
return nil
}
func (t *taskUpdateTool) addBlocksToTask(ctx context.Context, targetTaskID, blockedTaskID string) error {
taskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf("%s.json", targetTaskID))
content, err := t.Backend.Read(ctx, &ReadRequest{FilePath: taskFilePath})
if err != nil {
return fmt.Errorf("failed to read task #%s for updating blocks: %w", targetTaskID, err)
}
targetTask := &task{}
if unmarshalErr := sonic.UnmarshalString(content.Content, targetTask); unmarshalErr != nil {
return fmt.Errorf("failed to parse task #%s: %w", targetTaskID, unmarshalErr)
}
targetTask.Blocks = appendUnique(targetTask.Blocks, blockedTaskID)
updatedContent, err := sonic.MarshalString(targetTask)
if err != nil {
return fmt.Errorf("failed to marshal task #%s: %w", targetTaskID, err)
}
if err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil {
return fmt.Errorf("failed to write task #%s: %w", targetTaskID, err)
}
return nil
}
// checkIfNeedDeleteAllTasks checks if all tasks are completed, if so, it deletes all tasks
func (t *taskUpdateTool) checkIfNeedDeleteAllTasks(ctx context.Context) error {
tasks, err := listTasks(ctx, t.Backend, t.BaseDir)
if err != nil {
return err
}
for _, task := range tasks {
if task.Status != taskStatusCompleted {
return nil
}
}
for _, task := range tasks {
err := t.Backend.Delete(ctx, &DeleteRequest{
FilePath: filepath.Join(t.BaseDir, task.ID+".json"),
})
if err != nil {
return err
}
}
return nil
}
const TaskUpdateToolName = "TaskUpdate"
const taskUpdateToolDesc = `Use this tool to update a task in the task list.
## When to Use This Tool
**Mark tasks as resolved:**
- When you have completed the work described in a task
- When a task is no longer needed or has been superseded
- IMPORTANT: Always mark your assigned tasks as resolved when you finish them
- After resolving, call TaskList to find your next task
- ONLY mark a task as completed when you have FULLY accomplished it
- If you encounter errors, blockers, or cannot finish, keep the task as in_progress
- When blocked, create a new task describing what needs to be resolved
- Never mark a task as completed if:
- Tests are failing
- Implementation is partial
- You encountered unresolved errors
- You couldn't find necessary files or dependencies
**Delete tasks:**
- When a task is no longer relevant or was created in error
- Setting status to ` + "`deleted`" + ` permanently removes the task
**Update task details:**
- When requirements change or become clearer
- When establishing dependencies between tasks
## Fields You Can Update
- **status**: The task status (see Status Workflow below)
- **subject**: Change the task title (imperative form, e.g., "Run tests")
- **description**: Change the task description
- **activeForm**: Present continuous form shown in spinner when in_progress (e.g., "Running tests")
- **owner**: Change the task owner (agent name)
- **metadata**: Merge metadata keys into the task (set a key to null to delete it)
- **addBlocks**: Mark tasks that cannot start until this one completes
- **addBlockedBy**: Mark tasks that must complete before this one can start
## Status Workflow
Status progresses: ` + "`pending`" + `` + "`in_progress`" + `` + "`completed`" + `
Use ` + "`deleted`" + ` to permanently remove a task.
## Staleness
Make sure to read a task's latest state using ` + "`TaskGet`" + ` before updating it.
## Examples
Mark task as in progress when starting work:
` + "```json" + `
{"taskId": "1", "status": "in_progress"}
` + "```" + `
Mark task as completed after finishing work:
` + "```json" + `
{"taskId": "1", "status": "completed"}
` + "```" + `
Delete a task:
` + "```json" + `
{"taskId": "1", "status": "deleted"}
` + "```" + `
Claim a task by setting owner:
` + "```json" + `
{"taskId": "1", "owner": "my-name"}
` + "```" + `
Set up task dependencies:
` + "```json" + `
{"taskId": "2", "addBlockedBy": ["1"]}
` + "```" + `
`
const taskUpdateToolDescChinese = `使用此工具更新任务列表中的任务。
## 何时使用此工具
**将任务标记为已完成:**
- 当你完成了任务中描述的工作时
- 当任务不再需要或已被取代时
- 重要:完成分配给你的任务后,务必将其标记为已完成
- 完成后,调用 TaskList 查找下一个任务
- 只有在完全完成任务时才将其标记为已完成
- 如果遇到错误、阻塞或无法完成,请保持任务为 in_progress 状态
- 当被阻塞时,创建一个新任务描述需要解决的问题
- 在以下情况下不要将任务标记为已完成:
- 测试失败
- 实现不完整
- 遇到未解决的错误
- 找不到必要的文件或依赖项
**删除任务:**
- 当任务不再相关或创建错误时
- 将状态设置为 ` + "`deleted`" + ` 会永久删除任务
**更新任务详情:**
- 当需求变更或变得更清晰时
- 当建立任务之间的依赖关系时
## 可更新的字段
- **status**:任务状态(参见下方状态流程)
- **subject**:更改任务标题(使用祈使句形式,例如"运行测试")
- **description**:更改任务描述
- **activeForm**in_progress 状态时在加载动画中显示的现在进行时形式(例如"正在运行测试")
- **owner**:更改任务所有者(代理名称)
- **metadata**:将元数据键合并到任务中(将键设置为 null 可删除它)
- **addBlocks**:标记在此任务完成之前无法开始的任务
- **addBlockedBy**:标记必须在此任务开始之前完成的任务
## 状态流程
状态进展:` + "`pending`" + `` + "`in_progress`" + `` + "`completed`" + `
使用 ` + "`deleted`" + ` 永久删除任务。
## 过期性
更新任务前,请确保使用 ` + "`TaskGet`" + ` 读取任务的最新状态。
## 示例
开始工作时将任务标记为进行中:
` + "```json" + `
{"taskId": "1", "status": "in_progress"}
` + "```" + `
完成工作后将任务标记为已完成:
` + "```json" + `
{"taskId": "1", "status": "completed"}
` + "```" + `
删除任务:
` + "```json" + `
{"taskId": "1", "status": "deleted"}
` + "```" + `
通过设置 owner 认领任务:
` + "```json" + `
{"taskId": "1", "owner": "my-name"}
` + "```" + `
设置任务依赖关系:
` + "```json" + `
{"taskId": "2", "addBlockedBy": ["1"]}
` + "```" + `
`

Some files were not shown because too many files have changed in this diff Show More