Files
ragflow/internal/agent/component/production_chain_fixes_test.go
Zhichang Yu f58fae5fb7 feat(go-agent): Ported retrieval node, added Keenable web search tool (#16396)
Ported retrieval node, added Keenable web search tool
- [x] New Feature (non-breaking change which adds functionality)
2026-06-29 09:45:16 +08:00

808 lines
26 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// 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.
//
// production_chain_fixes_test.go — production-wiring regression tests.
//
// Each test pins one of 4 production bugs the production-chain
// review surfaced:
//
// ExeSQL v1 DSL param mismatch → TestExeSQL_V1DSLParamsAccepted
// Retrieval kb_ids vs dataset_ids → TestRetrieval_KbIDsTranslatedToDatasetIDs
// SearchMyDataset snake_case drift → TestSearchMyDataset_AllAliasesRegistered
// LLM retry multiplicative stack → TestLLM_RetryStackingSemantics
//
// If any of these tests starts failing, the corresponding production
// fix has regressed. Do not silently re-pin — investigate the diff.
package component
import (
"context"
"fmt"
"strings"
"testing"
"github.com/glebarez/sqlite"
"gorm.io/gorm"
"ragflow/internal/agent/runtime"
agenttool "ragflow/internal/agent/tool"
"ragflow/internal/dao"
"ragflow/internal/entity"
)
type codeExecSandboxRecorder struct {
req agenttool.SandboxRequest
resp *agenttool.SandboxResponse
err error
}
func (s *codeExecSandboxRecorder) ExecuteCode(_ context.Context, req agenttool.SandboxRequest) (*agenttool.SandboxResponse, error) {
s.req = req
return s.resp, s.err
}
// TestExeSQL_V1DSLParamsAccepted exercises the v1-DSL-compat
// translator that turns v1 DSL ExeSQL params (database/username/
// host/port/password/top_n, no db_type) into the tool's required shape
// (db_type/database/username/host/port/password/max_records).
//
// Without the translator, NewExeSQLConnParams would reject the v1
// shape with "missing required connection params (db_type/host/
// database/username)" and every legacy v1 ExeSQL canvas would fail
// at buildNodeBody time. With the translator in place, the v1 shape
// compiles cleanly (db_type defaults to "mysql"; port is coerced from
// float64; top_n is mapped to max_records).
func TestExeSQL_V1DSLParamsAccepted(t *testing.T) {
t.Parallel()
// v1 shape: no db_type, port as float64 (JSON-decoded),
// top_n present, no max_records. Should compile cleanly.
v1Params := map[string]any{
"database": "demo",
"username": "root",
"host": "127.0.0.1",
"port": float64(3306), // JSON-decoded numeric is float64.
"password": "secret",
"top_n": float64(50),
}
c, err := New(componentNameExeSQL, v1Params)
if err != nil {
t.Fatalf("New(ExeSQL) with v1 DSL params errored: %v\n"+
"(this is the regression — v1 shape must be translated to tool shape)", err)
}
if c == nil {
t.Fatalf("New(ExeSQL) returned nil")
}
if got := c.Name(); got != componentNameExeSQL {
t.Errorf("ExeSQL c.Name() = %q, want %q", got, componentNameExeSQL)
}
// translateExeSQLParamsToToolShape should also be directly
// testable as a pure function: the same shape, the same result.
got := translateExeSQLParamsToToolShape(v1Params)
if got["db_type"] != "mysql" {
t.Errorf("translated db_type = %v, want %q", got["db_type"], "mysql")
}
if v, ok := got["port"].(int); !ok || v != 3306 {
t.Errorf("translated port = %v (%T), want int 3306", got["port"], got["port"])
}
if v, ok := got["max_records"].(int); !ok || v != 50 {
t.Errorf("translated max_records = %v (%T), want int 50", got["max_records"], got["max_records"])
}
if _, ok := got["top_n"]; ok {
t.Errorf("translated map should drop top_n (mapped to max_records)")
}
// Idempotency: a second pass must not double-default.
got2 := translateExeSQLParamsToToolShape(got)
if got2["db_type"] != "mysql" {
t.Errorf("idempotent db_type = %v, want %q", got2["db_type"], "mysql")
}
if v, ok := got2["port"].(int); !ok || v != 3306 {
t.Errorf("idempotent port = %v (%T), want int 3306", got2["port"], got2["port"])
}
// Explicit override wins: passing db_type=postgres must be
// preserved through the translator.
override := translateExeSQLParamsToToolShape(map[string]any{
"db_type": "postgres",
"host": "10.0.0.1",
})
if override["db_type"] != "postgres" {
t.Errorf("override db_type = %v, want %q", override["db_type"], "postgres")
}
}
// TestRetrieval_KbIDsTranslatedToDatasetIDs pins the
// v1-DSL-compat translator that maps the v1 DSL name `kb_ids`
// (the Python surface) to the tool's expected name `dataset_ids`
// (the Go tool JSON schema).
//
// Without the fix, the wrapper marshals `kb_ids` into the
// argumentsInJSON sent to RetrievalTool.InvokableRun; the tool's
// retrievalArgs struct only declares `dataset_ids`, so the
// Unmarshal silently drops the v1 field and the resulting
// RetrievalRequest.DatasetIDs is empty — the search degrades to
// "no filter", which is a behavioural regression vs the v1
// Python canvas where kb_ids was honoured.
//
// The test reads the merged map directly through applyDefaults
// (same-package access). It does not exercise the wire — the
// wire-level equivalent is covered by the existence of the
// retrievalArgs.DatasetIDs field at internal/agent/tool/retrieval.go
// and the json.Marshal of `merged` in retrievalComponent.Invoke.
func TestRetrieval_KbIDsTranslatedToDatasetIDs(t *testing.T) {
t.Parallel()
// Case 1: kb_ids from the build-time node params (c.params.KbIDs)
// are surfaced as dataset_ids in the merged map.
c, err := newRetrievalComponent(map[string]any{
"kb_ids": []any{"kb-1", "kb-2"},
})
if err != nil {
t.Fatalf("newRetrievalComponent: %v", err)
}
rc, ok := c.(*retrievalComponent)
if !ok {
t.Fatalf("wrapper is %T, want *retrievalComponent", c)
}
merged := rc.applyDefaults(map[string]any{"query": "ragflow"})
if _, hasKB := merged["kb_ids"]; hasKB {
t.Errorf("kb_ids should be removed from merged map (canonical key is dataset_ids)")
}
ds, ok := merged["dataset_ids"].([]any)
if !ok {
t.Fatalf("dataset_ids missing or wrong type: %v (%T)", merged["dataset_ids"], merged["dataset_ids"])
}
if len(ds) != 2 || ds[0] != "kb-1" || ds[1] != "kb-2" {
t.Errorf("dataset_ids = %v, want [kb-1 kb-2]", ds)
}
// Case 2: kb_ids supplied at call time (inputs map) is also
// translated.
merged2 := rc.applyDefaults(map[string]any{
"query": "ragflow",
"kb_ids": []any{"kb-3"},
})
if _, hasKB := merged2["kb_ids"]; hasKB {
t.Errorf("inputs kb_ids should be removed after translation")
}
if ds, ok := merged2["dataset_ids"].([]any); !ok || len(ds) != 1 || ds[0] != "kb-3" {
t.Errorf("inputs dataset_ids = %v, want [kb-3]", merged2["dataset_ids"])
}
// Case 3: dataset_ids already present wins over kb_ids (no
// silent overwrite). Per-call input is the canonical source.
merged3 := rc.applyDefaults(map[string]any{
"kb_ids": []any{"kb-old"},
"dataset_ids": []any{"kb-new"},
})
ds3, ok := merged3["dataset_ids"].([]any)
if !ok || len(ds3) != 1 || ds3[0] != "kb-new" {
t.Errorf("dataset_ids should keep call-time value %v, got %v", "kb-new", merged3["dataset_ids"])
}
}
func TestRetrieval_LegacyQueryStringNormalized(t *testing.T) {
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{TranslateError: true})
if err != nil {
t.Fatalf("failed to open sqlite: %v", err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatalf("failed to unwrap sql db: %v", err)
}
sqlDB.SetMaxOpenConns(1)
if err := db.AutoMigrate(&entity.Knowledgebase{}); err != nil {
t.Fatalf("failed to migrate knowledgebase: %v", err)
}
if err := db.AutoMigrate(&entity.UserTenant{}); err != nil {
t.Fatalf("failed to migrate user_tenant: %v", err)
}
origDB := dao.DB
dao.DB = db
t.Cleanup(func() { dao.DB = origDB })
activeStatus := "1"
if err := db.Create(&entity.UserTenant{
ID: "ut-1",
UserID: "user-1",
TenantID: "tenant-1",
Role: "owner",
InvitedBy: "user-1",
Status: &activeStatus,
}).Error; err != nil {
t.Fatalf("failed to seed user_tenant: %v", err)
}
if err := db.Create(&entity.Knowledgebase{
ID: "kb-da1",
Name: "da1",
TenantID: "tenant-1",
EmbdID: "BAAI/bge-m3@yy2@SILICONFLOW",
Permission: "me",
CreatedBy: "user-1",
Status: func() *string { s := string(entity.StatusValid); return &s }(),
}).Error; err != nil {
t.Fatalf("failed to seed kb: %v", err)
}
c, err := newRetrievalComponent(nil)
if err != nil {
t.Fatalf("newRetrievalComponent: %v", err)
}
rc := c.(*retrievalComponent)
merged := rc.applyDefaults(map[string]any{
"query": "UserFillUp: da1\nInput diamond necklace\n",
})
state := runtime.NewCanvasState("run-1", "task-1")
state.Sys["user_id"] = "user-1"
normalizeLegacyRetrievalInputs(runtime.WithState(context.Background(), state), merged)
if got, _ := merged["query"].(string); got != "diamond necklace" {
t.Fatalf("query = %q, want diamond necklace", got)
}
ds, ok := merged["dataset_ids"].([]string)
if !ok || len(ds) != 1 || ds[0] != "kb-da1" {
t.Fatalf("dataset_ids = %#v, want []string{\"kb-da1\"}", merged["dataset_ids"])
}
}
func TestRetrieval_StructuredUserFillInputNormalized(t *testing.T) {
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{TranslateError: true})
if err != nil {
t.Fatalf("failed to open sqlite: %v", err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatalf("failed to unwrap sql db: %v", err)
}
sqlDB.SetMaxOpenConns(1)
if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.UserTenant{}); err != nil {
t.Fatalf("failed to migrate tables: %v", err)
}
origDB := dao.DB
dao.DB = db
t.Cleanup(func() { dao.DB = origDB })
activeStatus := "1"
if err := db.Create(&entity.UserTenant{
ID: "ut-1",
UserID: "user-1",
TenantID: "tenant-1",
Role: "owner",
InvitedBy: "user-1",
Status: &activeStatus,
}).Error; err != nil {
t.Fatalf("failed to seed user_tenant: %v", err)
}
if err := db.Create(&entity.Knowledgebase{
ID: "kb-da1",
Name: "da1",
TenantID: "tenant-1",
EmbdID: "BAAI/bge-m3@yy2@SILICONFLOW",
Permission: "me",
CreatedBy: "user-1",
Status: func() *string { s := string(entity.StatusValid); return &s }(),
}).Error; err != nil {
t.Fatalf("failed to seed kb: %v", err)
}
c, err := newRetrievalComponent(nil)
if err != nil {
t.Fatalf("newRetrievalComponent: %v", err)
}
rc := c.(*retrievalComponent)
merged := rc.applyDefaults(map[string]any{
"state": map[string]any{
"UserFillUp:KBInput": map[string]any{
"kb": "da1",
"query": "合同",
},
},
})
state := runtime.NewCanvasState("run-1", "task-1")
state.Sys["user_id"] = "user-1"
normalizeLegacyRetrievalInputs(runtime.WithState(context.Background(), state), merged)
if got, _ := merged["query"].(string); got != "合同" {
t.Fatalf("query = %q, want 合同", got)
}
ds, ok := merged["dataset_ids"].([]string)
if !ok || len(ds) != 1 || ds[0] != "kb-da1" {
t.Fatalf("dataset_ids = %#v, want []string{\"kb-da1\"}", merged["dataset_ids"])
}
}
func TestRetrieval_ResolveDatasetIDByTenantName(t *testing.T) {
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{TranslateError: true})
if err != nil {
t.Fatalf("failed to open sqlite: %v", err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatalf("failed to unwrap sql db: %v", err)
}
sqlDB.SetMaxOpenConns(1)
if err := db.AutoMigrate(&entity.Knowledgebase{}); err != nil {
t.Fatalf("failed to migrate knowledgebase: %v", err)
}
origDB := dao.DB
dao.DB = db
t.Cleanup(func() { dao.DB = origDB })
if err := db.Create(&entity.Knowledgebase{
ID: "kb-da1",
Name: "da1",
TenantID: "tenant-1",
EmbdID: "BAAI/bge-m3@yy2@SILICONFLOW",
Permission: "me",
CreatedBy: "user-1",
Status: func() *string { s := string(entity.StatusValid); return &s }(),
}).Error; err != nil {
t.Fatalf("failed to seed kb: %v", err)
}
state := runtime.NewCanvasState("run-1", "task-1")
state.Sys["tenant_id"] = "tenant-1"
ctx := runtime.WithState(context.Background(), state)
if got := resolveRetrievalDatasetID(ctx, "da1"); got != "kb-da1" {
t.Fatalf("resolveRetrievalDatasetID = %q, want kb-da1", got)
}
}
func TestRetrieval_StructuredInputPreservesQueryWhenDatasetIDsAlreadyPresent(t *testing.T) {
c, err := newRetrievalComponent(nil)
if err != nil {
t.Fatalf("newRetrievalComponent: %v", err)
}
rc := c.(*retrievalComponent)
merged := rc.applyDefaults(map[string]any{
"dataset_ids": []string{"kb-fixed"},
"state": map[string]any{
"UserFillUp:KBInput": map[string]any{
"kb": "da1",
"query": "合同",
},
},
})
consumed := normalizeStructuredRetrievalInputs(context.Background(), merged)
if !consumed {
t.Fatal("normalizeStructuredRetrievalInputs should consume structured query")
}
if got, _ := merged["query"].(string); got != "合同" {
t.Fatalf("query = %q, want 合同", got)
}
ds, ok := merged["dataset_ids"].([]string)
if !ok || len(ds) != 1 || ds[0] != "kb-fixed" {
t.Fatalf("dataset_ids = %#v, want []string{\"kb-fixed\"}", merged["dataset_ids"])
}
}
// TestRetrieval_KbIDsEndToEndThroughTool is the wire-level
// companion to TestRetrieval_KbIDsTranslatedToDatasetIDs: it
// installs the simple retrieval service, builds a wrapper with
// `kb_ids` in the build-time params, and confirms that Invoke
// flows through to the tool and returns non-empty content.
//
// Split from the pure-function test above so the
// SetSimpleRetrievalService / t.Cleanup dance doesn't race with
// the case-1..3 pure-function calls under t.Parallel().
//
// Not t.Parallel(): the global retrievalServiceImpl is shared
// across the package, and several other tests in this file and
// in retrieval_swap_test.go install/restore the simple service
// via t.Cleanup. Running in parallel makes the Cleanup chain
// racy — a peer's Cleanup can restore the stub between this
// test's SetSimpleRetrievalService and its Invoke. Marking this
// test serial avoids the race entirely.
func TestRetrieval_KbIDsEndToEndThroughTool(t *testing.T) {
prev := agenttool.GetRetrievalService()
agenttool.SetSimpleRetrievalService()
t.Cleanup(func() { agenttool.SetRetrievalService(prev) })
wrapper, err := New(componentNameRetrieval, map[string]any{
"kb_ids": []any{"kb-1"},
"top_n": 3,
})
if err != nil {
t.Fatalf("New(Retrieval): %v", err)
}
out, err := wrapper.Invoke(context.Background(), map[string]any{"query": "ragflow"})
if err != nil {
t.Fatalf("Retrieval.Invoke: %v", err)
}
if fc, _ := out["formalized_content"].(string); fc == "" {
t.Errorf("Retrieval.Invoke returned empty formalized_content (translation regression?)")
}
}
// TestSearchMyDataset_AllAliasesRegistered pins the
// SearchMyDataset alias surface: the snake_case and Python-typo
// variants must also resolve to the Universe B tool. The Universe
// B tool registry has always
// accepted all three spellings (search_my_dataset /
// search_my_dateset) plus PascalCase; Universe A now mirrors that
// surface so older DSLs don't fail with "unknown component" at
// buildNodeBody time.
//
// All four names must resolve to a registered factory; the
// wrappers' Name() returns the canonical "Retrieval" but the
// alias lookup succeeds regardless of which spelling the DSL used.
func TestSearchMyDataset_AllAliasesRegistered(t *testing.T) {
t.Parallel()
for _, name := range []string{
"Retrieval",
"SearchMyDataset",
"search_my_dataset",
"search_my_dateset",
} {
t.Run(name, func(t *testing.T) {
c, err := New(name, nil)
if err != nil {
t.Fatalf("New(%q) errored: %v", name, err)
}
if c == nil {
t.Fatalf("New(%q) returned nil", name)
}
// The canonical name is "Retrieval" — the aliases all
// resolve to the same wrapper.
if got := c.Name(); got != componentNameRetrieval {
t.Errorf("New(%q).Name() = %q, want %q", name, got, componentNameRetrieval)
}
})
}
// Also assert the registration list (used by introspection
// tooling) contains every alias.
have := map[string]bool{}
for _, n := range RegisteredNames() {
have[strings.ToLower(n)] = true
}
for _, expected := range []string{"retrieval", "searchmydataset", "search_my_dataset", "search_my_dateset"} {
if !have[expected] {
t.Errorf("RegisteredNames() missing %q (search-mydataset alias surface regression)", expected)
}
}
}
// TestCodeExec_LegacyDSLWrapperRegistered pins the Universe A
// registration for the legacy v1 DSL node label `CodeExec`.
// Without this wrapper, DSLs like internal/agent/dsl/testdata/all.json
// fail at buildNodeBody time with "unknown component".
func TestCodeExec_LegacyDSLWrapperRegistered(t *testing.T) {
t.Parallel()
c, err := New(componentNameCodeExec, map[string]any{
"lang": "python",
"script": "def main(): return 1",
})
if err != nil {
t.Fatalf("New(CodeExec) errored: %v", err)
}
if c == nil {
t.Fatal("New(CodeExec) returned nil")
}
if got := c.Name(); got != componentNameCodeExec {
t.Errorf("New(CodeExec).Name() = %q, want %q", got, componentNameCodeExec)
}
}
// TestCodeExec_LegacyDSLWrapperBridgesParamsAndOutputs verifies the
// component wrapper preserves the frontend DSL surface (`lang`,
// `script`, `arguments`) while translating the tool envelope back to
// the legacy `result` field consumed by downstream templates.
//
// Not t.Parallel(): this test swaps the package-global sandbox client.
func TestCodeExec_LegacyDSLWrapperBridgesParamsAndOutputs(t *testing.T) {
prev := agenttool.GetSandboxClient()
recorder := &codeExecSandboxRecorder{
resp: &agenttool.SandboxResponse{
ExitCode: 0,
Stdout: "ok",
StructuredResult: map[string]any{
"present": true,
"value": float64(14),
},
},
}
agenttool.SetSandboxClient(recorder)
t.Cleanup(func() { agenttool.SetSandboxClient(prev) })
c, err := New(componentNameCodeExec, map[string]any{
"lang": "python",
"script": "def main(x):\n" +
" return int(x) * 2\n",
"arguments": map[string]any{
"x": "from-params",
},
})
if err != nil {
t.Fatalf("New(CodeExec): %v", err)
}
out, err := c.Invoke(context.Background(), map[string]any{
"arguments": map[string]any{
"x": 7,
},
})
if err != nil {
t.Fatalf("CodeExec.Invoke: %v", err)
}
if recorder.req.Lang != "python" {
t.Errorf("sandbox lang = %q, want python", recorder.req.Lang)
}
if recorder.req.Script == "" {
t.Error("sandbox script should not be empty")
}
switch got := recorder.req.Arguments["x"].(type) {
case int:
if got != 7 {
t.Errorf("sandbox arguments[x] = %v, want 7", got)
}
case float64:
if got != 7 {
t.Errorf("sandbox arguments[x] = %v, want 7", got)
}
default:
t.Errorf("sandbox arguments[x] type = %T, want int/float64 carrying 7", got)
}
if got := out["result"]; got != float64(14) {
t.Errorf("CodeExec result = %v, want numeric 14", got)
}
if got := out["content"]; got != "14" {
t.Errorf("CodeExec content = %v, want 14", got)
}
if got := out["actual_type"]; got != "Number" {
t.Errorf("CodeExec actual_type = %v, want Number", got)
}
if got := out["_ERROR"]; got != "" {
t.Errorf("CodeExec _ERROR = %v, want empty string", got)
}
}
func TestCodeExec_LegacyDSLWrapperResolvesArgumentRefsFromState(t *testing.T) {
prev := agenttool.GetSandboxClient()
recorder := &codeExecSandboxRecorder{
resp: &agenttool.SandboxResponse{
ExitCode: 0,
StructuredResult: map[string]any{
"present": true,
"value": float64(16),
},
},
}
agenttool.SetSandboxClient(recorder)
t.Cleanup(func() { agenttool.SetSandboxClient(prev) })
c, err := New(componentNameCodeExec, map[string]any{
"lang": "python",
"script": "def main(x):\n" +
" return int(x) * 2\n",
"arguments": map[string]any{
"x": "UserFillUp:CodeInput@x",
},
"outputs": map[string]any{
"result": map[string]any{
"type": "Number",
},
},
})
if err != nil {
t.Fatalf("New(CodeExec): %v", err)
}
out, err := c.Invoke(context.Background(), map[string]any{
"state": map[string]map[string]any{
"UserFillUp:CodeInput": {
"x": "8",
},
},
})
if err != nil {
t.Fatalf("CodeExec.Invoke: %v", err)
}
if got := recorder.req.Arguments["x"]; got != "8" {
t.Fatalf("sandbox arguments[x] = %#v, want \"8\"", got)
}
if got := out["result"]; got != float64(16) {
t.Fatalf("CodeExec result = %v, want 16", got)
}
}
func TestCodeExec_LegacyDSLWrapperContractMismatchSetsError(t *testing.T) {
prev := agenttool.GetSandboxClient()
recorder := &codeExecSandboxRecorder{
resp: &agenttool.SandboxResponse{
ExitCode: 0,
StructuredResult: map[string]any{
"present": true,
"value": "not-a-number",
},
},
}
agenttool.SetSandboxClient(recorder)
t.Cleanup(func() { agenttool.SetSandboxClient(prev) })
c, err := New(componentNameCodeExec, map[string]any{
"lang": "python",
"script": "def main():\n return \"not-a-number\"\n",
"outputs": map[string]any{
"result": map[string]any{
"type": "Number",
},
},
})
if err != nil {
t.Fatalf("New(CodeExec): %v", err)
}
out, err := c.Invoke(context.Background(), map[string]any{})
if err != nil {
t.Fatalf("CodeExec.Invoke: %v", err)
}
if got := out["result"]; got != nil {
t.Errorf("CodeExec result = %v, want nil on contract mismatch", got)
}
if got := out["actual_type"]; got != "String" {
t.Errorf("CodeExec actual_type = %v, want String", got)
}
if got, _ := out["_ERROR"].(string); !strings.Contains(got, "expected type Number") {
t.Errorf("CodeExec _ERROR = %v, want contract mismatch message", out["_ERROR"])
}
if got := out["content"]; got != "not-a-number" {
t.Errorf("CodeExec content = %v, want raw canonical content", got)
}
}
func TestCodeExec_LegacyDSLWrapperPreservesExecutionError(t *testing.T) {
prev := agenttool.GetSandboxClient()
recorder := &codeExecSandboxRecorder{
resp: nil,
err: fmt.Errorf("Container pool is busy"),
}
agenttool.SetSandboxClient(recorder)
t.Cleanup(func() { agenttool.SetSandboxClient(prev) })
c, err := New(componentNameCodeExec, map[string]any{
"lang": "python",
"script": "def main(): return 16\n",
"outputs": map[string]any{
"result": map[string]any{
"type": "Number",
},
},
})
if err != nil {
t.Fatalf("New(CodeExec): %v", err)
}
out, err := c.Invoke(context.Background(), map[string]any{})
if err == nil {
t.Fatal("CodeExec.Invoke: want wrapped execution error, got nil")
}
if got, _ := out["_ERROR"].(string); got != "Container pool is busy" {
t.Errorf("CodeExec _ERROR = %v, want sandbox execution error", out["_ERROR"])
}
if got := out["result"]; got != nil {
t.Errorf("CodeExec result = %v, want nil on execution error", got)
}
}
// countingInvoker records every call and returns an error every
// time. Used by TestLLM_RetryStackingSemantics to assert the total
// number of invocations the LLM component makes.
type countingInvoker struct {
calls int
}
func (c *countingInvoker) Invoke(_ context.Context, _ ChatInvokeRequest) (*ChatInvokeResponse, error) {
c.calls++
return nil, errLLMRetryTestAlwaysFail
}
// errLLMRetryTestAlwaysFail is a sentinel so the test can match
// without depending on a specific retryInvoker error string.
var errLLMRetryTestAlwaysFail = &fakeRetryErr{}
type fakeRetryErr struct{}
func (e *fakeRetryErr) Error() string { return "TestLLM_RetryStackingSemantics: forced failure" }
// TestLLM_RetryStackingSemantics pins the multiplicative
// retry stacking semantics (when a future change accidentally
// re-introduces the stacking bug, this test fails).
//
// Contract: when LLMParam.MaxRetries > 0, the LLM component
// re-wraps the default invoker in a fresh retryInvoker with that
// budget. Each "attempt" of the outer retryInvoker is one full
// call to the inner invoker. Because production boots wrap the
// default invoker in another retryInvoker (see
// cmd/server_main.go), the TOTAL invocations for a per-call
// MaxRetries=N setting is (N+1) × (boot_attempts+1) in production.
//
// This test runs against the test default (bare einoChatInvoker
// replaced with a counting invoker), so the test sees exactly
// MaxRetries+1 attempts. If a future change makes the LLM
// component replace the boot retry layer when MaxRetries is
// explicitly set, this test must be updated to match the new
// semantics — that change would itself warrant a new comment
// block in llm.go.
//
// Not t.Parallel(): the package-level defaultChatInvoker is
// shared across tests in this package, and several LLM tests
// (and downstream consumers like agent_test.go) swap it via
// SetDefaultChatInvoker. Serialising this test avoids races
// with peer cleanups.
func TestLLM_RetryStackingSemantics(t *testing.T) {
counter := &countingInvoker{}
prev := getDefaultChatInvoker()
SetDefaultChatInvoker(counter)
t.Cleanup(func() { SetDefaultChatInvoker(prev) })
cases := []struct {
name string
maxRetries int
delayAfterError int // nanoseconds; 0 → use default backoff
wantMinAttempts int // minimum total invoker.Invokes expected
}{
{"no override → exactly one attempt", 0, 0, 1},
{"MaxRetries=2 → three attempts", 2, 0, 3},
{"MaxRetries=5 → six attempts", 5, 0, 6},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
counter.calls = 0
comp := NewLLMComponent(LLMParam{
ModelID: "stub-model",
UserPrompt: "hi",
MaxRetries: tc.maxRetries,
DelayAfterError: 0, // use the retryInvoker default backoff (small)
})
// Use a fast backoff so the test does not block
// on the default 2s delay.
if tc.delayAfterError > 0 {
comp.param.DelayAfterError = 1 // 1ns — effectively instant
} else if tc.maxRetries > 0 {
comp.param.DelayAfterError = 1
}
// The retry chain always returns errLLMRetryTestAlwaysFail
// so we can count attempts deterministically.
_, _ = comp.Invoke(context.Background(), nil)
if counter.calls < tc.wantMinAttempts {
t.Errorf("counter.calls = %d, want >= %d (MaxRetries=%d stacking regression?)",
counter.calls, tc.wantMinAttempts, tc.maxRetries)
}
// Without an override, exactly 1 attempt; with an
// override, MaxRetries+1 attempts. (In tests the
// inner invoker has no retry chain, so the count
// is linear, not multiplicative.)
wantExact := tc.wantMinAttempts
if tc.maxRetries > 0 {
wantExact = tc.maxRetries + 1
}
if counter.calls != wantExact {
t.Errorf("counter.calls = %d, want exactly %d (MaxRetries=%d stacking regression?)",
counter.calls, wantExact, tc.maxRetries)
}
})
}
}