Enterprise-Knowledge-Base/internal/agent/observe_test.go

package agent

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	agenttools "github.com/Tencent/WeKnora/internal/agent/tools"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestAnalyzeResponse_ToolCall_DoesNotTerminate is a regression guard: the
// agent has no dedicated terminal tool — any round that requests tool calls is
// non-terminal and must keep the loop running. The agent ends only by stopping
// naturally with its answer as plain text.
func TestAnalyzeResponse_ToolCall_DoesNotTerminate(t *testing.T) {
	engine := newTestEngine(t, &mockChat{})
	resp := &types.ChatResponse{
		FinishReason: "tool_calls",
		ToolCalls: []types.LLMToolCall{
			{
				ID:   "call-1",
				Type: "function",
				Function: types.FunctionCall{
					Name:      agenttools.ToolKnowledgeSearch,
					Arguments: `{"query": "hi"}`,
				},
			},
		},
	}

	verdict := engine.analyzeResponse(
		context.Background(), resp, types.AgentStep{}, 0, "sess-1", time.Now(),
	)

	assert.False(t, verdict.isDone,
		"non-terminal tool calls must keep the loop running")
}

// TestAnalyzeResponse_NaturalStop_Terminates guards the sole termination path:
// finish_reason == "stop" with no tool calls ends the loop and surfaces the
// plain content as the final answer.
func TestAnalyzeResponse_NaturalStop_Terminates(t *testing.T) {
	engine := newTestEngine(t, &mockChat{})
	resp := &types.ChatResponse{
		FinishReason: "stop",
		Content:      "Here is the answer.",
	}

	verdict := engine.analyzeResponse(
		context.Background(), resp, types.AgentStep{}, 0, "sess-1", time.Now(),
	)

	assert.True(t, verdict.isDone, "a natural stop with no tool calls must terminate the loop")
	assert.Equal(t, "Here is the answer.", verdict.finalAnswer)
}

// TestAppendToolResults_PreservesReasoningContent verifies that the assistant
// message produced by appendToolResults carries the reasoning_content emitted
// by the model in the same round. Without this, MiMo and DeepSeek V3.2+
// thinking-mode reject the next ReAct round with HTTP 400
// "The reasoning_content in the thinking mode must be passed back to the API."
// (issue #1302).
func TestAppendToolResults_PreservesReasoningContent(t *testing.T) {
	engine := &AgentEngine{}

	t.Run("assistant message carries reasoning_content alongside thought and tool_calls", func(t *testing.T) {
		step := types.AgentStep{
			Iteration:        0,
			Thought:          "I will call search.",
			ReasoningContent: "Detailed chain of thought from MiMo/DeepSeek.",
			ToolCalls: []types.ToolCall{{
				ID:               "call_1",
				Name:             "knowledge_search",
				Args:             map[string]interface{}{"query": "hi"},
				ProviderMetadata: types.ToolCallMetadata{"google": json.RawMessage(`{"thought_signature":"gemini-thought-signature"}`)},
				Result: &types.ToolResult{
					Success: true,
					Output:  "result text",
				},
			}},
			Timestamp: time.Now(),
		}

		out := engine.appendToolResults(nil, step)

		require.Len(t, out, 2, "expect one assistant + one tool message")
		assert.Equal(t, "assistant", out[0].Role)
		assert.Equal(t, "I will call search.", out[0].Content)
		assert.Equal(t, "Detailed chain of thought from MiMo/DeepSeek.", out[0].ReasoningContent,
			"reasoning_content must be propagated to the assistant message so providers like MiMo "+
				"and DeepSeek thinking-mode see it on the next round (issue #1302)")
		require.Len(t, out[0].ToolCalls, 1)
		assert.Equal(t, "call_1", out[0].ToolCalls[0].ID)
		assert.JSONEq(t, `{"thought_signature":"gemini-thought-signature"}`,
			string(out[0].ToolCalls[0].ProviderMetadata["google"]))

		assert.Equal(t, "tool", out[1].Role)
		assert.Equal(t, "result text", out[1].Content)
	})

	t.Run("reasoning_content alone produces an assistant message", func(t *testing.T) {
		// A pure thinking emission with no visible content / tool calls is
		// unusual but legal — preserve it so the next round's request still
		// carries reasoning_content for strict providers.
		step := types.AgentStep{
			Iteration:        0,
			ReasoningContent: "reasoning only",
			Timestamp:        time.Now(),
		}

		out := engine.appendToolResults(nil, step)

		require.Len(t, out, 1)
		assert.Equal(t, "assistant", out[0].Role)
		assert.Equal(t, "reasoning only", out[0].ReasoningContent)
		assert.Empty(t, out[0].Content)
		assert.Empty(t, out[0].ToolCalls)
	})

	t.Run("step without thought/tool_calls/reasoning produces no assistant message", func(t *testing.T) {
		step := types.AgentStep{Iteration: 0, Timestamp: time.Now()}
		out := engine.appendToolResults(nil, step)
		assert.Empty(t, out, "empty steps must not inject empty assistant messages")
	})

	t.Run("appends to existing message slice", func(t *testing.T) {
		prior := []chat.Message{
			{Role: "system", Content: "sys"},
			{Role: "user", Content: "hi"},
		}
		step := types.AgentStep{
			Iteration:        1,
			Thought:          "answer",
			ReasoningContent: "thinking",
			Timestamp:        time.Now(),
		}
		out := engine.appendToolResults(prior, step)
		require.Len(t, out, 3)
		assert.Equal(t, "system", out[0].Role)
		assert.Equal(t, "user", out[1].Role)
		assert.Equal(t, "assistant", out[2].Role)
		assert.Equal(t, "thinking", out[2].ReasoningContent)
	})
}

