From 4f19a2df158379f07d9d0a0898cff7a9f382f94a Mon Sep 17 00:00:00 2001
From: John Doe <miromaxxs@gmail.com>
Date: Mon, 7 Apr 2025 13:32:49 +0300
Subject: [PATCH] feat: add tool draft msg

---
 README.md            |  2 +-
 internal/tg/draft.go | 54 ++++++++++++++++++++++++++++++++++++++++++++
 serve.go             | 12 ++++++++++
 3 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 internal/tg/draft.go

diff --git a/README.md b/README.md
index 7a645ef..710191c 100644
--- a/README.md
+++ b/README.md
@@ -28,10 +28,10 @@ The Model Context Protocol (MCP) is a system that lets AI apps, like Claude Desk
 - [x] Get current user data
 - [x] Get the list of dialogs (chats, channels, groups)
 - [x] Get the list of (unread) messages in the given dialog
+- [x] Draft a message
 - [ ] Mark chanel as read
 - [ ] Retrieve messages by date and time
 - [ ] Get the list of contacts
-- [ ] Draft a message
 
 ## Installation
 
diff --git a/internal/tg/draft.go b/internal/tg/draft.go
new file mode 100644
index 0000000..926f9d4
--- /dev/null
+++ b/internal/tg/draft.go
@@ -0,0 +1,54 @@
+package tg
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/gotd/td/telegram/message"
+	"github.com/gotd/td/tg"
+	mcp "github.com/metoro-io/mcp-golang"
+	"github.com/pkg/errors"
+)
+
+type DraftArguments struct {
+	Name string `json:"name" jsonschema:"required,description=Name of the dialog"`
+	Text string `json:"text" jsonschema:"required,description=Plain text of the message"`
+}
+
+type DraftResponse struct {
+	Success bool `json:"success"`
+}
+
+func (c *Client) SendDraft(args DraftArguments) (*mcp.ToolResponse, error) {
+	var ok bool
+	client := c.T()
+	if err := client.Run(context.Background(), func(ctx context.Context) (err error) {
+		api := client.API()
+
+		sender := message.NewSender(api)
+		inputPeer, err := sender.Resolve(args.Name).AsInputPeer(ctx)
+		if err != nil {
+			return fmt.Errorf("failed to resolve name: %w", err)
+		}
+
+		ok, err = api.MessagesSaveDraft(ctx, &tg.MessagesSaveDraftRequest{
+			Peer:    inputPeer,
+			Message: args.Text,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to get history: %w", err)
+		}
+
+		return nil
+	}); err != nil {
+		return nil, errors.Wrap(err, "failed to get history")
+	}
+
+	jsonData, err := json.Marshal(DraftResponse{Success: ok})
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to marshal response")
+	}
+
+	return mcp.NewToolResponse(mcp.NewTextContent(string(jsonData))), nil
+}
diff --git a/serve.go b/serve.go
index be6d38f..3b15561 100644
--- a/serve.go
+++ b/serve.go
@@ -55,6 +55,13 @@ func serve(ctx context.Context, cmd *cli.Command) error {
 
 		log.Info().RawJSON("answer", []byte(answer.Content[0].TextContent.Text)).Msg("Check GetHistory: OK")
 
+		answer, err = client.SendDraft(tg.DraftArguments{Name: os.Getenv("TG_TEST_USERNAME"), Text: "test draft"})
+		if err != nil {
+			return fmt.Errorf("send draft: %w", err)
+		}
+
+		log.Info().RawJSON("answer", []byte(answer.Content[0].TextContent.Text)).Msg("Check SendDraft: OK")
+
 		return nil
 	}
 
@@ -73,6 +80,11 @@ func serve(ctx context.Context, cmd *cli.Command) error {
 		return fmt.Errorf("register dialogs tool: %w", err)
 	}
 
+	err = server.RegisterTool("tg_send", "Send draft message to dialog (channel, user)", client.SendDraft)
+	if err != nil {
+		return fmt.Errorf("register dialogs tool: %w", err)
+	}
+
 	if err := server.Serve(); err != nil {
 		return fmt.Errorf("serve: %w", err)
 	}