feat: select explicit msglib receive source

This commit is contained in:
zhaoyilun
2026-07-10 09:04:06 +08:00
parent 0820717767
commit 3f270b454f
6 changed files with 187 additions and 15 deletions

View File

@@ -35,6 +35,10 @@ func RegisterISphereReadTools(server *mcp.Server, source ReceiveMessagesSource)
}
func RegisterISphereReadToolsWithDisplayEntities(server *mcp.Server, source ReceiveMessagesSource, displaySource DisplayEntitySource) {
RegisterISphereReadToolsWithReceiveSources(server, source, nil, displaySource)
}
func RegisterISphereReadToolsWithReceiveSources(server *mcp.Server, source ReceiveMessagesSource, msglibSource ReceiveMessagesSource, displaySource DisplayEntitySource) {
if source == nil {
source = isphere.EncryptedPacketLogSource{}
}
@@ -43,11 +47,15 @@ func RegisterISphereReadToolsWithDisplayEntities(server *mcp.Server, source Rece
Name: ToolNameReceiveMessages,
Description: "Return iSphere messages from the configured read-only log-backed source.",
}, func(ctx context.Context, _ *mcp.CallToolRequest, input ReceiveMessagesArgs) (*mcp.CallToolResult, map[string]any, error) {
if err := validateReceiveMessagesArgs(input); err != nil {
if err := validateReceiveMessagesCursor(input); err != nil {
return nil, nil, err
}
selectedSource, err := receiveMessagesSourceForPreference(input.SourcePreference, source, msglibSource)
if err != nil {
return nil, nil, err
}
started := time.Now().UTC()
result, err := source.ReceiveMessages(ctx, isphere.ReceiveMessagesQuery{
result, err := selectedSource.ReceiveMessages(ctx, isphere.ReceiveMessagesQuery{
ConversationID: input.ConversationID,
Query: input.Query,
Since: input.Since,
@@ -71,8 +79,26 @@ func RegisterISphereReadToolsWithDisplayEntities(server *mcp.Server, source Rece
})
}
func validateReceiveMessagesArgs(input ReceiveMessagesArgs) error {
return validateLocalReadonlySourceAndCursor(ToolNameReceiveMessages, input.SourcePreference, input.Cursor)
func validateReceiveMessagesCursor(input ReceiveMessagesArgs) error {
if strings.TrimSpace(input.Cursor) != "" {
return fmt.Errorf("%s cursor pagination is not available yet", ToolNameReceiveMessages)
}
return nil
}
func receiveMessagesSourceForPreference(sourcePreferenceValue string, primarySource ReceiveMessagesSource, msglibSource ReceiveMessagesSource) (ReceiveMessagesSource, error) {
sourcePreference := strings.ToLower(strings.TrimSpace(sourcePreferenceValue))
switch sourcePreference {
case "", "auto", "local_readonly":
return primarySource, nil
case "msglib_readonly":
if msglibSource == nil {
return nil, fmt.Errorf("%s source_preference %q is not configured", ToolNameReceiveMessages, sourcePreferenceValue)
}
return msglibSource, nil
default:
return nil, fmt.Errorf("%s source_preference %q is not available; only auto, local_readonly, and msglib_readonly are supported", ToolNameReceiveMessages, sourcePreferenceValue)
}
}
func validateLocalReadonlySourceAndCursor(toolName string, sourcePreferenceValue string, cursorValue string) error {

View File

@@ -213,6 +213,79 @@ func TestISphereReceiveMessagesToolUsesInjectedDisplayEntities(t *testing.T) {
}
}
func TestISphereReceiveMessagesToolSelectsExplicitMsgLibSource(t *testing.T) {
primary := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{
ID: "packet-msg-1",
ConversationID: "packet-conversation",
ConversationType: "chat",
SenderID: "packet-sender",
ReceiverID: "packet-receiver",
Text: "packet body",
Timestamp: "1783423807000",
}}},
}
msglibSource := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{
ID: "db-msg-1",
ConversationID: "db-conversation",
ConversationType: "group",
SenderID: "db-sender",
ReceiverID: "",
Text: "db body",
Timestamp: "1783423808000",
Source: "local_readonly",
RawRef: "msglib:tblMsgGroupPersonMsg",
}}},
}
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
RegisterISphereReadToolsWithReceiveSources(server, primary, msglibSource, nil)
})
defer cleanup()
defaultResult, err := session.CallTool(context.Background(), &mcp.CallToolParams{
Name: ToolNameReceiveMessages,
Arguments: map[string]any{"limit": 5},
})
if err != nil {
t.Fatalf("default call %s: %v", ToolNameReceiveMessages, err)
}
defaultPayload := decodeReceiveMessagesForSelectionTest(t, defaultResult.StructuredContent)
if len(defaultPayload.Messages) != 1 || defaultPayload.Messages[0].ID != "packet-msg-1" {
t.Fatalf("default messages = %+v, want packet source", defaultPayload.Messages)
}
if len(primary.queries) != 1 || len(msglibSource.queries) != 0 {
t.Fatalf("default source calls primary=%d msglib=%d", len(primary.queries), len(msglibSource.queries))
}
dbResult, err := session.CallTool(context.Background(), &mcp.CallToolParams{
Name: ToolNameReceiveMessages,
Arguments: map[string]any{"source_preference": "msglib_readonly", "query": "db", "limit": 3},
})
if err != nil {
t.Fatalf("msglib call %s: %v", ToolNameReceiveMessages, err)
}
dbPayload := decodeReceiveMessagesForSelectionTest(t, dbResult.StructuredContent)
if len(dbPayload.Messages) != 1 || dbPayload.Messages[0].ID != "db-msg-1" || dbPayload.Messages[0].RawRef != "msglib:tblMsgGroupPersonMsg" {
t.Fatalf("db messages = %+v, want explicit msglib source", dbPayload.Messages)
}
if len(msglibSource.queries) != 1 || msglibSource.queries[0].Query != "db" || msglibSource.queries[0].Limit != 3 {
t.Fatalf("msglib source queries = %+v", msglibSource.queries)
}
missingMsgLibSession, missingCleanup := connectToolsTestSession(t, func(server *mcp.Server) {
RegisterISphereReadToolsWithReceiveSources(server, primary, nil, nil)
})
defer missingCleanup()
missingResult, err := missingMsgLibSession.CallTool(context.Background(), &mcp.CallToolParams{
Name: ToolNameReceiveMessages,
Arguments: map[string]any{"source_preference": "msglib_readonly", "limit": 1},
})
if err == nil && !missingResult.IsError {
t.Fatalf("explicit msglib source was accepted without a configured MsgLib receive source")
}
}
func TestISphereReceiveMessagesToolValidatesRemainingContractArgs(t *testing.T) {
fake := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{
@@ -282,6 +355,29 @@ func TestISphereReceiveMessagesToolValidatesRemainingContractArgs(t *testing.T)
}
}
func decodeReceiveMessagesForSelectionTest(t *testing.T, structured any) struct {
Messages []struct {
ID string `json:"id"`
RawRef string `json:"raw_ref"`
} `json:"messages"`
} {
t.Helper()
var decoded struct {
Messages []struct {
ID string `json:"id"`
RawRef string `json:"raw_ref"`
} `json:"messages"`
}
payload, err := json.Marshal(structured)
if err != nil {
t.Fatalf("marshal structured content: %v", err)
}
if err := json.Unmarshal(payload, &decoded); err != nil {
t.Fatalf("decode structured content %s: %v", payload, err)
}
return decoded
}
type fakeReceiveMessagesSource struct {
queries []isphere.ReceiveMessagesQuery
result isphere.ReceiveMessagesResult