func TestBuildRuntimeContextBlock_PinnedDocuments(t *testing.T) {
	block := buildRuntimeContextBlock(
		"sess-1",
		nil,
		[]*SelectedDocumentInfo{{
			KnowledgeID: "kid-1",
			Title:       "Report.pdf",
			FileType:    "pdf",
		}},
	)

	assert.Contains(t, block, "<pinned_documents")
	assert.Contains(t, block, `knowledge_id="kid-1"`)
	assert.Contains(t, block, `title="Report.pdf"`)
	assert.Contains(t, block, `file_type="pdf"`)
	assert.Contains(t, block, "list_knowledge_chunks")
	assert.NotContains(t, block, "<must_use>")
}

func TestBuildMustUseBlock_MCPAndSkills(t *testing.T) {
	block := buildMustUseBlock(
		[]*PinnedMCPServiceInfo{{
			ID:        "mcp-1",
			Name:      "ChemDB",
			ToolNames: []string{"mcp_chemdb_search"},
		}},
		[]*PinnedSkillInfo{{
			Name: "data-analysis",
		}},
	)

	assert.Contains(t, block, "<must_use>")
	assert.NotContains(t, block, "<runtime_context")
	assert.NotContains(t, block, "<instruction>")
	assert.Contains(t, block, "Must use MCP tools whose names start with mcp_chemdb_")
	assert.Contains(t, block, "@ChemDB")
	assert.Contains(t, block, `Must call read_skill(skill_name="data-analysis")`)
	assert.Contains(t, block, `@Skill "data-analysis"`)
}

func TestBuildMustUseBlock_MCPToolPrefixOnly(t *testing.T) {
	block := buildMustUseBlock(
		[]*PinnedMCPServiceInfo{{
			ID:        "mcp-1",
			Name:      "iwiki",
			ToolNames: []string{"mcp_iwiki_aisearchdocument", "mcp_iwiki_getdocument"},
		}},
		nil,
	)
	assert.Contains(t, block, "mcp_iwiki_")
	assert.NotContains(t, block, "aisearchdocument")
	assert.NotContains(t, block, `tools="`)
}

func TestBuildMustUseBlock_SkipsMCPWithoutTools(t *testing.T) {
	block := buildMustUseBlock(
		[]*PinnedMCPServiceInfo{{
			ID:   "mcp-1",
			Name: "DisabledMCP",
		}},
		[]*PinnedSkillInfo{{Name: "data-analysis"}},
	)
	assert.Contains(t, block, `Must call read_skill(skill_name="data-analysis")`)
	assert.NotContains(t, block, "DisabledMCP")
}

func TestRenderUserTurnContent_IncludesScopeBlocks(t *testing.T) {
	engine := &AgentEngine{
		knowledgeBasesInfo: []*KnowledgeBaseInfo{{ID: "kb-1", Name: "Docs"}},
		pinnedSkills:       []*PinnedSkillInfo{{Name: "analysis"}},
	}
	out := engine.RenderUserTurnContent("sess-1", "hello")
	assert.Contains(t, out, "<runtime_context")
	assert.Contains(t, out, "<must_use>")
	assert.Contains(t, out, "hello")
}