diff --git a/shortcuts/doc/doc_media_insert.go b/shortcuts/doc/doc_media_insert.go index 5be0eed1..3b5017d3 100644 --- a/shortcuts/doc/doc_media_insert.go +++ b/shortcuts/doc/doc_media_insert.go @@ -7,6 +7,7 @@ import ( "context" "fmt" "path/filepath" + "strings" "github.com/larksuite/cli/internal/output" "github.com/larksuite/cli/internal/validate" @@ -23,7 +24,7 @@ var alignMap = map[string]int{ var DocMediaInsert = common.Shortcut{ Service: "docs", Command: "+media-insert", - Description: "Insert a local image or file at the end of a Lark document (4-step orchestration + auto-rollback)", + Description: "Insert a local image or file into a Lark document (4-step orchestration + auto-rollback); appends to end by default, or inserts relative to a text selection with --selection-with-ellipsis", Risk: "write", Scopes: []string{"docs:document.media:upload", "docx:document:write_only", "docx:document:readonly"}, AuthTypes: []string{"user", "bot"}, @@ -33,6 +34,8 @@ var DocMediaInsert = common.Shortcut{ {Name: "type", Default: "image", Desc: "type: image | file"}, {Name: "align", Desc: "alignment: left | center | right"}, {Name: "caption", Desc: "image caption text"}, + {Name: "selection-with-ellipsis", Desc: "plain text (or 'start...end') that identifies the target block; the media is inserted after that block by default"}, + {Name: "before", Type: "bool", Desc: "insert before the matched block instead of after (requires --selection-with-ellipsis)"}, }, Validate: func(ctx context.Context, runtime *common.RuntimeContext) error { docRef, err := parseDocumentRef(runtime.Str("doc")) @@ -42,6 +45,9 @@ var DocMediaInsert = common.Shortcut{ if docRef.Kind == "doc" { return output.ErrValidation("docs +media-insert only supports docx documents; use a docx token/URL or a wiki URL that resolves to docx") } + if runtime.Bool("before") && strings.TrimSpace(runtime.Str("selection-with-ellipsis")) == "" { + return output.ErrValidation("--before requires --selection-with-ellipsis") + } return nil }, DryRun: func(ctx context.Context, runtime *common.RuntimeContext) *common.DryRunAPI { @@ -55,29 +61,70 @@ var DocMediaInsert = common.Shortcut{ filePath := runtime.Str("file") mediaType := runtime.Str("type") caption := runtime.Str("caption") + selection := strings.TrimSpace(runtime.Str("selection-with-ellipsis")) + hasSelection := selection != "" parentType := parentTypeForMediaType(mediaType) createBlockData := buildCreateBlockData(mediaType, 0) - createBlockData["index"] = "" + if hasSelection { + createBlockData["index"] = "" + } else { + createBlockData["index"] = "" + } batchUpdateData := buildBatchUpdateData("", mediaType, "", runtime.Str("align"), caption) d := common.NewDryRunAPI() + totalSteps := 4 + if docRef.Kind == "wiki" { + totalSteps++ + } + if hasSelection { + totalSteps++ + } + + positionLabel := map[bool]string{true: "before", false: "after"}[runtime.Bool("before")] + if docRef.Kind == "wiki" { documentID = "" stepBase = 2 - d.Desc("5-step orchestration: resolve wiki → query root → create block → upload file → bind to block (auto-rollback on failure)"). + d.Desc(fmt.Sprintf("%d-step orchestration: resolve wiki → query root →%s create block → upload file → bind to block (auto-rollback on failure)", + totalSteps, map[bool]string{true: " locate-doc →", false: ""}[hasSelection])). GET("/open-apis/wiki/v2/spaces/get_node"). Desc("[1] Resolve wiki node to docx document"). Params(map[string]interface{}{"token": docRef.Token}) } else { - d.Desc("4-step orchestration: query root → create block → upload file → bind to block (auto-rollback on failure)") + d.Desc(fmt.Sprintf("%d-step orchestration: query root →%s create block → upload file → bind to block (auto-rollback on failure)", + totalSteps, map[bool]string{true: " locate-doc →", false: ""}[hasSelection])) } d. GET("/open-apis/docx/v1/documents/:document_id/blocks/:document_id"). - Desc(fmt.Sprintf("[%d] Get document root block", stepBase)). + Desc(fmt.Sprintf("[%d] Get document root block", stepBase)) + + if hasSelection { + mcpEndpoint := common.MCPEndpoint(runtime.Config.Brand) + mcpArgs := map[string]interface{}{ + "doc_id": documentID, + "selection_with_ellipsis": selection, + "limit": 1, + } + d.POST(mcpEndpoint). + Desc(fmt.Sprintf("[%d] MCP locate-doc: find block matching selection (%s)", stepBase+1, positionLabel)). + Body(map[string]interface{}{ + "method": "tools/call", + "params": map[string]interface{}{ + "name": "locate-doc", + "arguments": mcpArgs, + }, + }). + Set("mcp_tool", "locate-doc"). + Set("args", mcpArgs) + stepBase++ + } + + d. POST("/open-apis/docx/v1/documents/:document_id/blocks/:document_id/children"). - Desc(fmt.Sprintf("[%d] Create empty block at document end", stepBase+1)). + Desc(fmt.Sprintf("[%d] Create empty block at target position", stepBase+1)). Body(createBlockData) appendDocMediaInsertUploadDryRun(d, filePath, parentType, stepBase+2) d.PATCH("/open-apis/docx/v1/documents/:document_id/blocks/batch_update"). @@ -126,13 +173,29 @@ var DocMediaInsert = common.Shortcut{ return err } - parentBlockID, insertIndex, err := extractAppendTarget(rootData, documentID) + parentBlockID, insertIndex, rootChildren, err := extractAppendTarget(rootData, documentID) if err != nil { return err } fmt.Fprintf(runtime.IO().ErrOut, "Root block ready: %s (%d children)\n", parentBlockID, insertIndex) - // Step 2: Create an empty block at the end of the document + selection := strings.TrimSpace(runtime.Str("selection-with-ellipsis")) + if selection != "" { + before := runtime.Bool("before") + fmt.Fprintf(runtime.IO().ErrOut, "Locating block matching selection: %q\n", selection) + idx, err := locateInsertIndex(runtime, documentID, selection, rootChildren, before) + if err != nil { + return err + } + insertIndex = idx + posLabel := "after" + if before { + posLabel = "before" + } + fmt.Fprintf(runtime.IO().ErrOut, "locate-doc matched: inserting %s at index %d\n", posLabel, insertIndex) + } + + // Step 2: Create an empty block at the target position fmt.Fprintf(runtime.IO().ErrOut, "Creating block at index %d\n", insertIndex) createData, err := runtime.CallAPI("POST", @@ -304,19 +367,116 @@ func buildBatchUpdateData(blockID, mediaType, fileToken, alignStr, caption strin } } -func extractAppendTarget(rootData map[string]interface{}, fallbackBlockID string) (string, int, error) { +func extractAppendTarget(rootData map[string]interface{}, fallbackBlockID string) (parentBlockID string, insertIndex int, children []interface{}, err error) { block, _ := rootData["block"].(map[string]interface{}) if len(block) == 0 { - return "", 0, output.Errorf(output.ExitAPI, "api_error", "failed to query document root block") + return "", 0, nil, output.Errorf(output.ExitAPI, "api_error", "failed to query document root block") } - parentBlockID := fallbackBlockID + parentBlockID = fallbackBlockID if blockID, _ := block["block_id"].(string); blockID != "" { parentBlockID = blockID } - children, _ := block["children"].([]interface{}) - return parentBlockID, len(children), nil + children, _ = block["children"].([]interface{}) + return parentBlockID, len(children), children, nil +} + +// locateInsertIndex uses the MCP locate-doc tool to find the root-level index +// at which to insert relative to the block matching selection. It walks the +// parent_id chain (using single-block GET calls when needed) to resolve nested +// blocks to their top-level ancestor in rootChildren. +func locateInsertIndex(runtime *common.RuntimeContext, documentID string, selection string, rootChildren []interface{}, before bool) (int, error) { + args := map[string]interface{}{ + "doc_id": documentID, + "selection_with_ellipsis": selection, + "limit": 1, + } + result, err := common.CallMCPTool(runtime, "locate-doc", args) + if err != nil { + return 0, err + } + + matches := common.GetSlice(result, "matches") + if len(matches) == 0 { + return 0, output.ErrWithHint( + output.ExitValidation, + "no_match", + fmt.Sprintf("locate-doc did not find any block matching selection %q", selection), + "check spelling or use 'start...end' syntax to narrow the selection", + ) + } + + matchMap, _ := matches[0].(map[string]interface{}) + anchorBlockID := common.GetString(matchMap, "anchor_block_id") + if anchorBlockID == "" { + // Fall back to first block entry if anchor_block_id is absent. + blocks := common.GetSlice(matchMap, "blocks") + if len(blocks) > 0 { + if b, ok := blocks[0].(map[string]interface{}); ok { + anchorBlockID = common.GetString(b, "block_id") + } + } + } + if anchorBlockID == "" { + return 0, output.Errorf(output.ExitAPI, "api_error", "locate-doc response missing anchor_block_id") + } + parentBlockID := common.GetString(matchMap, "parent_block_id") + + // Build root children set for O(1) lookup. + rootSet := make(map[string]int, len(rootChildren)) + for i, c := range rootChildren { + if id, ok := c.(string); ok { + rootSet[id] = i + } + } + + // Walk up the parent chain. locate-doc already gives us one level of parent, + // so most cases need zero extra API calls. + cur := anchorBlockID + nextParent := parentBlockID + visited := map[string]bool{} + const maxDepth = 8 + for depth := 0; depth < maxDepth; depth++ { + if visited[cur] { + break + } + visited[cur] = true + + if idx, ok := rootSet[cur]; ok { + if before { + return idx, nil + } + return idx + 1, nil + } + + // Advance: use the parent hint we already have, or fetch from API. + parent := nextParent + nextParent = "" // clear hint after first use + if parent == "" || parent == cur { + // Need to fetch this block to find its parent. + data, err := runtime.CallAPI("GET", + fmt.Sprintf("/open-apis/docx/v1/documents/%s/blocks/%s", + validate.EncodePathSegment(documentID), validate.EncodePathSegment(cur)), + nil, nil) + if err != nil { + return 0, err + } + block := common.GetMap(data, "block") + parent = common.GetString(block, "parent_id") + } + if parent == "" || parent == cur { + break + } + cur = parent + } + + return 0, output.ErrWithHint( + output.ExitValidation, + "block_not_reachable", + fmt.Sprintf("block matching selection %q is not reachable from document root", selection), + "try a top-level heading or paragraph as the selection", + ) } func extractCreatedBlockTargets(createData map[string]interface{}, mediaType string) (blockID, uploadParentNode, replaceBlockID string) { diff --git a/shortcuts/doc/doc_media_insert_test.go b/shortcuts/doc/doc_media_insert_test.go index 0e4f9bad..3355e33a 100644 --- a/shortcuts/doc/doc_media_insert_test.go +++ b/shortcuts/doc/doc_media_insert_test.go @@ -4,8 +4,17 @@ package doc import ( + "context" + "encoding/json" "reflect" + "strings" "testing" + + "github.com/spf13/cobra" + + "github.com/larksuite/cli/internal/cmdutil" + "github.com/larksuite/cli/internal/httpmock" + "github.com/larksuite/cli/shortcuts/common" ) func TestBuildCreateBlockDataUsesConcreteAppendIndex(t *testing.T) { @@ -109,7 +118,7 @@ func TestExtractAppendTargetUsesRootChildrenCount(t *testing.T) { }, } - blockID, index, err := extractAppendTarget(rootData, "fallback") + blockID, index, children, err := extractAppendTarget(rootData, "fallback") if err != nil { t.Fatalf("extractAppendTarget() unexpected error: %v", err) } @@ -119,6 +128,298 @@ func TestExtractAppendTargetUsesRootChildrenCount(t *testing.T) { if index != 3 { t.Fatalf("extractAppendTarget() index = %d, want 3", index) } + if len(children) != 3 { + t.Fatalf("extractAppendTarget() children len = %d, want 3", len(children)) + } +} + +// buildLocateDocMCPResponse builds a JSON-RPC 2.0 response for a locate-doc MCP call. +func buildLocateDocMCPResponse(matches []map[string]interface{}) map[string]interface{} { + resultJSON, _ := json.Marshal(map[string]interface{}{"matches": matches}) + return map[string]interface{}{ + "jsonrpc": "2.0", + "id": "test-id", + "result": map[string]interface{}{ + "content": []interface{}{ + map[string]interface{}{ + "type": "text", + "text": string(resultJSON), + }, + }, + }, + } +} + +func registerInsertWithSelectionStubs(reg interface { + Register(*httpmock.Stub) +}, docID, anchorBlockID, parentBlockID string) { + // Root block + reg.Register(&httpmock.Stub{ + Method: "GET", + URL: "/open-apis/docx/v1/documents/" + docID + "/blocks/" + docID, + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{ + "block": map[string]interface{}{ + "block_id": docID, + "children": []interface{}{"blk_a", "blk_b"}, + }, + }, + }, + }) + // MCP locate-doc + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "mcp.feishu.cn/mcp", + Body: buildLocateDocMCPResponse([]map[string]interface{}{ + {"anchor_block_id": anchorBlockID, "parent_block_id": parentBlockID}, + }), + }) + // Create block + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/docx/v1/documents/" + docID + "/blocks/" + docID + "/children", + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{ + "children": []interface{}{ + map[string]interface{}{"block_id": "blk_new", "block_type": 27, "image": map[string]interface{}{}}, + }, + }, + }, + }) + // Upload + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_all", + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{"file_token": "ftok_test"}, + }, + }) + // Batch update + reg.Register(&httpmock.Stub{ + Method: "PATCH", + URL: "/open-apis/docx/v1/documents/" + docID + "/blocks/batch_update", + Body: map[string]interface{}{"code": 0, "msg": "ok", "data": map[string]interface{}{}}, + }) +} + +// TestLocateInsertIndexAfterModeViaExecute verifies that --selection-with-ellipsis +// inserts after the matched root-level block (index = root index + 1). +func TestLocateInsertIndexAfterModeViaExecute(t *testing.T) { + f, _, _, reg := cmdutil.TestFactory(t, docsTestConfigWithAppID("locate-after-app")) + registerInsertWithSelectionStubs(reg, "doxcnSEL", "blk_a", "doxcnSEL") + + tmpDir := t.TempDir() + withDocsWorkingDir(t, tmpDir) + writeSizedDocTestFile(t, "img.png", 100) + + err := mountAndRunDocs(t, DocMediaInsert, []string{ + "+media-insert", + "--doc", "doxcnSEL", + "--file", "img.png", + "--selection-with-ellipsis", "Introduction", + "--as", "bot", + }, f, nil) + if err != nil { + t.Fatalf("Execute() error: %v", err) + } +} + +// TestLocateInsertIndexBeforeModeViaExecute verifies that --before inserts before +// the matched root-level block. +func TestLocateInsertIndexBeforeModeViaExecute(t *testing.T) { + f, _, _, reg := cmdutil.TestFactory(t, docsTestConfigWithAppID("locate-before-app")) + registerInsertWithSelectionStubs(reg, "doxcnSEL2", "blk_b", "doxcnSEL2") + + tmpDir := t.TempDir() + withDocsWorkingDir(t, tmpDir) + writeSizedDocTestFile(t, "img.png", 100) + + err := mountAndRunDocs(t, DocMediaInsert, []string{ + "+media-insert", + "--doc", "doxcnSEL2", + "--file", "img.png", + "--selection-with-ellipsis", "Architecture", + "--before", + "--as", "bot", + }, f, nil) + if err != nil { + t.Fatalf("Execute() error: %v", err) + } +} + +// TestLocateInsertIndexNestedBlockViaExecute verifies that a nested block's +// parent_block_id hint is used to walk to the root-level ancestor. +func TestLocateInsertIndexNestedBlockViaExecute(t *testing.T) { + f, _, _, reg := cmdutil.TestFactory(t, docsTestConfigWithAppID("locate-nested-app")) + + docID := "doxcnNESTED" + // Root block with blk_section and blk_other as children + reg.Register(&httpmock.Stub{ + Method: "GET", + URL: "/open-apis/docx/v1/documents/" + docID + "/blocks/" + docID, + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{ + "block": map[string]interface{}{ + "block_id": docID, + "children": []interface{}{"blk_section", "blk_other"}, + }, + }, + }, + }) + // MCP locate-doc returns blk_child nested under blk_section + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "mcp.feishu.cn/mcp", + Body: buildLocateDocMCPResponse([]map[string]interface{}{ + {"anchor_block_id": "blk_child", "parent_block_id": "blk_section"}, + }), + }) + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/docx/v1/documents/" + docID + "/blocks/" + docID + "/children", + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{ + "children": []interface{}{ + map[string]interface{}{"block_id": "blk_new", "block_type": 27, "image": map[string]interface{}{}}, + }, + }, + }, + }) + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_all", + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{"file_token": "ftok_nested"}, + }, + }) + reg.Register(&httpmock.Stub{ + Method: "PATCH", + URL: "/open-apis/docx/v1/documents/" + docID + "/blocks/batch_update", + Body: map[string]interface{}{"code": 0, "msg": "ok", "data": map[string]interface{}{}}, + }) + + tmpDir := t.TempDir() + withDocsWorkingDir(t, tmpDir) + writeSizedDocTestFile(t, "img.png", 100) + + err := mountAndRunDocs(t, DocMediaInsert, []string{ + "+media-insert", + "--doc", docID, + "--file", "img.png", + "--selection-with-ellipsis", "nested content", + "--as", "bot", + }, f, nil) + if err != nil { + t.Fatalf("Execute() error: %v", err) + } +} + +// TestLocateInsertIndexNoMatchReturnsError verifies that when locate-doc returns +// no matches, Execute returns a descriptive error. +func TestLocateInsertIndexNoMatchReturnsError(t *testing.T) { + f, _, _, reg := cmdutil.TestFactory(t, docsTestConfigWithAppID("locate-nomatch-app")) + + docID := "doxcnNOMATCH" + reg.Register(&httpmock.Stub{ + Method: "GET", + URL: "/open-apis/docx/v1/documents/" + docID + "/blocks/" + docID, + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{ + "block": map[string]interface{}{ + "block_id": docID, + "children": []interface{}{"blk_a"}, + }, + }, + }, + }) + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "mcp.feishu.cn/mcp", + Body: buildLocateDocMCPResponse([]map[string]interface{}{}), + }) + + tmpDir := t.TempDir() + withDocsWorkingDir(t, tmpDir) + writeSizedDocTestFile(t, "img.png", 100) + + err := mountAndRunDocs(t, DocMediaInsert, []string{ + "+media-insert", + "--doc", docID, + "--file", "img.png", + "--selection-with-ellipsis", "nonexistent text", + "--as", "bot", + }, f, nil) + if err == nil { + t.Fatal("expected no-match error, got nil") + } + if !strings.Contains(err.Error(), "no_match") && !strings.Contains(err.Error(), "did not find") { + t.Fatalf("unexpected error: %v", err) + } +} + +// TestLocateInsertIndexDryRunIncludesMCPStep verifies that the dry-run output +// includes a locate-doc MCP step when --selection-with-ellipsis is provided. +func TestLocateInsertIndexDryRunIncludesMCPStep(t *testing.T) { + t.Parallel() + + cmd := &cobra.Command{Use: "docs +media-insert"} + cmd.Flags().String("file", "", "") + cmd.Flags().String("doc", "", "") + cmd.Flags().String("type", "image", "") + cmd.Flags().String("align", "", "") + cmd.Flags().String("caption", "", "") + cmd.Flags().String("selection-with-ellipsis", "", "") + cmd.Flags().Bool("before", false, "") + _ = cmd.Flags().Set("file", "img.png") + _ = cmd.Flags().Set("doc", "doxcnABCDEF") + _ = cmd.Flags().Set("selection-with-ellipsis", "Introduction") + + rt := common.TestNewRuntimeContext(cmd, docsTestConfigWithAppID("dry-run-app")) + dryAPI := DocMediaInsert.DryRun(context.Background(), rt) + raw, _ := json.Marshal(dryAPI) + + var dry struct { + Description string `json:"description"` + API []struct { + Desc string `json:"desc"` + URL string `json:"url"` + Body map[string]interface{} `json:"body"` + } `json:"api"` + } + if err := json.Unmarshal(raw, &dry); err != nil { + t.Fatalf("decode dry-run: %v", err) + } + + foundMCP := false + for _, step := range dry.API { + if strings.Contains(step.Desc, "locate-doc") { + foundMCP = true + } + } + if !foundMCP { + t.Fatalf("dry-run should include a locate-doc step, got: %+v", dry.API) + } + if !strings.Contains(dry.Description, "locate-doc") { + t.Fatalf("dry-run description should mention 'locate-doc', got: %s", dry.Description) + } + + // Verify create-block step shows not + for _, step := range dry.API { + if strings.Contains(step.URL, "/children") && step.Body != nil { + if idx, ok := step.Body["index"]; ok { + if idx != "" { + t.Fatalf("create-block index in selection mode = %q, want ", idx) + } + } + } + } } func TestExtractCreatedBlockTargetsForImage(t *testing.T) {