nullkey
9bb83b47fd
`weknora mcp serve` — long-lived stdio MCP (Model Context Protocol) transport that exposes a fixed, curated tool surface to MCP-aware agents (Claude Desktop, Claude Code, custom MCP clients). Curated tool set (readonly by default): - whoami — active context + tenant - search (hybrid retrieval against a KB) - kb list / view - doc list / view - agent list / view / invoke - session list / view The list is intentionally narrow to the read + agent-invoke surface; destructive verbs (`delete` / `empty` / `upload`) are gated behind `--write`. Schema is built from each leaf cobra command's flags so adding a new tool is a single registry entry plus a Service interface. Includes the simplify post-review polish + a second simplify pass to fold the resulting feedback (typed schemas, agent_help wording, unify chat / agent invoke option names).
275 lines
8.8 KiB
Go
275 lines
8.8 KiB
Go
package agentcmd
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/Tencent/WeKnora/cli/internal/cmdutil"
|
|
"github.com/Tencent/WeKnora/cli/internal/iostreams"
|
|
sdk "github.com/Tencent/WeKnora/client"
|
|
)
|
|
|
|
// scriptedInvokeSvc serves a canned stream of agent events to runInvoke.
|
|
type scriptedInvokeSvc struct {
|
|
createResp *sdk.Session
|
|
createErr error
|
|
events []*sdk.AgentStreamResponse
|
|
streamErr error
|
|
got struct {
|
|
sessionID string
|
|
req *sdk.AgentQARequest
|
|
}
|
|
}
|
|
|
|
func (s *scriptedInvokeSvc) CreateSession(_ context.Context, req *sdk.CreateSessionRequest) (*sdk.Session, error) {
|
|
if s.createResp == nil && s.createErr == nil {
|
|
return &sdk.Session{ID: "sess_auto", Title: req.Title}, nil
|
|
}
|
|
return s.createResp, s.createErr
|
|
}
|
|
|
|
func (s *scriptedInvokeSvc) AgentQAStreamWithRequest(_ context.Context, sessionID string, req *sdk.AgentQARequest, cb sdk.AgentEventCallback) error {
|
|
s.got.sessionID = sessionID
|
|
s.got.req = req
|
|
for _, e := range s.events {
|
|
if err := cb(e); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return s.streamErr
|
|
}
|
|
|
|
func answerEvent(content string) *sdk.AgentStreamResponse {
|
|
return &sdk.AgentStreamResponse{ResponseType: sdk.AgentResponseTypeAnswer, Content: content}
|
|
}
|
|
func doneEvent() *sdk.AgentStreamResponse {
|
|
return &sdk.AgentStreamResponse{ResponseType: sdk.AgentResponseTypeAnswer, Done: true}
|
|
}
|
|
func toolCallEvent(id, name string) *sdk.AgentStreamResponse {
|
|
return &sdk.AgentStreamResponse{
|
|
ResponseType: sdk.AgentResponseTypeToolCall,
|
|
ID: id,
|
|
Content: name,
|
|
}
|
|
}
|
|
func referencesEvent(refs []*sdk.SearchResult) *sdk.AgentStreamResponse {
|
|
return &sdk.AgentStreamResponse{
|
|
ResponseType: sdk.AgentResponseTypeReferences,
|
|
KnowledgeReferences: refs,
|
|
}
|
|
}
|
|
|
|
func TestInvoke_AccumulateMode_EmitsJSONEnvelope(t *testing.T) {
|
|
out, _ := iostreams.SetForTest(t)
|
|
svc := &scriptedInvokeSvc{
|
|
events: []*sdk.AgentStreamResponse{
|
|
answerEvent("Hello "),
|
|
answerEvent("world."),
|
|
referencesEvent([]*sdk.SearchResult{{KnowledgeID: "k1", KnowledgeTitle: "Doc 1"}}),
|
|
doneEvent(),
|
|
},
|
|
}
|
|
opts := &InvokeOptions{AgentID: "ag_x", Query: "ping"}
|
|
if err := runInvoke(context.Background(), opts, &cmdutil.JSONOptions{}, svc); err != nil {
|
|
t.Fatalf("runInvoke: %v", err)
|
|
}
|
|
var env struct {
|
|
OK bool `json:"ok"`
|
|
Data invokeData `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(out.Bytes(), &env); err != nil {
|
|
t.Fatalf("parse: %v\n%s", err, out.String())
|
|
}
|
|
if !env.OK {
|
|
t.Fatalf("ok=false: %s", out.String())
|
|
}
|
|
if env.Data.Answer != "Hello world." {
|
|
t.Errorf("answer = %q, want %q", env.Data.Answer, "Hello world.")
|
|
}
|
|
if env.Data.AgentID != "ag_x" {
|
|
t.Errorf("agent_id = %q, want ag_x", env.Data.AgentID)
|
|
}
|
|
if env.Data.Query != "ping" {
|
|
t.Errorf("query = %q, want ping", env.Data.Query)
|
|
}
|
|
if env.Data.SessionID != "sess_auto" {
|
|
t.Errorf("session_id = %q, want sess_auto", env.Data.SessionID)
|
|
}
|
|
if len(env.Data.References) != 1 || env.Data.References[0].KnowledgeID != "k1" {
|
|
t.Errorf("references missing: %+v", env.Data.References)
|
|
}
|
|
}
|
|
|
|
func TestInvoke_AutoCreatedSessionID_PassedAsAgentRequest(t *testing.T) {
|
|
_, _ = iostreams.SetForTest(t)
|
|
svc := &scriptedInvokeSvc{events: []*sdk.AgentStreamResponse{doneEvent()}}
|
|
opts := &InvokeOptions{AgentID: "ag_42", Query: "x"}
|
|
if err := runInvoke(context.Background(), opts, &cmdutil.JSONOptions{}, svc); err != nil {
|
|
t.Fatalf("runInvoke: %v", err)
|
|
}
|
|
if svc.got.sessionID != "sess_auto" {
|
|
t.Errorf("agent-chat got sessionID=%q, want sess_auto", svc.got.sessionID)
|
|
}
|
|
if svc.got.req == nil || svc.got.req.AgentID != "ag_42" {
|
|
t.Errorf("AgentID not forwarded: %+v", svc.got.req)
|
|
}
|
|
if !svc.got.req.AgentEnabled {
|
|
t.Error("AgentEnabled must be true for agent invoke")
|
|
}
|
|
}
|
|
|
|
func TestInvoke_ExistingSessionID_SkipsCreate(t *testing.T) {
|
|
_, _ = iostreams.SetForTest(t)
|
|
created := false
|
|
svc := &scriptedInvokeSvc{events: []*sdk.AgentStreamResponse{doneEvent()}}
|
|
// Wrap CreateSession to detect call.
|
|
svc.createResp = &sdk.Session{ID: "should_not_be_used"}
|
|
wrapped := &createSessionTracker{InvokeService: svc, called: &created}
|
|
opts := &InvokeOptions{AgentID: "ag", Query: "x", SessionID: "sess_existing"}
|
|
if err := runInvoke(context.Background(), opts, &cmdutil.JSONOptions{}, wrapped); err != nil {
|
|
t.Fatalf("runInvoke: %v", err)
|
|
}
|
|
if created {
|
|
t.Error("CreateSession should not be called when --session is set")
|
|
}
|
|
if svc.got.sessionID != "sess_existing" {
|
|
t.Errorf("agent-chat got sessionID=%q, want sess_existing", svc.got.sessionID)
|
|
}
|
|
}
|
|
|
|
type createSessionTracker struct {
|
|
InvokeService
|
|
called *bool
|
|
}
|
|
|
|
func (c *createSessionTracker) CreateSession(ctx context.Context, req *sdk.CreateSessionRequest) (*sdk.Session, error) {
|
|
*c.called = true
|
|
return c.InvokeService.CreateSession(ctx, req)
|
|
}
|
|
|
|
func TestInvoke_ToolEventsCaptured(t *testing.T) {
|
|
out, _ := iostreams.SetForTest(t)
|
|
svc := &scriptedInvokeSvc{events: []*sdk.AgentStreamResponse{
|
|
toolCallEvent("call_1", "knowledge_search"),
|
|
answerEvent("answer text"),
|
|
doneEvent(),
|
|
}}
|
|
opts := &InvokeOptions{AgentID: "ag", Query: "x"}
|
|
if err := runInvoke(context.Background(), opts, &cmdutil.JSONOptions{}, svc); err != nil {
|
|
t.Fatalf("runInvoke: %v", err)
|
|
}
|
|
var env struct {
|
|
Data invokeData `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(out.Bytes(), &env); err != nil {
|
|
t.Fatalf("parse: %v", err)
|
|
}
|
|
if len(env.Data.ToolEvents) != 1 {
|
|
t.Fatalf("expected 1 tool call, got %d", len(env.Data.ToolEvents))
|
|
}
|
|
if env.Data.ToolEvents[0].ID != "call_1" {
|
|
t.Errorf("tool_calls[0].id = %q, want call_1", env.Data.ToolEvents[0].ID)
|
|
}
|
|
}
|
|
|
|
func TestInvoke_EmptyQuery_Rejected(t *testing.T) {
|
|
_, _ = iostreams.SetForTest(t)
|
|
svc := &scriptedInvokeSvc{}
|
|
opts := &InvokeOptions{AgentID: "ag", Query: ""}
|
|
err := runInvoke(context.Background(), opts, nil, svc)
|
|
if err == nil {
|
|
t.Fatal("expected input.invalid_argument, got nil")
|
|
}
|
|
var typed *cmdutil.Error
|
|
if !errors.As(err, &typed) || typed.Code != cmdutil.CodeInputInvalidArgument {
|
|
t.Errorf("expected input.invalid_argument, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestInvoke_StreamAbortBeforeDone_MapsToSSEStreamAborted(t *testing.T) {
|
|
_, _ = iostreams.SetForTest(t)
|
|
svc := &scriptedInvokeSvc{
|
|
events: []*sdk.AgentStreamResponse{
|
|
answerEvent("partial"),
|
|
},
|
|
streamErr: errors.New("connection reset"),
|
|
}
|
|
opts := &InvokeOptions{AgentID: "ag", Query: "x"}
|
|
err := runInvoke(context.Background(), opts, &cmdutil.JSONOptions{}, svc)
|
|
if err == nil {
|
|
t.Fatal("expected stream-aborted error")
|
|
}
|
|
var typed *cmdutil.Error
|
|
if !errors.As(err, &typed) || typed.Code != cmdutil.CodeSSEStreamAborted {
|
|
t.Errorf("expected local.sse_stream_aborted, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestInvoke_NoDoneEvent_MapsToSSEStreamAborted(t *testing.T) {
|
|
_, _ = iostreams.SetForTest(t)
|
|
svc := &scriptedInvokeSvc{events: []*sdk.AgentStreamResponse{answerEvent("incomplete")}}
|
|
opts := &InvokeOptions{AgentID: "ag", Query: "x"}
|
|
err := runInvoke(context.Background(), opts, &cmdutil.JSONOptions{}, svc)
|
|
if err == nil {
|
|
t.Fatal("expected stream-aborted error")
|
|
}
|
|
var typed *cmdutil.Error
|
|
if !errors.As(err, &typed) || typed.Code != cmdutil.CodeSSEStreamAborted {
|
|
t.Errorf("expected local.sse_stream_aborted, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestInvoke_CreateSessionFails_MapsToSessionCreateFailed(t *testing.T) {
|
|
_, _ = iostreams.SetForTest(t)
|
|
svc := &scriptedInvokeSvc{createErr: errors.New("connection refused")}
|
|
opts := &InvokeOptions{AgentID: "ag", Query: "x"}
|
|
err := runInvoke(context.Background(), opts, nil, svc)
|
|
if err == nil {
|
|
t.Fatal("expected session_create_failed")
|
|
}
|
|
var typed *cmdutil.Error
|
|
if !errors.As(err, &typed) || typed.Code != cmdutil.CodeSessionCreateFailed {
|
|
t.Errorf("expected server.session_create_failed, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestInvoke_Cancellation_MapsToUserAborted(t *testing.T) {
|
|
_, _ = iostreams.SetForTest(t)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel() // pre-cancel
|
|
svc := &scriptedInvokeSvc{streamErr: context.Canceled}
|
|
opts := &InvokeOptions{AgentID: "ag", Query: "x"}
|
|
err := runInvoke(ctx, opts, &cmdutil.JSONOptions{}, svc)
|
|
if err == nil {
|
|
t.Fatal("expected user_aborted")
|
|
}
|
|
var typed *cmdutil.Error
|
|
if !errors.As(err, &typed) || typed.Code != cmdutil.CodeUserAborted {
|
|
t.Errorf("expected local.user_aborted, got %v", err)
|
|
}
|
|
}
|
|
|
|
// Sanity: human-mode output writes the answer body and a tool-trace footer.
|
|
func TestInvoke_Human_Accumulate_PrintsAnswerAndFooter(t *testing.T) {
|
|
out, _ := iostreams.SetForTest(t)
|
|
svc := &scriptedInvokeSvc{events: []*sdk.AgentStreamResponse{
|
|
answerEvent("hello"),
|
|
toolCallEvent("c1", "knowledge_search"),
|
|
doneEvent(),
|
|
}}
|
|
opts := &InvokeOptions{AgentID: "ag", Query: "x", NoStream: true}
|
|
if err := runInvoke(context.Background(), opts, nil, svc); err != nil {
|
|
t.Fatalf("runInvoke: %v", err)
|
|
}
|
|
got := out.String()
|
|
if !strings.Contains(got, "hello") {
|
|
t.Errorf("answer body missing: %q", got)
|
|
}
|
|
if !strings.Contains(got, "Tool trace") {
|
|
t.Errorf("tool trace footer missing: %q", got)
|
|
}
|
|
}
|