Files
ragflow/internal/agent/dsl/extract_test.go
Zhichang Yu dfe2dc346d feat[Go]: port agent attachment download, chatbot + agentbot completion/info endpoints from Python (#16405)
## Summary

Ports five Python agent APIs to Go under the v1 Gin router:

- `GET  /api/v1/agents/attachments/<attachment_id>/download`
- `POST /api/v1/chatbots/<dialog_id>/completions`  (SSE)
- `GET  /api/v1/chatbots/<dialog_id>/info`
- `POST /api/v1/agentbots/<agent_id>/completions` (SSE)
- `GET  /api/v1/agentbots/<agent_id>/inputs`

Mirrors the existing Python wire shape (`{code, message,
data:{answer,reference,...}}` per Python `canvas_service.completion`) so
the iframe SDK and existing JS widgets keep working.

## Behavioural parity with Python

| # | Concern | How it's met |
|---|---------|--------------|
| R0 | Bot routes must not require regular user session | Routes mount
on `apiNoAuth` (router.go:198-202), with `BetaAuthMiddleware` only |
| R3 | Two SSE formats in Go drift | F2: `AgentChatCompletions` and
`AgentbotCompletion` share `service.WriteChatbotRunEvent` |
| R7 | `GetBySessionID` returns `(nil, nil)` on miss | Defensive
nil-check before `session.UserID != tenantID` |
| R8 | Begin component name vs ID | `FindBeginComponentID` resolves name
→ ID first, then `ExtractComponentInputForm(dsl, beginID)` |
| R9 | Defensive PromptConfig parsing | `stringFromMap` helper used for
`prologue` and `tavily_api_key` |
| R10 | `BetaAuthMiddleware` Bearer-prefix pre-filter | Removed —
`GetUserByToken` is called unconditionally, falls back to
`GetUserByBetaAPIToken` |
| F8 | Multi-turn chatbot history | `ChatbotCompletion` reads prior
turns from `session.Message`, appends user turn, calls LLM, persists new
pair via new `API4ConversationDAO.Update` |
| F9 | UUID gate stricter than plan | Removed — only `filepath.Base` +
CR/LF/quote header sanitization remains |
| H2 | Defence-in-depth IDOR | `AgentbotCompletion` calls `loadCanvas`
before delegating to `RunAgent` |
| M2 | SSE error leakage | `WriteChatbotFrame` emits generic `"an
internal error occurred"`; real error logged via `common.Error` |

## Verification

```bash
$ go vet ./...                                     # clean (only pre-existing issues)
$ go build ./...                                   # success
$ go test ./internal/handler/ ./internal/service/ ./internal/agent/dsl/ ./internal/common/ ./internal/dao/
ok  ragflow/internal/handler     0.617s
ok  ragflow/internal/service     1.729s
ok  ragflow/internal/agent/dsl   0.008s
ok  ragflow/internal/common      0.087s
ok  ragflow/internal/dao         0.083s
```

1199 tests pass across 5 packages.

## Known follow-ups (out of scope for this PR)

- **F1**: token-level streaming in `ChatbotCompletion` (currently emits
one frame per turn)
- **F3**: per-route `auth_types` attribute in Go (currently applied via
route group middleware)

---------

Co-authored-by: Claude <noreply@anthropic.com>
2026-06-29 09:45:16 +08:00

267 lines
7.3 KiB
Go

//
// 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.
//
package dsl
import (
"errors"
"testing"
)
// happyDSL returns a 2-component dsl suitable for the happy-path
// extractors.
func happyDSL() map[string]any {
return map[string]any{
"components": map[string]any{
"begin": map[string]any{
"obj": map[string]any{
"component_name": "Begin",
"params": map[string]any{
"mode": "Manual",
},
"input_form": map[string]any{
"query": map[string]any{
"type": "string",
},
},
},
},
"answer": map[string]any{
"obj": map[string]any{
"component_name": "Answer",
},
},
},
}
}
func TestExtractComponentInputForm_HappyPath(t *testing.T) {
got, err := ExtractComponentInputForm(happyDSL(), "begin")
if err != nil {
t.Fatalf("err = %v, want nil", err)
}
if q, ok := got["query"].(map[string]any); !ok || q["type"] != "string" {
t.Errorf("query type = %v, want string", got["query"])
}
}
func TestExtractComponentInputForm_NotFound(t *testing.T) {
_, err := ExtractComponentInputForm(happyDSL(), "missing")
if !errors.Is(err, ErrComponentNotFound) {
t.Errorf("err = %v, want ErrComponentNotFound", err)
}
}
func TestExtractComponentInputForm_MissingObj(t *testing.T) {
dsl := map[string]any{
"components": map[string]any{
"bare": map[string]any{}, // no obj
},
}
_, err := ExtractComponentInputForm(dsl, "bare")
if !errors.Is(err, ErrMalformedDSL) {
t.Errorf("err = %v, want ErrMalformedDSL", err)
}
}
func TestExtractComponentInputForm_MissingInputForm(t *testing.T) {
// "answer" has obj but no input_form.
_, err := ExtractComponentInputForm(happyDSL(), "answer")
if !errors.Is(err, ErrMissingInputForm) {
t.Errorf("err = %v, want ErrMissingInputForm", err)
}
}
func TestExtractComponentInputForm_NilDSL(t *testing.T) {
_, err := ExtractComponentInputForm(nil, "anything")
if !errors.Is(err, ErrMalformedDSL) {
t.Errorf("err = %v, want ErrMalformedDSL", err)
}
}
func TestExtractComponentParams_HappyPath(t *testing.T) {
got, err := ExtractComponentParams(happyDSL(), "begin")
if err != nil {
t.Fatalf("err = %v, want nil", err)
}
if got["mode"] != "Manual" {
t.Errorf("mode = %v, want Manual", got["mode"])
}
}
func TestExtractComponentParams_NoParams(t *testing.T) {
got, err := ExtractComponentParams(happyDSL(), "answer")
if err != nil {
t.Fatalf("err = %v, want nil (params is optional)", err)
}
if got != nil && len(got) != 0 {
t.Errorf("params = %v, want empty/nil", got)
}
}
// TestExtractComponentParams_WrongType pins that a present-but-
// wrongly-typed params field is ErrMalformedDSL. CodeRabbit PR #1.
func TestExtractComponentParams_WrongType(t *testing.T) {
dsl := map[string]any{
"components": map[string]any{
"bad": map[string]any{
"obj": map[string]any{
"component_name": "Begin",
"params": "this is a string, not a dict",
},
},
},
}
_, err := ExtractComponentParams(dsl, "bad")
if !errors.Is(err, ErrMalformedDSL) {
t.Errorf("err = %v, want ErrMalformedDSL", err)
}
}
func TestExtractComponentName_HappyPath(t *testing.T) {
got, err := ExtractComponentName(happyDSL(), "begin")
if err != nil {
t.Fatalf("err = %v, want nil", err)
}
if got != "Begin" {
t.Errorf("name = %q, want Begin", got)
}
}
func TestExtractComponentName_NotFound(t *testing.T) {
_, err := ExtractComponentName(happyDSL(), "missing")
if !errors.Is(err, ErrComponentNotFound) {
t.Errorf("err = %v, want ErrComponentNotFound", err)
}
}
// TestFindBeginComponentID_HappyPath covers the common case where the
// component ID is literally "begin" (mirrors the
// internal/agent/dsl/testdata fixtures).
func TestFindBeginComponentID_HappyPath(t *testing.T) {
id, err := FindBeginComponentID(happyDSL())
if err != nil {
t.Fatalf("err = %v, want nil", err)
}
if id != "begin" {
t.Errorf("id = %q, want begin", id)
}
}
// TestFindBeginComponentID_DifferentID ensures the helper resolves
// the name to whatever ID the canvas uses (mirrors real-world
// canvases where IDs are sally:0 / jack:0 etc.).
func TestFindBeginComponentID_DifferentID(t *testing.T) {
dsl := map[string]any{
"components": map[string]any{
"sally:0": map[string]any{
"obj": map[string]any{
"component_name": "Begin",
},
},
"jack:0": map[string]any{
"obj": map[string]any{
"component_name": "LLM",
},
},
},
}
id, err := FindBeginComponentID(dsl)
if err != nil {
t.Fatalf("err = %v, want nil", err)
}
if id != "sally:0" {
t.Errorf("id = %q, want sally:0", id)
}
}
// TestFindBeginComponentID_NotFound pins that a canvas with no begin
// component returns ErrComponentNotFound. The service layer maps this
// to an empty fallback (degrades gracefully, no panic).
func TestFindBeginComponentID_NotFound(t *testing.T) {
dsl := map[string]any{
"components": map[string]any{
"jack:0": map[string]any{
"obj": map[string]any{
"component_name": "LLM",
},
},
},
}
_, err := FindBeginComponentID(dsl)
if !errors.Is(err, ErrComponentNotFound) {
t.Errorf("err = %v, want ErrComponentNotFound", err)
}
}
// TestFindBeginComponentID_NilDSL pins that a nil dsl returns
// ErrMalformedDSL (not a nil-deref panic).
func TestFindBeginComponentID_NilDSL(t *testing.T) {
_, err := FindBeginComponentID(nil)
if !errors.Is(err, ErrMalformedDSL) {
t.Errorf("err = %v, want ErrMalformedDSL", err)
}
}
// TestExtractPrologue_HappyPath pins the prologue lookup path.
func TestExtractPrologue_HappyPath(t *testing.T) {
dsl := happyDSL()
dsl["components"].(map[string]any)["begin"].(map[string]any)["obj"].(map[string]any)["prologue"] = "hello"
got, err := ExtractPrologue(dsl)
if err != nil {
t.Fatalf("err = %v, want nil", err)
}
if got != "hello" {
t.Errorf("prologue = %q, want hello", got)
}
}
// TestExtractPrologue_NotFound pins that a missing begin component
// returns ErrComponentNotFound (the service layer turns this into
// empty-string fallback).
func TestExtractPrologue_NotFound(t *testing.T) {
_, err := ExtractPrologue(map[string]any{
"components": map[string]any{},
})
if !errors.Is(err, ErrComponentNotFound) {
t.Errorf("err = %v, want ErrComponentNotFound", err)
}
}
// TestExtractMode_HappyPath pins the mode lookup path.
func TestExtractMode_HappyPath(t *testing.T) {
dsl := happyDSL()
dsl["components"].(map[string]any)["begin"].(map[string]any)["obj"].(map[string]any)["mode"] = "Agent"
got, err := ExtractMode(dsl)
if err != nil {
t.Fatalf("err = %v, want nil", err)
}
if got != "Agent" {
t.Errorf("mode = %q, want Agent", got)
}
}
// TestExtractMode_NotFound pins that a missing begin component
// returns ErrComponentNotFound.
func TestExtractMode_NotFound(t *testing.T) {
_, err := ExtractMode(map[string]any{
"components": map[string]any{},
})
if !errors.Is(err, ErrComponentNotFound) {
t.Errorf("err = %v, want ErrComponentNotFound", err)
}
}