From fde008a25d3f9b98939015d0a49df23691fbf1c5 Mon Sep 17 00:00:00 2001 From: Dennis Date: Tue, 14 Jul 2026 11:20:55 +0800 Subject: [PATCH] fix(shortcut): address review safety notes --- internal/app/legacy.go | 5 ++- internal/app/usage_recording.go | 14 ++++++- internal/app/usage_recording_test.go | 39 ++++++++++++++++++ internal/shortcut/runner.go | 51 +++++++++++++++++++++--- internal/shortcut/shortcut_test.go | 39 ++++++++++++++++++ internal/shortcut/smart/book.go | 10 ++--- internal/shortcut/smart/broadcast.go | 2 +- internal/shortcut/smart/replace_batch.go | 2 +- internal/shortcut/smart/wiki_new_doc.go | 7 ++-- 9 files changed, 149 insertions(+), 20 deletions(-) create mode 100644 internal/app/usage_recording_test.go diff --git a/internal/app/legacy.go b/internal/app/legacy.go index 726800f4..3eaab247 100644 --- a/internal/app/legacy.go +++ b/internal/app/legacy.go @@ -14,6 +14,7 @@ package app import ( + "log/slog" "sort" "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/cobracmd" @@ -33,7 +34,9 @@ func newLegacyPublicCommands(runner executor.Runner, caller edition.ToolCaller) // Load user-defined shortcuts (~/.dws/shortcuts/*.yaml) BEFORE compiling the // command tree, so distilled high-frequency operations mount alongside the // built-ins. Conflicts with built-ins are skipped inside Load. - _, _ = userdef.Load() + if _, err := userdef.Load(); err != nil { + slog.Warn("shortcut: failed to load user-defined shortcuts", "error", err) + } // Built-in + user shortcuts (`dws +`) share the same // command tree; mergeTopLevelCommands folds each shortcut's service parent // into the matching helper command so the `+leaf` sits alongside existing diff --git a/internal/app/usage_recording.go b/internal/app/usage_recording.go index b8218f76..4ef9f19e 100644 --- a/internal/app/usage_recording.go +++ b/internal/app/usage_recording.go @@ -32,8 +32,9 @@ func newRecordingToolCaller(inner edition.ToolCaller) edition.ToolCaller { } func (r recordingToolCaller) CallTool(ctx context.Context, product, tool string, args map[string]any) (*edition.ToolResult, error) { + recordedArgs := cloneToolArgs(args) res, err := r.inner.CallTool(ctx, product, tool, args) - usage.Append(product, tool, args, err == nil, r.inner.DryRun()) + usage.Append(product, tool, recordedArgs, err == nil, r.inner.DryRun()) return res, err } @@ -41,3 +42,14 @@ func (r recordingToolCaller) Format() string { return r.inner.Format() } func (r recordingToolCaller) DryRun() bool { return r.inner.DryRun() } func (r recordingToolCaller) Fields() string { return r.inner.Fields() } func (r recordingToolCaller) JQ() string { return r.inner.JQ() } + +func cloneToolArgs(args map[string]any) map[string]any { + if len(args) == 0 { + return nil + } + out := make(map[string]any, len(args)) + for k, v := range args { + out[k] = v + } + return out +} diff --git a/internal/app/usage_recording_test.go b/internal/app/usage_recording_test.go new file mode 100644 index 00000000..a7e63a3b --- /dev/null +++ b/internal/app/usage_recording_test.go @@ -0,0 +1,39 @@ +// Copyright 2026 Alibaba Group +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package app + +import "testing" + +func TestCloneToolArgsDefensiveCopy(t *testing.T) { + args := map[string]any{"page": 1, "query": "keep"} + cloned := cloneToolArgs(args) + args["page"] = 2 + args["extra"] = true + + if got := cloned["page"]; got != 1 { + t.Fatalf("cloned page = %#v, want 1", got) + } + if _, ok := cloned["extra"]; ok { + t.Fatal("clone changed after source map mutation") + } +} + +func TestCloneToolArgsEmpty(t *testing.T) { + if got := cloneToolArgs(nil); got != nil { + t.Fatalf("nil clone = %#v, want nil", got) + } + if got := cloneToolArgs(map[string]any{}); got != nil { + t.Fatalf("empty clone = %#v, want nil", got) + } +} diff --git a/internal/shortcut/runner.go b/internal/shortcut/runner.go index b5347a12..61738905 100644 --- a/internal/shortcut/runner.go +++ b/internal/shortcut/runner.go @@ -88,16 +88,55 @@ func (rt *RuntimeContext) CallMCP(tool string, params map[string]any) error { return helpers.CallMCPToolOnServer(rt.shortcut.product(), tool, params) } -// CallMCPData dispatches a tool call to an explicit MCP product and returns the -// PARSED response as data WITHOUT printing. This is the building block for -// multi-step ("smart") shortcuts: call a tool, read its output, feed it into the -// next call. Errors carry DWS's auth/PAT/business classification. +// CallMCPData dispatches a read-only tool call to an explicit MCP product and +// returns the PARSED response as data WITHOUT printing. This is the building +// block for multi-step ("smart") shortcuts: call a tool, read its output, feed +// it into the next call. Errors carry DWS's auth/PAT/business classification. // // The product is explicit (not the shortcut's own) because smart shortcuts // routinely cross services — e.g. resolve a name via `contact` then act via -// `chat`. Reads run even under --dry-run so a preview can still resolve inputs; -// the terminal write via CallMCP still honours --dry-run. +// `chat`. Reads run even under --dry-run so a preview can still resolve inputs. +// Write tools that need parsed responses must use CallMCPWriteData instead; as +// a backstop, obvious write-like tool names are rejected here under --dry-run. func (rt *RuntimeContext) CallMCPData(product, tool string, params map[string]any) (map[string]any, error) { + if rt.DryRun() && looksWriteTool(tool) { + return nil, dryRunWriteError(product, tool) + } + return rt.callMCPData(product, tool, params) +} + +// CallMCPWriteData dispatches a write tool call and returns its parsed response. +// Unlike CallMCPData, it refuses to run under --dry-run so smart shortcuts cannot +// accidentally perform writes while rendering a preview. +func (rt *RuntimeContext) CallMCPWriteData(product, tool string, params map[string]any) (map[string]any, error) { + if rt.DryRun() { + return nil, dryRunWriteError(product, tool) + } + return rt.callMCPData(product, tool, params) +} + +func dryRunWriteError(product, tool string) error { + return apperrors.NewValidation(fmt.Sprintf( + "--dry-run 下禁止执行写操作 %s/%s；请在 shortcut 中输出 preview 后返回", product, tool)) +} + +func looksWriteTool(tool string) bool { + tool = strings.TrimSpace(strings.ToLower(tool)) + for _, prefix := range []string{ + "add_", "append_", "approve_", "archive_", "cancel_", "create_", + "delete_", "disable_", "enable_", "grant_", "import_", "insert_", + "invite_", "move_", "publish_", "reject_", "remove_", "replace_", + "respond", "revoke_", "send_", "set_", "submit_", "update_", + "upload_", "write_", + } { + if strings.HasPrefix(tool, prefix) { + return true + } + } + return false +} + +func (rt *RuntimeContext) callMCPData(product, tool string, params map[string]any) (map[string]any, error) { if params == nil { params = map[string]any{} } diff --git a/internal/shortcut/shortcut_test.go b/internal/shortcut/shortcut_test.go index ef03b784..5cf8ddcb 100644 --- a/internal/shortcut/shortcut_test.go +++ b/internal/shortcut/shortcut_test.go @@ -14,6 +14,7 @@ package shortcut import ( + "strings" "testing" "github.com/spf13/cobra" @@ -123,6 +124,44 @@ func TestBuildGroupsByService(t *testing.T) { func noop(_ *RuntimeContext) error { return nil } +func TestCallMCPWriteDataRejectsDryRun(t *testing.T) { + root := &cobra.Command{Use: "dws"} + root.PersistentFlags().Bool("dry-run", false, "") + cmd := &cobra.Command{Use: "x"} + root.AddCommand(cmd) + if err := root.PersistentFlags().Set("dry-run", "true"); err != nil { + t.Fatalf("set dry-run: %v", err) + } + + rt := &RuntimeContext{cmd: cmd, shortcut: Shortcut{Service: "calendar"}} + _, err := rt.CallMCPWriteData("calendar", "create_calendar_event", map[string]any{"summary": "x"}) + if err == nil { + t.Fatal("expected dry-run write guard error") + } + if !strings.Contains(err.Error(), "calendar/create_calendar_event") { + t.Fatalf("error = %q, want tool name", err.Error()) + } +} + +func TestCallMCPDataRejectsLikelyWriteUnderDryRun(t *testing.T) { + root := &cobra.Command{Use: "dws"} + root.PersistentFlags().Bool("dry-run", false, "") + cmd := &cobra.Command{Use: "x"} + root.AddCommand(cmd) + if err := root.PersistentFlags().Set("dry-run", "true"); err != nil { + t.Fatalf("set dry-run: %v", err) + } + + rt := &RuntimeContext{cmd: cmd, shortcut: Shortcut{Service: "chat"}} + _, err := rt.CallMCPData("chat", "send_personal_message", map[string]any{"receiverUserId": "u"}) + if err == nil { + t.Fatal("expected likely write guard error") + } + if !strings.Contains(err.Error(), "chat/send_personal_message") { + t.Fatalf("error = %q, want tool name", err.Error()) + } +} + func TestValidationHelpers(t *testing.T) { s := Shortcut{ Service: "x", Command: "+y", diff --git a/internal/shortcut/smart/book.go b/internal/shortcut/smart/book.go index 7ac9fdb6..e8a397fa 100644 --- a/internal/shortcut/smart/book.go +++ b/internal/shortcut/smart/book.go @@ -87,9 +87,7 @@ var Book = shortcut.Shortcut{ } // Under --dry-run we resolved names (reads) to validate them, but must not - // create the event or send invites: the create/invite/rollback steps below - // go through CallMCPData (which, unlike CallMCP, does not honour --dry-run). - // Preview what would happen and return. + // create the event or send invites. Preview what would happen and return. if rt.DryRun() { return rt.Output(map[string]any{ "dryRun": true, @@ -99,7 +97,7 @@ var Book = shortcut.Shortcut{ } // Step 2 — create the event and pull the eventId from the response. - created, err := rt.CallMCPData("calendar", "create_calendar_event", createArgs) + created, err := rt.CallMCPWriteData("calendar", "create_calendar_event", createArgs) if err != nil { return err } @@ -110,13 +108,13 @@ var Book = shortcut.Shortcut{ // Step 3 — add all participants in one batch. attendeesToAdd copied // verbatim from the helper's `attendee add` call site. - if _, err := rt.CallMCPData("calendar", "add_calendar_participant", map[string]any{ + if _, err := rt.CallMCPWriteData("calendar", "add_calendar_participant", map[string]any{ "eventId": eventID, "attendeesToAdd": userIDs, }); err != nil { // Rollback: delete the just-created event so we never leave a // half-built event without its intended attendees. - _, delErr := rt.CallMCPData("calendar", "delete_calendar_event", map[string]any{ + _, delErr := rt.CallMCPWriteData("calendar", "delete_calendar_event", map[string]any{ "eventId": eventID, }) if delErr != nil { diff --git a/internal/shortcut/smart/broadcast.go b/internal/shortcut/smart/broadcast.go index 6a723fc1..13f5569d 100644 --- a/internal/shortcut/smart/broadcast.go +++ b/internal/shortcut/smart/broadcast.go @@ -79,7 +79,7 @@ var Broadcast = shortcut.Shortcut{ sent = append(sent, user.name) continue } - if _, err := rt.CallMCPData("chat", "send_personal_message", map[string]any{ + if _, err := rt.CallMCPWriteData("chat", "send_personal_message", map[string]any{ "receiverUserId": user.userID, "msgType": "markdown", "content": string(content), diff --git a/internal/shortcut/smart/replace_batch.go b/internal/shortcut/smart/replace_batch.go index 54fbbf88..204bb08d 100644 --- a/internal/shortcut/smart/replace_batch.go +++ b/internal/shortcut/smart/replace_batch.go @@ -85,7 +85,7 @@ var ReplaceBatch = shortcut.Shortcut{ for _, p := range pairs { // Params mirror helpers replace_minutes_text call site: taskUuid / // originalText / replacedText. - _, callErr := rt.CallMCPData("minutes", "replace_minutes_text", map[string]any{ + _, callErr := rt.CallMCPWriteData("minutes", "replace_minutes_text", map[string]any{ "taskUuid": taskUUID, "originalText": p.orig, "replacedText": p.repl, diff --git a/internal/shortcut/smart/wiki_new_doc.go b/internal/shortcut/smart/wiki_new_doc.go index 8eaf9622..ebac09f6 100644 --- a/internal/shortcut/smart/wiki_new_doc.go +++ b/internal/shortcut/smart/wiki_new_doc.go @@ -78,9 +78,8 @@ var WikiNewDoc = shortcut.Shortcut{ // Step 3 — create the document node under the space root. workspaceId / // name / type params copied verbatim from the helper's `wiki node create` // call site (create_file lives on the doc MCP server, so route there - // explicitly via CallMCPData rather than this shortcut's own product; - // adoc = 在线文档). CallMCPData does not honour --dry-run, so guard the - // create explicitly: the search above is a read and already ran. + // explicitly rather than this shortcut's own product; adoc = 在线文档). + // Guard the create explicitly: the search above is a read and already ran. if rt.DryRun() { return rt.Output(map[string]any{ "dryRun": true, @@ -89,7 +88,7 @@ var WikiNewDoc = shortcut.Shortcut{ "wouldCreateIn": workspaceID, }) } - created, err := rt.CallMCPData("doc", "create_file", map[string]any{ + created, err := rt.CallMCPWriteData("doc", "create_file", map[string]any{ "workspaceId": workspaceID, "name": title, "type": "adoc",