WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,82 @@
|
||||
// Package mcp wires the curated weknora tool set to an
|
||||
// modelcontextprotocol/go-sdk server. RunStdio is the entry point invoked
|
||||
// by `weknora mcp serve`.
|
||||
//
|
||||
// Design notes:
|
||||
//
|
||||
// - Tool surface is hand-curated rather than auto-derived from the cobra
|
||||
// tree (which would expose auth/link/completion/destructive verbs that
|
||||
// don't belong on an agent-callable surface).
|
||||
// - Long-running tools (chat / session_ask) accumulate the LLM SSE
|
||||
// stream server-side and return a single CallToolResult - MCP spec
|
||||
// 2025-06-18 does not define streamed tool-result content, so the
|
||||
// accumulate-and-return pattern is the canonical path.
|
||||
// - Handlers receive ctx for cancellation; mid-LLM-stream cancellation
|
||||
// propagates to the SDK via context, which closes the SSE connection.
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/Tencent/WeKnora/cli/internal/build"
|
||||
sdk "github.com/Tencent/WeKnora/client"
|
||||
)
|
||||
|
||||
// ServiceClient bundles the SDK methods the tool registry needs. *sdk.Client
|
||||
// satisfies it; tests substitute a fake to exercise the tool handlers
|
||||
// in-process without standing up a real WeKnora server.
|
||||
//
|
||||
// Embedding the full SDK Client would couple every tool test to every SDK
|
||||
// method; declaring the narrow surface here keeps the seam tight.
|
||||
type ServiceClient interface {
|
||||
knowledgeBaseService
|
||||
knowledgeService
|
||||
chatService
|
||||
agentService
|
||||
chunkListService
|
||||
}
|
||||
|
||||
var _ ServiceClient = (*sdk.Client)(nil)
|
||||
|
||||
// RunStdio constructs the MCP server, registers the curated 10 tools, and
|
||||
// blocks reading JSON-RPC from stdin until the client disconnects or ctx
|
||||
// is cancelled. Returns the underlying transport error (if any); the cobra
|
||||
// RunE caller maps it through the usual cmdutil exit-code path.
|
||||
func RunStdio(ctx context.Context, svc ServiceClient) error {
|
||||
v, _, _ := build.Info()
|
||||
server := mcpsdk.NewServer(
|
||||
&mcpsdk.Implementation{
|
||||
Name: "weknora",
|
||||
Version: v,
|
||||
},
|
||||
nil,
|
||||
)
|
||||
registerTools(server, svc)
|
||||
if err := server.Run(ctx, &mcpsdk.StdioTransport{}); err != nil {
|
||||
// An MCP host shuts the server down by closing stdin (EOF) or cancelling
|
||||
// the context — that is normal termination, not a failure, so exit 0.
|
||||
// Only a genuine transport fault is surfaced as an error.
|
||||
if isCleanShutdown(err) {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("mcp serve: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// isCleanShutdown reports whether err is the ordinary "client went away"
|
||||
// signal: EOF on stdin or a cancelled/expired context. The go-sdk wraps the
|
||||
// stdin EOF in a formatted string ("server is closing: EOF") that does not
|
||||
// always unwrap to io.EOF, so we also match the message suffix as a fallback.
|
||||
func isCleanShutdown(err error) bool {
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return true
|
||||
}
|
||||
return strings.HasSuffix(strings.TrimSpace(err.Error()), "EOF")
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestIsCleanShutdown pins that an MCP host disconnecting (stdin EOF) or a
|
||||
// cancelled context is treated as normal termination (exit 0), while a genuine
|
||||
// transport fault is not.
|
||||
func TestIsCleanShutdown(t *testing.T) {
|
||||
clean := []error{
|
||||
io.EOF,
|
||||
context.Canceled,
|
||||
context.DeadlineExceeded,
|
||||
fmt.Errorf("jsonrpc: %w", io.EOF), // wrapped EOF
|
||||
errors.New("server is closing: EOF"), // go-sdk formatted string (no unwrap)
|
||||
}
|
||||
for _, e := range clean {
|
||||
if !isCleanShutdown(e) {
|
||||
t.Errorf("isCleanShutdown(%v) = false, want true", e)
|
||||
}
|
||||
}
|
||||
notClean := []error{
|
||||
errors.New("connection reset by peer"),
|
||||
errors.New("write: broken pipe"),
|
||||
errors.New("EOF while reading header but more expected"), // EOF not at end
|
||||
}
|
||||
for _, e := range notClean {
|
||||
if isCleanShutdown(e) {
|
||||
t.Errorf("isCleanShutdown(%v) = true, want false", e)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestSkillReferencedMCPToolsExist is part of the K6 drift guard: the MCP tool
|
||||
// names the bundled `weknora-shared` skill advertises (in its "CLI vs MCP"
|
||||
// section) must all be registered by registerTools. Guards against a skill
|
||||
// naming a renamed/removed tool (e.g. the agent_invoke→session_ask rename).
|
||||
func TestSkillReferencedMCPToolsExist(t *testing.T) {
|
||||
_, file, _, ok := runtime.Caller(0)
|
||||
if !ok {
|
||||
t.Fatal("runtime.Caller failed")
|
||||
}
|
||||
skillPath := filepath.Join(filepath.Dir(file), "..", "..", "skills", "weknora-shared", "SKILL.md")
|
||||
raw, err := os.ReadFile(skillPath)
|
||||
if err != nil {
|
||||
t.Fatalf("read skill: %v", err)
|
||||
}
|
||||
body := string(raw)
|
||||
|
||||
// Slice the "CLI vs MCP" section (heading → next "## " heading).
|
||||
start := strings.Index(body, "CLI vs MCP")
|
||||
if start < 0 {
|
||||
t.Fatal("weknora-shared SKILL.md has no 'CLI vs MCP' section to validate")
|
||||
}
|
||||
section := body[start:]
|
||||
if end := strings.Index(section, "\n## "); end >= 0 {
|
||||
section = section[:end]
|
||||
}
|
||||
|
||||
// Tool names are backtick-quoted snake_case identifiers in that section.
|
||||
tokenRe := regexp.MustCompile("`([a-z][a-z_]*)`")
|
||||
referenced := map[string]bool{}
|
||||
for _, m := range tokenRe.FindAllStringSubmatch(section, -1) {
|
||||
tok := m[1]
|
||||
if strings.Contains(tok, "_") || tok == "chat" { // tool-shaped (skip prose like `cli`)
|
||||
referenced[tok] = true
|
||||
}
|
||||
}
|
||||
if len(referenced) == 0 {
|
||||
t.Fatal("found no MCP tool names in the skill's CLI vs MCP section")
|
||||
}
|
||||
|
||||
// Live registered tool set.
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.ListTools(ctx, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("ListTools: %v", err)
|
||||
}
|
||||
registered := map[string]bool{}
|
||||
for _, tool := range res.Tools {
|
||||
registered[tool.Name] = true
|
||||
}
|
||||
|
||||
for name := range referenced {
|
||||
if !registered[name] {
|
||||
t.Errorf("weknora-shared skill references MCP tool %q which is NOT registered (renamed/removed?)", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,713 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/Tencent/WeKnora/cli/internal/cmdutil"
|
||||
"github.com/Tencent/WeKnora/cli/internal/sse"
|
||||
sdk "github.com/Tencent/WeKnora/client"
|
||||
)
|
||||
|
||||
// toolErrorResult builds an error CallToolResult with IsError=true, a
|
||||
// human-readable text fallback (Content), and the error envelope payload
|
||||
// (StructuredContent). Reuses cmdutil.ErrorToDetail so the hint / retry /
|
||||
// risk / detail fallback table stays single-source.
|
||||
//
|
||||
// Returns nil when err is nil; callers should only invoke on a real error.
|
||||
//
|
||||
// CAVEAT (go-sdk v1.6.0): SetError(err) clobbers Content; we manually
|
||||
// build CallToolResult instead.
|
||||
func toolErrorResult(err error) *mcpsdk.CallToolResult {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
detail := cmdutil.ErrorToDetail(err)
|
||||
textLine := detail.Type + ": " + detail.Message
|
||||
if detail.Hint != "" {
|
||||
textLine += "\nhint: " + detail.Hint
|
||||
}
|
||||
if len(detail.RetryArgv) != 0 {
|
||||
textLine += "\nretry: " + strings.Join(detail.RetryArgv, " ")
|
||||
}
|
||||
// StructuredContent accepts any; pass *ErrDetail directly (no round-trip).
|
||||
return &mcpsdk.CallToolResult{
|
||||
IsError: true,
|
||||
Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: textLine}},
|
||||
StructuredContent: detail,
|
||||
}
|
||||
}
|
||||
|
||||
func streamErrorDetail(sessionID, assistantMessageID string) map[string]any {
|
||||
detail := map[string]any{"session_id": sessionID}
|
||||
if assistantMessageID != "" {
|
||||
detail["assistant_message_id"] = assistantMessageID
|
||||
}
|
||||
return detail
|
||||
}
|
||||
|
||||
func toolStreamError(err *cmdutil.Error, sessionID, assistantMessageID string) *mcpsdk.CallToolResult {
|
||||
return toolErrorResult(err.WithDetail(streamErrorDetail(sessionID, assistantMessageID)))
|
||||
}
|
||||
|
||||
// successResult builds a CallToolResult with StructuredContent = payload.
|
||||
// Using Out=any in all handlers disables the SDK auto-marshal path (which
|
||||
// would overwrite our StructuredContent with a zero-struct when the handler
|
||||
// returns a typed nil). We manually populate both StructuredContent and a
|
||||
// text Content fallback so the shape is identical on success and error.
|
||||
func successResult(payload any) *mcpsdk.CallToolResult {
|
||||
return &mcpsdk.CallToolResult{
|
||||
StructuredContent: payload,
|
||||
Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: marshalToString(payload)}},
|
||||
}
|
||||
}
|
||||
|
||||
func marshalToString(v any) string {
|
||||
b, _ := json.Marshal(v)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// Narrow per-domain service interfaces. ServiceClient (server.go) embeds
|
||||
// them all; *sdk.Client satisfies the union implicitly.
|
||||
|
||||
type knowledgeBaseService interface {
|
||||
ListKnowledgeBases(ctx context.Context) ([]sdk.KnowledgeBase, error)
|
||||
GetKnowledgeBase(ctx context.Context, id string) (*sdk.KnowledgeBase, error)
|
||||
}
|
||||
|
||||
type knowledgeService interface {
|
||||
ListKnowledgeWithFilter(ctx context.Context, kbID string, page, pageSize int, filter sdk.KnowledgeListFilter) ([]sdk.Knowledge, int64, error)
|
||||
GetKnowledge(ctx context.Context, knowledgeID string) (*sdk.Knowledge, error)
|
||||
OpenKnowledgeFile(ctx context.Context, knowledgeID string) (string, io.ReadCloser, error)
|
||||
HybridSearch(
|
||||
ctx context.Context,
|
||||
kbID string,
|
||||
params *sdk.SearchParams,
|
||||
opts ...sdk.ResourceURLOptions,
|
||||
) ([]*sdk.SearchResult, error)
|
||||
}
|
||||
|
||||
type chatService interface {
|
||||
CreateSession(ctx context.Context, req *sdk.CreateSessionRequest) (*sdk.Session, error)
|
||||
KnowledgeQAStream(ctx context.Context, sessionID string, req *sdk.KnowledgeQARequest, cb func(*sdk.StreamResponse) error, opts ...sdk.ResourceURLOptions) error
|
||||
}
|
||||
|
||||
type agentService interface {
|
||||
ListAgents(ctx context.Context) ([]sdk.Agent, error)
|
||||
GetAgent(ctx context.Context, agentID string) (*sdk.Agent, error)
|
||||
AgentQAStreamWithRequest(ctx context.Context, sessionID string, req *sdk.AgentQARequest, cb sdk.AgentEventCallback, opts ...sdk.ResourceURLOptions) error
|
||||
}
|
||||
|
||||
// chunkListService is the narrow surface chunk_list depends on. Kept
|
||||
// separate from knowledgeService because the chunk subtree is its own
|
||||
// domain on the server side (/api/v1/chunks/...).
|
||||
type chunkListService interface {
|
||||
ListKnowledgeChunks(ctx context.Context, knowledgeID string, page, pageSize int, chunkTypes ...string) ([]sdk.Chunk, int64, error)
|
||||
}
|
||||
|
||||
// sessionAskService composes the two SDK methods session_ask needs
|
||||
// (CreateSession for the auto-session path + AgentQAStreamWithRequest
|
||||
// for the run itself). Declared here alongside the per-domain
|
||||
// interfaces above so ServiceClient (server.go) - which embeds the
|
||||
// four domain interfaces - also satisfies it.
|
||||
type sessionAskService interface {
|
||||
CreateSession(ctx context.Context, req *sdk.CreateSessionRequest) (*sdk.Session, error)
|
||||
AgentQAStreamWithRequest(ctx context.Context, sessionID string, req *sdk.AgentQARequest, cb sdk.AgentEventCallback, opts ...sdk.ResourceURLOptions) error
|
||||
}
|
||||
|
||||
// registerTools wires the curated 10 tools onto server. Adding a tool here
|
||||
// is a deliberate API expansion - the agent-callable surface is the
|
||||
// reason this CLI ships an MCP server, not its CLI command list, so this
|
||||
// list must be maintained by hand.
|
||||
//
|
||||
// TODO: add OutputSchema to each mcpsdk.Tool registration so agents can
|
||||
// type-check responses without structural probing. Currently omitted
|
||||
// because the Out type is `any` on all handlers (required to suppress
|
||||
// the SDK's auto-marshal which clobbers our manually-populated
|
||||
// StructuredContent). When go-sdk exposes a typed OutputSchema field
|
||||
// independent of the handler Out type, populate it from the
|
||||
// corresponding *Output struct.
|
||||
func registerTools(server *mcpsdk.Server, svc ServiceClient) {
|
||||
addKBList(server, svc)
|
||||
addKBView(server, svc)
|
||||
addDocList(server, svc)
|
||||
addDocView(server, svc)
|
||||
addDocDownload(server, svc)
|
||||
addSearchChunks(server, svc)
|
||||
addChat(server, svc)
|
||||
addAgentList(server, svc)
|
||||
addSessionAsk(server, svc)
|
||||
addChunkList(server, svc)
|
||||
}
|
||||
|
||||
// ---- kb_list -------------------------------------------------------------
|
||||
|
||||
type kbListInput struct{}
|
||||
|
||||
type kbListOutput struct {
|
||||
Items []sdk.KnowledgeBase `json:"items"`
|
||||
}
|
||||
|
||||
func addKBList(server *mcpsdk.Server, svc knowledgeBaseService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "kb_list",
|
||||
Description: "List all knowledge bases visible to the active WeKnora tenant. No arguments. Returns items[]: each item carries id, name, description, knowledge_count, is_pinned, updated_at - useful for selecting a kb_id to pass to other tools.",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "List Knowledge Bases",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: true,
|
||||
IdempotentHint: true,
|
||||
OpenWorldHint: bptr(false),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, _ kbListInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
items, err := svc.ListKnowledgeBases(ctx)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "list knowledge bases")), nil, nil
|
||||
}
|
||||
if items == nil {
|
||||
items = []sdk.KnowledgeBase{}
|
||||
}
|
||||
return successResult(kbListOutput{Items: items}), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- kb_view -------------------------------------------------------------
|
||||
|
||||
type kbViewInput struct {
|
||||
KBID string `json:"kb_id" jsonschema:"knowledge base ID"`
|
||||
}
|
||||
|
||||
func addKBView(server *mcpsdk.Server, svc knowledgeBaseService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "kb_view",
|
||||
Description: "Fetch a knowledge base by ID. Returns the full record including chunking config, embedding/summary model IDs, knowledge_count, and chunk_count.",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "View Knowledge Base",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: true,
|
||||
IdempotentHint: true,
|
||||
OpenWorldHint: bptr(false),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, in kbViewInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
if in.KBID == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeKBIDRequired, "kb_id is required")), nil, nil
|
||||
}
|
||||
kb, err := svc.GetKnowledgeBase(ctx, in.KBID)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "get knowledge base")), nil, nil
|
||||
}
|
||||
return successResult(kb), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- doc_list ------------------------------------------------------------
|
||||
|
||||
type docListInput struct {
|
||||
KBID string `json:"kb_id" jsonschema:"knowledge base ID"`
|
||||
Page int `json:"page,omitempty" jsonschema:"1-indexed page number; defaults to 1"`
|
||||
PageSize int `json:"page_size,omitempty" jsonschema:"items per page (1..1000); defaults to 20"`
|
||||
Status string `json:"status,omitempty" jsonschema:"filter by parse status: pending | processing | completed | failed"`
|
||||
Keyword string `json:"keyword,omitempty" jsonschema:"server-side substring filter (case-sensitive LIKE against title / file_name); leave empty to skip"`
|
||||
FileType string `json:"file_type,omitempty" jsonschema:"filter by file extension (e.g. pdf, md)"`
|
||||
Source string `json:"source,omitempty" jsonschema:"filter by ingestion source (e.g. api, web)"`
|
||||
TagID string `json:"tag_id,omitempty" jsonschema:"filter by tag association"`
|
||||
StartTime string `json:"start_time,omitempty" jsonschema:"include docs with updated_at >= this RFC3339 timestamp (e.g. 2006-01-02T15:04:05Z)"`
|
||||
EndTime string `json:"end_time,omitempty" jsonschema:"include docs with updated_at <= this RFC3339 timestamp (e.g. 2006-01-02T15:04:05Z)"`
|
||||
}
|
||||
|
||||
type docListOutput struct {
|
||||
Items []sdk.Knowledge `json:"items"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"page_size"`
|
||||
Total int64 `json:"total"`
|
||||
}
|
||||
|
||||
func addDocList(server *mcpsdk.Server, svc knowledgeService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "doc_list",
|
||||
Description: "List documents in a knowledge base, with pagination and optional filters (parse-status, keyword, file_type, source, tag_id, start_time/end_time on updated_at). Returns items[] with id, file_name, title, parse_status, size, updated_at - plus the page/total metadata.",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "List Documents",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: true,
|
||||
IdempotentHint: true,
|
||||
OpenWorldHint: bptr(false),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, in docListInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
if in.KBID == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeKBIDRequired, "kb_id is required")), nil, nil
|
||||
}
|
||||
page := in.Page
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
size := in.PageSize
|
||||
if size < 1 {
|
||||
size = 20
|
||||
}
|
||||
if size > 1000 {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputInvalidArgument, "page_size must be in 1..1000")), nil, nil
|
||||
}
|
||||
filter := sdk.KnowledgeListFilter{
|
||||
ParseStatus: in.Status,
|
||||
Keyword: in.Keyword,
|
||||
FileType: in.FileType,
|
||||
Source: in.Source,
|
||||
TagID: in.TagID,
|
||||
}
|
||||
if in.StartTime != "" {
|
||||
t, err := time.Parse(time.RFC3339, in.StartTime)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputInvalidArgument, fmt.Sprintf("start_time must be RFC3339 (e.g. 2006-01-02T15:04:05Z), got %q", in.StartTime))), nil, nil
|
||||
}
|
||||
filter.StartTime = t
|
||||
}
|
||||
if in.EndTime != "" {
|
||||
t, err := time.Parse(time.RFC3339, in.EndTime)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputInvalidArgument, fmt.Sprintf("end_time must be RFC3339 (e.g. 2006-01-02T15:04:05Z), got %q", in.EndTime))), nil, nil
|
||||
}
|
||||
filter.EndTime = t
|
||||
}
|
||||
items, total, err := svc.ListKnowledgeWithFilter(ctx, in.KBID, page, size, filter)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "list documents")), nil, nil
|
||||
}
|
||||
if items == nil {
|
||||
items = []sdk.Knowledge{}
|
||||
}
|
||||
return successResult(docListOutput{Items: items, Page: page, PageSize: size, Total: total}), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- doc_view ------------------------------------------------------------
|
||||
|
||||
type docViewInput struct {
|
||||
DocID string `json:"doc_id" jsonschema:"document ID (same value as the doc-id positional in CLI commands)"`
|
||||
}
|
||||
|
||||
func addDocView(server *mcpsdk.Server, svc knowledgeService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "doc_view",
|
||||
Description: "Fetch a single document by ID. Returns the Knowledge record (file_name, title, type, parse_status, size, embedding_model_id, source URL if any, etc.).",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "View Document",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: true,
|
||||
IdempotentHint: true,
|
||||
OpenWorldHint: bptr(false),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, in docViewInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
if in.DocID == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputMissingFlag, "doc_id is required")), nil, nil
|
||||
}
|
||||
k, err := svc.GetKnowledge(ctx, in.DocID)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "get knowledge")), nil, nil
|
||||
}
|
||||
return successResult(k), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- doc_download --------------------------------------------------------
|
||||
|
||||
type docDownloadInput struct {
|
||||
DocID string `json:"doc_id" jsonschema:"document ID (same value as the doc-id positional in CLI commands)"`
|
||||
}
|
||||
|
||||
type docDownloadOutput struct {
|
||||
DocID string `json:"doc_id"`
|
||||
FileName string `json:"file_name"`
|
||||
Bytes int `json:"bytes"`
|
||||
// Content is the file contents (UTF-8 if text, base64 if the SDK
|
||||
// reports a binary-looking blob). For binary, agents should decode
|
||||
// before consuming.
|
||||
Content string `json:"content"`
|
||||
IsBase64 bool `json:"is_base64"`
|
||||
}
|
||||
|
||||
// maxDocDownloadBytes caps the per-call payload to keep an agent's context
|
||||
// window safe; agents needing larger documents should chunk via doc_view +
|
||||
// search_chunks. 1 MiB matches a typical LLM context-window budget for
|
||||
// inline content (~250k tokens) while remaining cheap to serialize.
|
||||
const maxDocDownloadBytes = 1 << 20
|
||||
|
||||
func addDocDownload(server *mcpsdk.Server, svc knowledgeService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "doc_download",
|
||||
Description: "Download a document's raw bytes by ID. Capped at 1 MiB per call - for larger documents, use search_chunks to find the relevant excerpts. is_base64 reports whether content was base64-encoded (heuristic: presence of NUL byte in the first 512 bytes).",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "Download Document",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: true,
|
||||
IdempotentHint: true,
|
||||
OpenWorldHint: bptr(false),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, in docDownloadInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
if in.DocID == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputMissingFlag, "doc_id is required")), nil, nil
|
||||
}
|
||||
name, body, err := svc.OpenKnowledgeFile(ctx, in.DocID)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "open knowledge file")), nil, nil
|
||||
}
|
||||
defer body.Close()
|
||||
buf, err := io.ReadAll(io.LimitReader(body, maxDocDownloadBytes+1))
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.Wrapf(cmdutil.CodeLocalFileIO, err, "read knowledge file")), nil, nil
|
||||
}
|
||||
if len(buf) > maxDocDownloadBytes {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputInvalidArgument, fmt.Sprintf("document exceeds the %d-byte per-call cap; use search_chunks for excerpts", maxDocDownloadBytes))), nil, nil
|
||||
}
|
||||
content, isBase64 := encodeDownload(buf)
|
||||
return successResult(docDownloadOutput{
|
||||
DocID: in.DocID,
|
||||
FileName: name,
|
||||
Bytes: len(buf),
|
||||
Content: content,
|
||||
IsBase64: isBase64,
|
||||
}), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- search_chunks -------------------------------------------------------
|
||||
|
||||
type searchChunksInput struct {
|
||||
KBID string `json:"kb_id" jsonschema:"knowledge base ID to search"`
|
||||
Query string `json:"query" jsonschema:"natural-language search query"`
|
||||
Limit int `json:"limit,omitempty" jsonschema:"client-side cap on results (1..1000); defaults to 10"`
|
||||
VectorThreshold float64 `json:"vector_threshold,omitempty" jsonschema:"minimum vector similarity (0..1)"`
|
||||
KeywordThreshold float64 `json:"keyword_threshold,omitempty" jsonschema:"minimum keyword score (0..1)"`
|
||||
}
|
||||
|
||||
type searchChunksOutput struct {
|
||||
Results []*sdk.SearchResult `json:"results"`
|
||||
}
|
||||
|
||||
func addSearchChunks(server *mcpsdk.Server, svc knowledgeService) {
|
||||
// Out = any: SDK output schema would derive from searchChunksOutput,
|
||||
// which embeds *sdk.SearchResult - and SearchResult.Metadata is a
|
||||
// nilable map[string]any that violates the auto-generated
|
||||
// type=object constraint when empty. Skipping derivation by using
|
||||
// `any` keeps the structured JSON shape identical while bypassing
|
||||
// the over-eager validator. Same pattern applied to chat / session_ask
|
||||
// below.
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "search_chunks",
|
||||
Description: "Hybrid (vector + keyword) retrieval against a knowledge base. Returns the top chunks ranked by RRF; use this before chat to ground an answer in cited context. Results include knowledge_id, content, score - feed back into chat as context or display directly.",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "Search Knowledge Chunks",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: true,
|
||||
IdempotentHint: true,
|
||||
OpenWorldHint: bptr(false),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, in searchChunksInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
if in.KBID == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeKBIDRequired, "kb_id is required")), nil, nil
|
||||
}
|
||||
if strings.TrimSpace(in.Query) == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputMissingFlag, "query cannot be empty")), nil, nil
|
||||
}
|
||||
limit := in.Limit
|
||||
if limit < 1 {
|
||||
limit = 10
|
||||
}
|
||||
if limit > 1000 {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputInvalidArgument, "limit must be in 1..1000")), nil, nil
|
||||
}
|
||||
results, err := svc.HybridSearch(ctx, in.KBID, &sdk.SearchParams{
|
||||
QueryText: in.Query,
|
||||
MatchCount: limit,
|
||||
VectorThreshold: in.VectorThreshold,
|
||||
KeywordThreshold: in.KeywordThreshold,
|
||||
})
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "hybrid search")), nil, nil
|
||||
}
|
||||
if len(results) > limit {
|
||||
results = results[:limit]
|
||||
}
|
||||
if results == nil {
|
||||
results = []*sdk.SearchResult{}
|
||||
}
|
||||
return successResult(searchChunksOutput{Results: results}), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- chat ----------------------------------------------------------------
|
||||
|
||||
type chatInput struct {
|
||||
KBID string `json:"kb_id" jsonschema:"knowledge base ID to chat against"`
|
||||
Query string `json:"query" jsonschema:"user query"`
|
||||
SessionID string `json:"session_id,omitempty" jsonschema:"existing session to continue; auto-created when empty"`
|
||||
Reference bool `json:"reference,omitempty" jsonschema:"include indexed references"`
|
||||
Verbose bool `json:"verbose,omitempty" jsonschema:"include reasoning, tools, and lifecycle events"`
|
||||
}
|
||||
|
||||
type chatOutput struct {
|
||||
Events []sse.ProjectedEvent `json:"events"`
|
||||
SessionID string `json:"session_id"`
|
||||
AssistantMessageID string `json:"assistant_message_id,omitempty"`
|
||||
KBID string `json:"kb_id"`
|
||||
Query string `json:"query"`
|
||||
}
|
||||
|
||||
func addChat(server *mcpsdk.Server, svc chatService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "chat",
|
||||
Description: "Run a RAG answer against a knowledge base. Returns a bounded answer-event projection by default; reference=true adds indexed citations and verbose=true adds reasoning, tool, and lifecycle events. MCP tools/call is buffered rather than streaming. Pass session_id to continue a conversation.",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "Chat with KB (Streaming RAG)",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: false,
|
||||
IdempotentHint: false,
|
||||
OpenWorldHint: bptr(true),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, in chatInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
if in.KBID == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeKBIDRequired, "kb_id is required")), nil, nil
|
||||
}
|
||||
if strings.TrimSpace(in.Query) == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputMissingFlag, "query cannot be empty")), nil, nil
|
||||
}
|
||||
sessionID := in.SessionID
|
||||
if sessionID == "" {
|
||||
sess, err := svc.CreateSession(ctx, &sdk.CreateSessionRequest{Title: "weknora mcp chat"})
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "create chat session")), nil, nil
|
||||
}
|
||||
sessionID = sess.ID
|
||||
}
|
||||
req := &sdk.KnowledgeQARequest{
|
||||
Query: in.Query,
|
||||
KnowledgeBaseIDs: []string{in.KBID},
|
||||
AgentEnabled: false,
|
||||
Channel: "api",
|
||||
}
|
||||
projector := sse.NewProjector(in.Verbose, in.Reference, in.KBID)
|
||||
events := make([]sse.ProjectedEvent, 0)
|
||||
streamErr := svc.KnowledgeQAStream(ctx, sessionID, req, func(r *sdk.StreamResponse) error {
|
||||
if event, include := projector.Chat(r); include {
|
||||
events = append(events, event)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if streamErr != nil {
|
||||
return toolStreamError(cmdutil.WrapStream(streamErr, "knowledge qa stream"), sessionID, projector.AssistantMessageID()), nil, nil
|
||||
}
|
||||
if !projector.Done() {
|
||||
return toolStreamError(cmdutil.NewError(cmdutil.CodeSSEStreamAborted, "stream ended without a terminal event"), sessionID, projector.AssistantMessageID()), nil, nil
|
||||
}
|
||||
sid := projector.SessionID()
|
||||
if sid == "" {
|
||||
sid = sessionID
|
||||
}
|
||||
return successResult(chatOutput{
|
||||
Events: events,
|
||||
SessionID: sid,
|
||||
AssistantMessageID: projector.AssistantMessageID(),
|
||||
KBID: in.KBID,
|
||||
Query: in.Query,
|
||||
}), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- agent_list ----------------------------------------------------------
|
||||
|
||||
type agentListInput struct{}
|
||||
|
||||
type agentListOutput struct {
|
||||
Items []sdk.Agent `json:"items"`
|
||||
}
|
||||
|
||||
func addAgentList(server *mcpsdk.Server, svc agentService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "agent_list",
|
||||
Description: "List the tenant's custom agents. Returns items[] with id, name, description, is_builtin - use to discover an agent_id before session_ask.",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "List Custom Agents",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: true,
|
||||
IdempotentHint: true,
|
||||
OpenWorldHint: bptr(false),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, _ agentListInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
items, err := svc.ListAgents(ctx)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "list agents")), nil, nil
|
||||
}
|
||||
if items == nil {
|
||||
items = []sdk.Agent{}
|
||||
}
|
||||
return successResult(agentListOutput{Items: items}), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- session_ask ---------------------------------------------------------
|
||||
|
||||
type sessionAskInput struct {
|
||||
AgentID string `json:"agent_id" jsonschema:"custom agent ID"`
|
||||
Query string `json:"query" jsonschema:"user query"`
|
||||
SessionID string `json:"session_id,omitempty" jsonschema:"existing session to continue; auto-created when empty"`
|
||||
Reference bool `json:"reference,omitempty" jsonschema:"include indexed references"`
|
||||
Verbose bool `json:"verbose,omitempty" jsonschema:"include reasoning, tools, and lifecycle events"`
|
||||
}
|
||||
|
||||
type sessionAskOutput struct {
|
||||
Events []sse.ProjectedEvent `json:"events"`
|
||||
SessionID string `json:"session_id"`
|
||||
AgentID string `json:"agent_id"`
|
||||
Query string `json:"query"`
|
||||
}
|
||||
|
||||
func addSessionAsk(server *mcpsdk.Server, svc sessionAskService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "session_ask",
|
||||
Description: "Run a query through a custom agent. Returns a bounded answer-event projection by default; reference=true adds indexed citations and verbose=true adds reasoning, tool, and lifecycle events. MCP tools/call is buffered rather than streaming.",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "Ask a Custom Agent (session ask --agent)",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: false,
|
||||
IdempotentHint: false,
|
||||
OpenWorldHint: bptr(true),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, in sessionAskInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
if in.AgentID == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputMissingFlag, "agent_id is required")), nil, nil
|
||||
}
|
||||
if strings.TrimSpace(in.Query) == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputMissingFlag, "query cannot be empty")), nil, nil
|
||||
}
|
||||
projector := sse.NewProjector(in.Verbose, in.Reference, "")
|
||||
events := make([]sse.ProjectedEvent, 0)
|
||||
req := &sdk.AgentQARequest{
|
||||
Query: in.Query,
|
||||
AgentEnabled: true,
|
||||
AgentID: in.AgentID,
|
||||
Channel: "api",
|
||||
}
|
||||
// Auto-create session if not supplied. Sessions are agent-
|
||||
// agnostic at creation (verified against server source).
|
||||
sessionID := in.SessionID
|
||||
if sessionID == "" {
|
||||
sess, err := svc.CreateSession(ctx, &sdk.CreateSessionRequest{Title: "weknora mcp session_ask"})
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "create chat session")), nil, nil
|
||||
}
|
||||
sessionID = sess.ID
|
||||
}
|
||||
streamErr := svc.AgentQAStreamWithRequest(ctx, sessionID, req, func(r *sdk.AgentStreamResponse) error {
|
||||
if event, include := projector.Agent(r); include {
|
||||
events = append(events, event)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if streamErr != nil {
|
||||
return toolStreamError(cmdutil.WrapStream(streamErr, "agent-chat stream"), sessionID, ""), nil, nil
|
||||
}
|
||||
if !projector.Done() {
|
||||
return toolStreamError(cmdutil.NewError(cmdutil.CodeSSEStreamAborted, "stream ended without a terminal event"), sessionID, ""), nil, nil
|
||||
}
|
||||
return successResult(sessionAskOutput{
|
||||
Events: events,
|
||||
SessionID: sessionID,
|
||||
AgentID: in.AgentID,
|
||||
Query: in.Query,
|
||||
}), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- chunk_list ----------------------------------------------------------
|
||||
|
||||
type chunkListInput struct {
|
||||
DocID string `json:"doc_id" jsonschema:"document (knowledge entry) ID"`
|
||||
Limit int `json:"limit,omitempty" jsonschema:"max chunks to return (1..1000); defaults to 50"`
|
||||
}
|
||||
|
||||
type chunkListOutput struct {
|
||||
Chunks []sdk.Chunk `json:"chunks"`
|
||||
Total int64 `json:"total"`
|
||||
TruncatedAtLimit bool `json:"truncated_at_limit"`
|
||||
}
|
||||
|
||||
// chunkListDefaultLimit + chunkListMaxLimit mirror the schema's default+max.
|
||||
// MCP schema deliberately exposes only `limit`, not the CLI's full
|
||||
// --limit/--page/--page-size triple: LLM agents typically need a single
|
||||
// bounded fetch, not pagination workflows. Above 1000, fall back to the CLI.
|
||||
const (
|
||||
chunkListDefaultLimit = 50
|
||||
chunkListMaxLimit = 1000
|
||||
)
|
||||
|
||||
func addChunkList(server *mcpsdk.Server, svc chunkListService) {
|
||||
mcpsdk.AddTool(server, &mcpsdk.Tool{
|
||||
Name: "chunk_list",
|
||||
Description: "List chunks of a knowledge document for RAG retrieval debug. Returns at most `limit` chunks starting from ChunkIndex 0; if total chunks exceed limit, truncated_at_limit=true signals the agent to fall back to the CLI for paginated retrieval.",
|
||||
Annotations: &mcpsdk.ToolAnnotations{
|
||||
Title: "List Knowledge Chunks",
|
||||
DestructiveHint: bptr(false),
|
||||
ReadOnlyHint: true,
|
||||
IdempotentHint: true,
|
||||
OpenWorldHint: bptr(false),
|
||||
},
|
||||
}, func(ctx context.Context, _ *mcpsdk.CallToolRequest, in chunkListInput) (*mcpsdk.CallToolResult, any, error) {
|
||||
if in.DocID == "" {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputMissingFlag, "doc_id is required")), nil, nil
|
||||
}
|
||||
// `limit` is typed as int by chunkListInput, so the SDK rejects
|
||||
// non-numeric values at schema validation (e.g. "limit":"50")
|
||||
// before this handler runs. Default when unset; reject over-max
|
||||
// (rather than silently clamping) so the agent's request is never
|
||||
// quietly changed — matching search_chunks in this same file.
|
||||
limit := in.Limit
|
||||
if limit < 1 {
|
||||
limit = chunkListDefaultLimit
|
||||
}
|
||||
if limit > chunkListMaxLimit {
|
||||
return toolErrorResult(cmdutil.NewError(cmdutil.CodeInputInvalidArgument, fmt.Sprintf("limit must be in 1..%d", chunkListMaxLimit))), nil, nil
|
||||
}
|
||||
chunks, total, err := svc.ListKnowledgeChunks(ctx, in.DocID, 1, limit)
|
||||
if err != nil {
|
||||
return toolErrorResult(cmdutil.WrapHTTP(err, "list knowledge chunks")), nil, nil
|
||||
}
|
||||
if chunks == nil {
|
||||
chunks = []sdk.Chunk{}
|
||||
}
|
||||
return successResult(chunkListOutput{
|
||||
Chunks: chunks,
|
||||
Total: total,
|
||||
TruncatedAtLimit: total > int64(limit),
|
||||
}), nil, nil
|
||||
})
|
||||
}
|
||||
|
||||
// encodeDownload returns (content, isBase64). Heuristic: if the first 512
|
||||
// bytes contain a NUL, treat as binary. Otherwise it's UTF-8-ish text.
|
||||
// Matches what /usr/bin/file's "binary" heuristic does at a coarse level -
|
||||
// good enough to spare an agent from base64-decoding obvious text.
|
||||
func encodeDownload(buf []byte) (string, bool) {
|
||||
probe := buf
|
||||
if len(probe) > 512 {
|
||||
probe = probe[:512]
|
||||
}
|
||||
for _, b := range probe {
|
||||
if b == 0 {
|
||||
return base64.StdEncoding.EncodeToString(buf), true
|
||||
}
|
||||
}
|
||||
return string(buf), false
|
||||
}
|
||||
|
||||
// bptr returns a pointer to a bool literal. MCP ToolAnnotations uses
|
||||
// pointer types for DestructiveHint and OpenWorldHint so that "explicit
|
||||
// false" can be distinguished from "field omitted (default true per MCP
|
||||
// spec 2025-06-18)".
|
||||
func bptr(b bool) *bool { return &b }
|
||||
@@ -0,0 +1,105 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/Tencent/WeKnora/cli/internal/cmdutil"
|
||||
"github.com/Tencent/WeKnora/cli/internal/output"
|
||||
sdk "github.com/Tencent/WeKnora/client"
|
||||
)
|
||||
|
||||
// TestToolErrorResult_TypedError_PopulatesStructuredContent verifies the
|
||||
// helper maps a typed cmdutil.Error to MCP StructuredContent as *output.ErrDetail.
|
||||
func TestToolErrorResult_TypedError_PopulatesStructuredContent(t *testing.T) {
|
||||
err := cmdutil.NewError(cmdutil.CodeAuthUnauthenticated, "no creds")
|
||||
res := toolErrorResult(err)
|
||||
if !res.IsError {
|
||||
t.Errorf("expected IsError true")
|
||||
}
|
||||
if res.StructuredContent == nil {
|
||||
t.Fatal("expected StructuredContent populated")
|
||||
}
|
||||
sc, ok := res.StructuredContent.(*output.ErrDetail)
|
||||
if !ok {
|
||||
t.Fatalf("StructuredContent has unexpected type %T", res.StructuredContent)
|
||||
}
|
||||
if sc.Type != "auth.unauthenticated" {
|
||||
t.Errorf("type: got %v want auth.unauthenticated", sc.Type)
|
||||
}
|
||||
// retry_argv should fall back to ["weknora","auth","login"] via cmdutil default
|
||||
if !reflect.DeepEqual(sc.RetryArgv, []string{"weknora", "auth", "login"}) {
|
||||
t.Errorf("retry_argv: got %v", sc.RetryArgv)
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolErrorResult_TextFallback verifies the human-readable Content[0]
|
||||
// contains code + message + (hint? + retry?) lines.
|
||||
func TestToolErrorResult_TextFallback(t *testing.T) {
|
||||
err := cmdutil.NewError(cmdutil.CodeKBIDRequired, "kb_id is required")
|
||||
res := toolErrorResult(err)
|
||||
if len(res.Content) == 0 {
|
||||
t.Fatal("expected Content[0]")
|
||||
}
|
||||
// Assertion is best-effort on the text shape; exact prose may vary
|
||||
// with hint/retry defaults. Lock the type code prefix.
|
||||
text := contentText(t, res.Content[0])
|
||||
if !strings.Contains(text, "local.kb_id_required: kb_id is required") {
|
||||
t.Errorf("text fallback missing code:message; got %q", text)
|
||||
}
|
||||
}
|
||||
|
||||
// contentText extracts the Text field from an mcpsdk.Content value.
|
||||
// Lives in the test file to avoid leaking helpers into production.
|
||||
func contentText(t *testing.T, c any) string {
|
||||
t.Helper()
|
||||
// Direct assertion to the concrete type used by toolErrorResult.
|
||||
if tc, ok := c.(*mcpsdk.TextContent); ok {
|
||||
return tc.Text
|
||||
}
|
||||
// Fallback: some go-sdk versions may expose a GetText() method.
|
||||
if tc, ok := c.(interface{ GetText() string }); ok {
|
||||
return tc.GetText()
|
||||
}
|
||||
t.Fatalf("cannot extract text from Content of type %T", c)
|
||||
return ""
|
||||
}
|
||||
|
||||
// TestToolErrorResult_IsErrorFlagSet verifies IsError is always true.
|
||||
func TestToolErrorResult_IsErrorFlagSet(t *testing.T) {
|
||||
res := toolErrorResult(cmdutil.NewError(cmdutil.CodeInputMissingFlag, "doc_id is required"))
|
||||
if !res.IsError {
|
||||
t.Error("expected IsError=true")
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolErrorResult_ContentPresent verifies Content slice is non-empty.
|
||||
func TestToolErrorResult_ContentPresent(t *testing.T) {
|
||||
res := toolErrorResult(cmdutil.NewError(cmdutil.CodeInputInvalidArgument, "page_size must be in 1..1000"))
|
||||
if len(res.Content) == 0 {
|
||||
t.Error("expected at least one Content item")
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolErrorResult_NetworkError verifies a network error code round-trips
|
||||
// through toolErrorResult with the correct type in StructuredContent.
|
||||
func TestToolErrorResult_NetworkError(t *testing.T) {
|
||||
err := cmdutil.NewError(cmdutil.CodeNetworkError, "dial failed")
|
||||
res := toolErrorResult(err)
|
||||
if !res.IsError {
|
||||
t.Error("expected IsError=true")
|
||||
}
|
||||
sc, ok := res.StructuredContent.(*output.ErrDetail)
|
||||
if !ok {
|
||||
t.Fatalf("StructuredContent has unexpected type %T", res.StructuredContent)
|
||||
}
|
||||
if sc.Type != "network.error" {
|
||||
t.Errorf("type: got %v want network.error", sc.Type)
|
||||
}
|
||||
}
|
||||
|
||||
// Suppress unused-import warning if sdk isn't used in this file.
|
||||
var _ = sdk.KnowledgeBase{}
|
||||
@@ -0,0 +1,737 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
sdk "github.com/Tencent/WeKnora/client"
|
||||
)
|
||||
|
||||
// fakeSvc implements every narrow service interface ServiceClient embeds.
|
||||
// Each method records the last call args; per-test setup populates the
|
||||
// return values it wants to assert against.
|
||||
type fakeSvc struct {
|
||||
listKBs []sdk.KnowledgeBase
|
||||
listKBsErr error
|
||||
getKB *sdk.KnowledgeBase
|
||||
getKBErr error
|
||||
listDocs []sdk.Knowledge
|
||||
listDocsTotal int64
|
||||
listDocsErr error
|
||||
getDoc *sdk.Knowledge
|
||||
getDocErr error
|
||||
openDocName string
|
||||
openDocBody io.ReadCloser
|
||||
openDocErr error
|
||||
hybridResults []*sdk.SearchResult
|
||||
hybridErr error
|
||||
createSess *sdk.Session
|
||||
createSessErr error
|
||||
kbStreamEvents []*sdk.StreamResponse
|
||||
kbStreamErr error
|
||||
agents []sdk.Agent
|
||||
agentsErr error
|
||||
agent *sdk.Agent
|
||||
agentErr error
|
||||
agentEvents []*sdk.AgentStreamResponse
|
||||
agentStreamErr error
|
||||
chunks []sdk.Chunk
|
||||
chunksTotal int64
|
||||
chunksErr error
|
||||
// Captured args:
|
||||
calls struct {
|
||||
listKBs int
|
||||
kbViewID string
|
||||
docListKBID string
|
||||
docListFilter sdk.KnowledgeListFilter
|
||||
docViewID string
|
||||
openDocID string
|
||||
hybridKBID string
|
||||
hybridParams *sdk.SearchParams
|
||||
createSessReq *sdk.CreateSessionRequest
|
||||
kbQAReq *sdk.KnowledgeQARequest
|
||||
kbQASess string
|
||||
agentListN int
|
||||
agentViewID string
|
||||
agentReq *sdk.AgentQARequest
|
||||
agentSess string
|
||||
chunkDocID string
|
||||
chunkPage int
|
||||
chunkPageSize int
|
||||
}
|
||||
}
|
||||
|
||||
func (f *fakeSvc) ListKnowledgeBases(_ context.Context) ([]sdk.KnowledgeBase, error) {
|
||||
f.calls.listKBs++
|
||||
return f.listKBs, f.listKBsErr
|
||||
}
|
||||
func (f *fakeSvc) GetKnowledgeBase(_ context.Context, id string) (*sdk.KnowledgeBase, error) {
|
||||
f.calls.kbViewID = id
|
||||
return f.getKB, f.getKBErr
|
||||
}
|
||||
func (f *fakeSvc) ListKnowledgeWithFilter(_ context.Context, kbID string, _, _ int, filter sdk.KnowledgeListFilter) ([]sdk.Knowledge, int64, error) {
|
||||
f.calls.docListKBID = kbID
|
||||
f.calls.docListFilter = filter
|
||||
return f.listDocs, f.listDocsTotal, f.listDocsErr
|
||||
}
|
||||
func (f *fakeSvc) GetKnowledge(_ context.Context, id string) (*sdk.Knowledge, error) {
|
||||
f.calls.docViewID = id
|
||||
return f.getDoc, f.getDocErr
|
||||
}
|
||||
func (f *fakeSvc) OpenKnowledgeFile(_ context.Context, id string) (string, io.ReadCloser, error) {
|
||||
f.calls.openDocID = id
|
||||
return f.openDocName, f.openDocBody, f.openDocErr
|
||||
}
|
||||
func (f *fakeSvc) HybridSearch(
|
||||
_ context.Context, kbID string, p *sdk.SearchParams, _ ...sdk.ResourceURLOptions,
|
||||
) ([]*sdk.SearchResult, error) {
|
||||
f.calls.hybridKBID, f.calls.hybridParams = kbID, p
|
||||
return f.hybridResults, f.hybridErr
|
||||
}
|
||||
func (f *fakeSvc) CreateSession(_ context.Context, req *sdk.CreateSessionRequest) (*sdk.Session, error) {
|
||||
f.calls.createSessReq = req
|
||||
if f.createSess == nil && f.createSessErr == nil {
|
||||
return &sdk.Session{ID: "sess_auto"}, nil
|
||||
}
|
||||
return f.createSess, f.createSessErr
|
||||
}
|
||||
func (f *fakeSvc) KnowledgeQAStream(_ context.Context, sess string, req *sdk.KnowledgeQARequest, cb func(*sdk.StreamResponse) error, opts ...sdk.ResourceURLOptions) error {
|
||||
f.calls.kbQASess, f.calls.kbQAReq = sess, req
|
||||
for _, e := range f.kbStreamEvents {
|
||||
if err := cb(e); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return f.kbStreamErr
|
||||
}
|
||||
func (f *fakeSvc) ListAgents(_ context.Context) ([]sdk.Agent, error) {
|
||||
f.calls.agentListN++
|
||||
return f.agents, f.agentsErr
|
||||
}
|
||||
func (f *fakeSvc) GetAgent(_ context.Context, id string) (*sdk.Agent, error) {
|
||||
f.calls.agentViewID = id
|
||||
return f.agent, f.agentErr
|
||||
}
|
||||
func (f *fakeSvc) AgentQAStreamWithRequest(_ context.Context, sess string, req *sdk.AgentQARequest, cb sdk.AgentEventCallback, opts ...sdk.ResourceURLOptions) error {
|
||||
f.calls.agentSess, f.calls.agentReq = sess, req
|
||||
for _, e := range f.agentEvents {
|
||||
if err := cb(e); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return f.agentStreamErr
|
||||
}
|
||||
func (f *fakeSvc) ListKnowledgeChunks(_ context.Context, docID string, page, pageSize int, _ ...string) ([]sdk.Chunk, int64, error) {
|
||||
f.calls.chunkDocID = docID
|
||||
f.calls.chunkPage = page
|
||||
f.calls.chunkPageSize = pageSize
|
||||
return f.chunks, f.chunksTotal, f.chunksErr
|
||||
}
|
||||
|
||||
// newTestServer wires svc to an in-process MCP server and returns a
|
||||
// connected client session ready to CallTool against it.
|
||||
func newTestServer(t *testing.T, svc ServiceClient) (*mcpsdk.ClientSession, context.CancelFunc) {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "weknora-test", Version: "v0.0.0-test"}, nil)
|
||||
registerTools(server, svc)
|
||||
|
||||
st, ct := mcpsdk.NewInMemoryTransports()
|
||||
serverSession, err := server.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
cancel()
|
||||
t.Fatalf("server.Connect: %v", err)
|
||||
}
|
||||
client := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "test-client", Version: "v0.0.0"}, nil)
|
||||
clientSession, err := client.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
_ = serverSession.Close()
|
||||
cancel()
|
||||
t.Fatalf("client.Connect: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = clientSession.Close()
|
||||
_ = serverSession.Close()
|
||||
cancel()
|
||||
})
|
||||
return clientSession, cancel
|
||||
}
|
||||
|
||||
// callTool invokes name with args and returns the parsed structured output.
|
||||
func callTool(t *testing.T, c *mcpsdk.ClientSession, name string, args any, out any) *mcpsdk.CallToolResult {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.CallTool(ctx, &mcpsdk.CallToolParams{Name: name, Arguments: args})
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool(%s): %v", name, err)
|
||||
}
|
||||
if res.IsError {
|
||||
if len(res.Content) > 0 {
|
||||
t.Fatalf("tool %s returned error: %+v", name, res.Content)
|
||||
}
|
||||
t.Fatalf("tool %s returned error (no content)", name)
|
||||
}
|
||||
if out != nil && res.StructuredContent != nil {
|
||||
b, _ := json.Marshal(res.StructuredContent)
|
||||
if err := json.Unmarshal(b, out); err != nil {
|
||||
t.Fatalf("decode %s output: %v\nraw=%s", name, err, b)
|
||||
}
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
func TestTool_ListsRegistered(t *testing.T) {
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.ListTools(ctx, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("ListTools: %v", err)
|
||||
}
|
||||
want := []string{"kb_list", "kb_view", "doc_list", "doc_view", "doc_download", "search_chunks", "chat", "agent_list", "session_ask", "chunk_list"}
|
||||
got := map[string]bool{}
|
||||
for _, tool := range res.Tools {
|
||||
got[tool.Name] = true
|
||||
}
|
||||
for _, name := range want {
|
||||
if !got[name] {
|
||||
t.Errorf("missing tool %q in ListTools response", name)
|
||||
}
|
||||
}
|
||||
if len(res.Tools) != len(want) {
|
||||
t.Errorf("registered %d tools, want exactly %d (no scope creep)", len(res.Tools), len(want))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTool_SessionAsk_NotAgentInvoke asserts the MCP rename landed: the
|
||||
// registered set must contain "session_ask" and must NOT contain the
|
||||
// stale "agent_invoke" name (clean break, no deprecation alias).
|
||||
func TestTool_SessionAsk_NotAgentInvoke(t *testing.T) {
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.ListTools(ctx, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("ListTools: %v", err)
|
||||
}
|
||||
names := map[string]bool{}
|
||||
for _, tool := range res.Tools {
|
||||
names[tool.Name] = true
|
||||
}
|
||||
if !names["session_ask"] {
|
||||
t.Error("expected tool 'session_ask' to be registered")
|
||||
}
|
||||
if names["agent_invoke"] {
|
||||
t.Error("stale tool 'agent_invoke' must NOT be registered (clean break, no alias)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_KBList(t *testing.T) {
|
||||
svc := &fakeSvc{listKBs: []sdk.KnowledgeBase{{ID: "kb1", Name: "Marketing"}}}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out kbListOutput
|
||||
callTool(t, c, "kb_list", map[string]any{}, &out)
|
||||
if len(out.Items) != 1 || out.Items[0].ID != "kb1" {
|
||||
t.Errorf("got %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_KBView_RequiresKBID(t *testing.T) {
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.CallTool(ctx, &mcpsdk.CallToolParams{Name: "kb_view", Arguments: map[string]any{}})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected transport error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatal("expected IsError=true on missing kb_id")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_KBView(t *testing.T) {
|
||||
svc := &fakeSvc{getKB: &sdk.KnowledgeBase{ID: "kb_x", Name: "Eng"}}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out sdk.KnowledgeBase
|
||||
callTool(t, c, "kb_view", map[string]any{"kb_id": "kb_x"}, &out)
|
||||
if out.ID != "kb_x" || out.Name != "Eng" {
|
||||
t.Errorf("got %+v", out)
|
||||
}
|
||||
if svc.calls.kbViewID != "kb_x" {
|
||||
t.Errorf("kb_id not forwarded: %s", svc.calls.kbViewID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_DocList_DefaultPagination(t *testing.T) {
|
||||
svc := &fakeSvc{listDocs: []sdk.Knowledge{{ID: "k1"}}, listDocsTotal: 1}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out docListOutput
|
||||
callTool(t, c, "doc_list", map[string]any{"kb_id": "kb_x"}, &out)
|
||||
if out.Page != 1 || out.PageSize != 20 {
|
||||
t.Errorf("default pagination not applied: %+v", out)
|
||||
}
|
||||
if svc.calls.docListKBID != "kb_x" {
|
||||
t.Errorf("kb_id not forwarded: %s", svc.calls.docListKBID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_DocList_StatusFilter_Forwarded(t *testing.T) {
|
||||
svc := &fakeSvc{}
|
||||
c, _ := newTestServer(t, svc)
|
||||
callTool(t, c, "doc_list", map[string]any{"kb_id": "kb_x", "status": "failed"}, nil)
|
||||
if svc.calls.docListFilter.ParseStatus != "failed" {
|
||||
t.Errorf("status not forwarded as filter.ParseStatus: %+v", svc.calls.docListFilter)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTool_DocList_PassesFilterFields drives every C11 filter field at once
|
||||
// and asserts they all land on filter struct (AND-combined server-side).
|
||||
func TestTool_DocList_PassesFilterFields(t *testing.T) {
|
||||
svc := &fakeSvc{}
|
||||
c, _ := newTestServer(t, svc)
|
||||
args := map[string]any{
|
||||
"kb_id": "kb_x",
|
||||
"status": "completed",
|
||||
"keyword": "spec",
|
||||
"file_type": "pdf",
|
||||
"source": "api",
|
||||
"tag_id": "tag_42",
|
||||
"start_time": "2026-01-01T00:00:00Z",
|
||||
"end_time": "2026-12-31T23:59:59Z",
|
||||
}
|
||||
callTool(t, c, "doc_list", args, nil)
|
||||
f := svc.calls.docListFilter
|
||||
assert.Equal(t, "completed", f.ParseStatus)
|
||||
assert.Equal(t, "spec", f.Keyword)
|
||||
assert.Equal(t, "pdf", f.FileType)
|
||||
assert.Equal(t, "api", f.Source)
|
||||
assert.Equal(t, "tag_42", f.TagID)
|
||||
assert.False(t, f.StartTime.IsZero(), "start_time RFC3339 must populate filter.StartTime")
|
||||
assert.False(t, f.EndTime.IsZero(), "end_time RFC3339 must populate filter.EndTime")
|
||||
}
|
||||
|
||||
// TestTool_DocList_InvalidStartTime asserts malformed RFC3339 is rejected
|
||||
// at the handler boundary (before the SDK is called).
|
||||
func TestTool_DocList_InvalidStartTime(t *testing.T) {
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.CallTool(ctx, &mcpsdk.CallToolParams{
|
||||
Name: "doc_list",
|
||||
Arguments: map[string]any{"kb_id": "kb_x", "start_time": "tomorrow"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.True(t, res.IsError, "expected IsError=true on malformed RFC3339 start_time")
|
||||
}
|
||||
|
||||
// TestTool_DocList_InvalidEndTime mirrors the start_time guard for end_time.
|
||||
func TestTool_DocList_InvalidEndTime(t *testing.T) {
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.CallTool(ctx, &mcpsdk.CallToolParams{
|
||||
Name: "doc_list",
|
||||
Arguments: map[string]any{"kb_id": "kb_x", "end_time": "2026-05-01"}, // date-only, not RFC3339
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.True(t, res.IsError, "expected IsError=true on malformed RFC3339 end_time")
|
||||
}
|
||||
|
||||
func TestTool_DocView(t *testing.T) {
|
||||
svc := &fakeSvc{getDoc: &sdk.Knowledge{ID: "k1", FileName: "a.pdf"}}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out sdk.Knowledge
|
||||
callTool(t, c, "doc_view", map[string]any{"doc_id": "k1"}, &out)
|
||||
if out.ID != "k1" {
|
||||
t.Errorf("got %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_DocDownload_Text(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
openDocName: "notes.txt",
|
||||
openDocBody: io.NopCloser(strings.NewReader("hello world")),
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out docDownloadOutput
|
||||
callTool(t, c, "doc_download", map[string]any{"doc_id": "k1"}, &out)
|
||||
if out.Content != "hello world" {
|
||||
t.Errorf("content = %q", out.Content)
|
||||
}
|
||||
if out.IsBase64 {
|
||||
t.Error("text content should not be base64-encoded")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_DocDownload_BinaryBase64(t *testing.T) {
|
||||
// First 512 bytes contain a NUL → encodeDownload returns base64.
|
||||
bin := []byte{0x00, 0x01, 0x02, 0x03}
|
||||
svc := &fakeSvc{
|
||||
openDocName: "blob.bin",
|
||||
openDocBody: io.NopCloser(strings.NewReader(string(bin))),
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out docDownloadOutput
|
||||
callTool(t, c, "doc_download", map[string]any{"doc_id": "k1"}, &out)
|
||||
if !out.IsBase64 {
|
||||
t.Errorf("binary should be base64; got is_base64=%v content=%q", out.IsBase64, out.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_SearchChunks(t *testing.T) {
|
||||
svc := &fakeSvc{hybridResults: []*sdk.SearchResult{{KnowledgeID: "k1", Score: 0.9}}}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out searchChunksOutput
|
||||
callTool(t, c, "search_chunks", map[string]any{"kb_id": "kb_x", "query": "what is RAG"}, &out)
|
||||
if len(out.Results) != 1 || out.Results[0].KnowledgeID != "k1" {
|
||||
t.Errorf("got %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_SearchChunks_LimitCap(t *testing.T) {
|
||||
// 5 results, limit 3 → 3 returned.
|
||||
svc := &fakeSvc{}
|
||||
for i := 0; i < 5; i++ {
|
||||
svc.hybridResults = append(svc.hybridResults, &sdk.SearchResult{KnowledgeID: "k", Score: float64(i)})
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out searchChunksOutput
|
||||
callTool(t, c, "search_chunks", map[string]any{"kb_id": "kb_x", "query": "x", "limit": 3}, &out)
|
||||
if len(out.Results) != 3 {
|
||||
t.Errorf("limit not honored: got %d, want 3", len(out.Results))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTool_SearchChunks_PassesMatchCountFromLimit is a regression guard for
|
||||
// the v0.5 audit bug: the search_chunks dispatch built SearchParams without
|
||||
// setting MatchCount, so the server fell back to its default cap and the
|
||||
// client-side trim (results[:limit]) was a no-op when limit > server default.
|
||||
// Verifies the limit arg is threaded into SearchParams.MatchCount.
|
||||
func TestTool_SearchChunks_PassesMatchCountFromLimit(t *testing.T) {
|
||||
svc := &fakeSvc{}
|
||||
c, _ := newTestServer(t, svc)
|
||||
callTool(t, c, "search_chunks", map[string]any{"kb_id": "kb_x", "query": "test", "limit": 50}, nil)
|
||||
require.NotNil(t, svc.calls.hybridParams, "HybridSearch must be called with non-nil SearchParams")
|
||||
assert.Equal(t, 50, svc.calls.hybridParams.MatchCount, "MCP search_chunks must thread limit into SearchParams.MatchCount")
|
||||
}
|
||||
|
||||
func TestTool_Chat_DefaultReturnsAnswerEventsOnly(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
kbStreamEvents: []*sdk.StreamResponse{
|
||||
{Content: "Hello "},
|
||||
{Content: "world."},
|
||||
{KnowledgeReferences: []*sdk.SearchResult{{
|
||||
ID: "c1",
|
||||
KnowledgeID: "k1",
|
||||
ParentChunkID: "p1",
|
||||
Content: "bulky passage",
|
||||
}}},
|
||||
{ResponseType: sdk.ResponseTypeComplete},
|
||||
},
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out chatOutput
|
||||
callTool(t, c, "chat", map[string]any{"kb_id": "kb_x", "query": "ping"}, &out)
|
||||
if len(out.Events) != 2 || out.Events[0].Content != "Hello " || out.Events[1].Content != "world." {
|
||||
t.Errorf("answer events=%+v", out.Events)
|
||||
}
|
||||
for _, event := range out.Events {
|
||||
if event.ResponseType != "answer" || len(event.KnowledgeReferences) != 0 {
|
||||
t.Errorf("default output leaked non-answer data: %+v", event)
|
||||
}
|
||||
}
|
||||
if out.SessionID != "sess_auto" {
|
||||
t.Errorf("session_id = %q, want sess_auto", out.SessionID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCP_ChatVerbosePreservesThinking(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
kbStreamEvents: []*sdk.StreamResponse{
|
||||
{ResponseType: sdk.ResponseTypeThinking, Content: "let me reason..."},
|
||||
{ResponseType: sdk.ResponseTypeAnswer, Content: "final answer"},
|
||||
{ResponseType: sdk.ResponseTypeComplete},
|
||||
},
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out chatOutput
|
||||
callTool(t, c, "chat", map[string]any{"kb_id": "kb_x", "query": "deep question", "verbose": true}, &out)
|
||||
want := []string{"thinking", "answer", "complete"}
|
||||
if len(out.Events) != len(want) {
|
||||
t.Fatalf("events=%+v", out.Events)
|
||||
}
|
||||
for i, responseType := range want {
|
||||
if out.Events[i].ResponseType != responseType {
|
||||
t.Errorf("events[%d]=%q, want %q", i, out.Events[i].ResponseType, responseType)
|
||||
}
|
||||
}
|
||||
if out.KBID != "kb_x" {
|
||||
t.Errorf("kb_id = %q, want %q", out.KBID, "kb_x")
|
||||
}
|
||||
if out.Query != "deep question" {
|
||||
t.Errorf("query = %q, want %q", out.Query, "deep question")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCP_SessionAskVerboseAndReferenceReturnsBothDetailClasses(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
agentEvents: []*sdk.AgentStreamResponse{
|
||||
{ResponseType: sdk.AgentResponseTypeThinking, Content: "agent thinks"},
|
||||
{ResponseType: sdk.AgentResponseTypeToolCall, ID: "tc1", Content: "knowledge_search"},
|
||||
{ResponseType: sdk.AgentResponseTypeReferences, KnowledgeReferences: []*sdk.SearchResult{{
|
||||
ID: "c1",
|
||||
ParentChunkID: "p1",
|
||||
Content: "bulky passage",
|
||||
}}},
|
||||
{ResponseType: sdk.AgentResponseTypeAnswer, Content: "agent answer"},
|
||||
{ResponseType: sdk.AgentResponseTypeComplete, Done: true},
|
||||
},
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out sessionAskOutput
|
||||
callTool(t, c, "session_ask", map[string]any{"agent_id": "ag1", "query": "tool question", "verbose": true, "reference": true}, &out)
|
||||
want := []string{"thinking", "tool_call", "references", "answer", "complete"}
|
||||
if len(out.Events) != len(want) {
|
||||
t.Fatalf("events=%+v", out.Events)
|
||||
}
|
||||
for i, responseType := range want {
|
||||
if out.Events[i].ResponseType != responseType {
|
||||
t.Errorf("events[%d]=%q, want %q", i, out.Events[i].ResponseType, responseType)
|
||||
}
|
||||
}
|
||||
refs := out.Events[2].KnowledgeReferences
|
||||
if len(refs) != 1 || refs[0].ChunkID != "c1" || refs[0].ParentChunkID != "p1" {
|
||||
t.Errorf("reference indexes=%+v", refs)
|
||||
}
|
||||
if out.Query != "tool question" {
|
||||
t.Errorf("query = %q, want %q", out.Query, "tool question")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_Chat_ExistingSessionSkipsCreate(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
kbStreamEvents: []*sdk.StreamResponse{{ResponseType: sdk.ResponseTypeComplete}},
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
callTool(t, c, "chat", map[string]any{"kb_id": "kb_x", "query": "x", "session_id": "sess_existing"}, nil)
|
||||
if svc.calls.createSessReq != nil {
|
||||
t.Error("CreateSession should not fire when session_id is supplied")
|
||||
}
|
||||
if svc.calls.kbQASess != "sess_existing" {
|
||||
t.Errorf("session id not forwarded to QA stream: %s", svc.calls.kbQASess)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_AgentList(t *testing.T) {
|
||||
svc := &fakeSvc{agents: []sdk.Agent{{ID: "ag1", Name: "Research"}}}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out agentListOutput
|
||||
callTool(t, c, "agent_list", map[string]any{}, &out)
|
||||
if len(out.Items) != 1 || out.Items[0].ID != "ag1" {
|
||||
t.Errorf("got %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_SessionAsk(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
agentEvents: []*sdk.AgentStreamResponse{
|
||||
{ResponseType: sdk.AgentResponseTypeAnswer, Content: "result"},
|
||||
{ResponseType: sdk.AgentResponseTypeToolCall, ID: "c1", Content: "knowledge_search"},
|
||||
{ResponseType: sdk.AgentResponseTypeComplete, Done: true},
|
||||
},
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out sessionAskOutput
|
||||
callTool(t, c, "session_ask", map[string]any{"agent_id": "ag1", "query": "x"}, &out)
|
||||
if len(out.Events) != 1 || out.Events[0].ResponseType != "answer" || out.Events[0].Content != "result" {
|
||||
t.Errorf("default events=%+v", out.Events)
|
||||
}
|
||||
if out.AgentID != "ag1" {
|
||||
t.Errorf("agent_id = %q", out.AgentID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_SessionAsk_StreamAbort(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
agentEvents: []*sdk.AgentStreamResponse{{ResponseType: sdk.AgentResponseTypeAnswer, Content: "partial"}},
|
||||
agentStreamErr: errors.New("connection reset"),
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.CallTool(ctx, &mcpsdk.CallToolParams{Name: "session_ask", Arguments: map[string]any{"agent_id": "ag1", "query": "x"}})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected transport error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatal("expected IsError=true on mid-stream abort")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_Chat_StreamErrorIncludesSessionDetail(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
kbStreamEvents: []*sdk.StreamResponse{
|
||||
{ResponseType: sdk.ResponseTypeAnswer, Content: "partial"},
|
||||
{ResponseType: sdk.ResponseTypeError, Content: "boom", Done: true},
|
||||
},
|
||||
kbStreamErr: sdk.NewSSEStreamError("boom"),
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.CallTool(ctx, &mcpsdk.CallToolParams{Name: "chat", Arguments: map[string]any{"kb_id": "kb_x", "query": "q"}})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected transport error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatal("expected IsError=true on terminal stream error")
|
||||
}
|
||||
b, err := json.Marshal(res.StructuredContent)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal structured content: %v", err)
|
||||
}
|
||||
var detail struct {
|
||||
Type string `json:"type"`
|
||||
Detail map[string]any `json:"detail"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &detail); err != nil {
|
||||
t.Fatalf("unmarshal error detail: %v", err)
|
||||
}
|
||||
if detail.Type != "server.error" {
|
||||
t.Errorf("type=%q, want server.error", detail.Type)
|
||||
}
|
||||
if detail.Detail["session_id"] != "sess_auto" {
|
||||
t.Errorf("detail=%v, want session_id sess_auto", detail.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_ChunkList_Happy(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
chunks: []sdk.Chunk{{ID: "c1", ChunkIndex: 0, Content: "hello"}},
|
||||
chunksTotal: 1,
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out chunkListOutput
|
||||
callTool(t, c, "chunk_list", map[string]any{"doc_id": "doc_abc", "limit": 50}, &out)
|
||||
require.Len(t, out.Chunks, 1)
|
||||
assert.Equal(t, "c1", out.Chunks[0].ID)
|
||||
assert.Equal(t, "doc_abc", svc.calls.chunkDocID)
|
||||
assert.Equal(t, 1, svc.calls.chunkPage)
|
||||
assert.Equal(t, 50, svc.calls.chunkPageSize) // SDK page=1, pageSize=limit
|
||||
}
|
||||
|
||||
func TestTool_ChunkList_TruncatedAtLimit(t *testing.T) {
|
||||
svc := &fakeSvc{
|
||||
chunks: []sdk.Chunk{{ID: "c1"}},
|
||||
chunksTotal: 100, // more than limit
|
||||
}
|
||||
c, _ := newTestServer(t, svc)
|
||||
var out chunkListOutput
|
||||
callTool(t, c, "chunk_list", map[string]any{"doc_id": "d", "limit": 1}, &out)
|
||||
assert.True(t, out.TruncatedAtLimit)
|
||||
}
|
||||
|
||||
func TestTool_ChunkList_MissingDocID(t *testing.T) {
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.CallTool(ctx, &mcpsdk.CallToolParams{Name: "chunk_list", Arguments: map[string]any{"limit": 50}})
|
||||
require.NoError(t, err)
|
||||
require.True(t, res.IsError, "expected IsError=true on missing doc_id")
|
||||
}
|
||||
|
||||
// TestTool_ChunkList_NonNumericLimit asserts the MCP framework rejects a
|
||||
// string-valued `limit`. The schema declares limit as integer (via the
|
||||
// chunkListInput struct tag), so non-numeric values fail validation
|
||||
// before the handler runs.
|
||||
func TestTool_ChunkList_NonNumericLimit(t *testing.T) {
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.CallTool(ctx, &mcpsdk.CallToolParams{
|
||||
Name: "chunk_list",
|
||||
Arguments: map[string]any{"doc_id": "d", "limit": "50"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.True(t, res.IsError, "expected IsError=true when limit is a string")
|
||||
}
|
||||
|
||||
// derefBool is the test-side counterpart to bptr: ToolAnnotations uses
|
||||
// *bool for DestructiveHint/OpenWorldHint to distinguish "unset" from
|
||||
// "false", but assertions here treat unset as false (a nil pointer means
|
||||
// the field was omitted from the JSON wire envelope, which clients should
|
||||
// read as the documented default).
|
||||
func derefBool(p *bool) bool {
|
||||
if p == nil {
|
||||
return false
|
||||
}
|
||||
return *p
|
||||
}
|
||||
|
||||
// TestToolAnnotations_AllToolsHaveExpectedHints locks the per-tool hint
|
||||
// table. Each of the 10 registered tools must surface the exact
|
||||
// DestructiveHint / ReadOnlyHint / IdempotentHint / OpenWorldHint + Title
|
||||
// values shown below. This guards against silent drift during future
|
||||
// refactors (e.g. someone marking chat as readOnly, or an invoke tool as
|
||||
// closed-world).
|
||||
//
|
||||
// Note on plain-bool fields: ReadOnlyHint and IdempotentHint are bool
|
||||
// (not *bool) with `omitempty`. For invoke-class tools that explicitly set
|
||||
// them to false in the builder, the JSON envelope omits the field and the
|
||||
// client-side decode surfaces the zero value (false), which matches the
|
||||
// table.
|
||||
func TestToolAnnotations_AllToolsHaveExpectedHints(t *testing.T) {
|
||||
expected := map[string]struct {
|
||||
destructive bool
|
||||
readOnly bool
|
||||
idempotent bool
|
||||
openWorld bool
|
||||
title string
|
||||
}{
|
||||
"kb_list": {destructive: false, readOnly: true, idempotent: true, openWorld: false, title: "List Knowledge Bases"},
|
||||
"kb_view": {destructive: false, readOnly: true, idempotent: true, openWorld: false, title: "View Knowledge Base"},
|
||||
"doc_list": {destructive: false, readOnly: true, idempotent: true, openWorld: false, title: "List Documents"},
|
||||
"doc_view": {destructive: false, readOnly: true, idempotent: true, openWorld: false, title: "View Document"},
|
||||
"doc_download": {destructive: false, readOnly: true, idempotent: true, openWorld: false, title: "Download Document"},
|
||||
"search_chunks": {destructive: false, readOnly: true, idempotent: true, openWorld: false, title: "Search Knowledge Chunks"},
|
||||
"chat": {destructive: false, readOnly: false, idempotent: false, openWorld: true, title: "Chat with KB (Streaming RAG)"},
|
||||
"agent_list": {destructive: false, readOnly: true, idempotent: true, openWorld: false, title: "List Custom Agents"},
|
||||
"session_ask": {destructive: false, readOnly: false, idempotent: false, openWorld: true, title: "Ask a Custom Agent (session ask --agent)"},
|
||||
"chunk_list": {destructive: false, readOnly: true, idempotent: true, openWorld: false, title: "List Knowledge Chunks"},
|
||||
}
|
||||
|
||||
c, _ := newTestServer(t, &fakeSvc{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
res, err := c.ListTools(ctx, nil)
|
||||
require.NoError(t, err, "ListTools must succeed")
|
||||
|
||||
got := map[string]*mcpsdk.Tool{}
|
||||
for _, tool := range res.Tools {
|
||||
got[tool.Name] = tool
|
||||
}
|
||||
|
||||
for name, want := range expected {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
tool, ok := got[name]
|
||||
require.True(t, ok, "tool %q not registered", name)
|
||||
require.NotNil(t, tool.Annotations, "tool %q must set Annotations", name)
|
||||
a := tool.Annotations
|
||||
assert.Equal(t, want.title, a.Title, "Title")
|
||||
assert.Equal(t, want.destructive, derefBool(a.DestructiveHint), "DestructiveHint")
|
||||
assert.Equal(t, want.readOnly, a.ReadOnlyHint, "ReadOnlyHint")
|
||||
assert.Equal(t, want.idempotent, a.IdempotentHint, "IdempotentHint")
|
||||
assert.Equal(t, want.openWorld, derefBool(a.OpenWorldHint), "OpenWorldHint")
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user