ttychart-mcp/main.go

// Package main ttychart-mcp 终端图表 MCP 服务入口
//
// 该服务提供终端图表绘制功能，支持折线图、柱状图、散点图等类型
// 通过 MCP 协议与客户端通信，支持 stdio 和 HTTP 两种传输方式
//
// 使用方法:
//   - stdio 模式: 直接运行二进制文件
//   - HTTP 模式: 设置 HTTP_PORT 环境变量后运行
//
// 作者: titor
// 创建日期: 2026-04-15
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"

	"hub.gaomia.site/titor/ttychart-mcp/tools"

	"github.com/mark3labs/mcp-go/server"
)

// 全局日志器
var (
	logger *log.Logger
)

// 命令行参数定义
var (
	HTTPPort = flag.Int("port", 3100, "HTTP 服务端口号")
	Stdio    = flag.Bool("stdio", false, "使用 stdio 模式")
	Version  = flag.Bool("version", false, "显示版本信息")
)

// 版本信息
const (
	// 项目名称
	ProjectName = "ttychart-mcp"
	// 项目描述
	ProjectDesc = "终端图表 MCP 服务"
	// 版本号
	VersionNum = "0.1.0"
)

// init 初始化函数
// 做一些启动前的准备工作
func init() {
	// 初始化日志输出
	logger = log.New(os.Stdout, "["+ProjectName+"] ", log.LstdFlags)
	// 设置日志格式
	log.SetFlags(log.LstdFlags)
}

// main 主入口函数
//
// 根据命令行参数决定启动模式:
// 1. --stdio: 使用 stdio 模式
// 2. --port: 使用 HTTP 模式
// 3. 默认: 优先 HTTP 模式
func main() {
	// 解析命令行参数
	flag.Parse()

	// 显示版本信息
	if *Version {
		fmt.Printf("%s %s - %s\n", ProjectName, VersionNum, ProjectDesc)
		return
	}

	// 设置优雅退出
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// 处理系统信号
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-signalChan
		logger.Println("收到退出信号正在关闭服务...")
		cancel()
	}()

	// 输出启动信息
	logger.Printf("启动 %s v%s 服务...", ProjectName, VersionNum)

	// 选择启动模式
	selectStdio := *Stdio
	httpPort := os.Getenv("HTTP_PORT")
	if httpPort != "" && !selectStdio {
		// HTTP 模式
		port := httpPort
		startHTTPServer(ctx, port)
	} else if selectStdio {
		// stdio 模式
		startStdioServer(ctx)
	} else {
		// 默认 HTTP 模式
		startHTTPServer(ctx, fmt.Sprintf(":%d", *HTTPPort))
	}
}

// startStdioServer 启动 stdio 模式服务
//
// 参数:
//   - ctx: 上下文
func startStdioServer(ctx context.Context) {
	logger.Println("使用 stdio 模式启动服务...")

	// 创建 MCP 服务器
	srv := tools.NewServer()

	// 运行 stdio 服务
	if err := server.ServeStdio(srv); err != nil {
		logger.Printf("stdio 服务错误: %v", err)
		os.Exit(1)
	}
}

// startHTTPServer 启动 HTTP 模式服务
//
// 参数:
//   - ctx: 上下文
//   - addr: 监听地址，格式如 ":3100" 或 "localhost:3100"
func startHTTPServer(ctx context.Context, addr string) {
	logger.Printf("使用 HTTP 模式启动服务，监听地址: %s", addr)

	// 创建 MCP 服务器
	srv := tools.NewServer()

	// 运行 HTTP 服务
	httpSrv := server.NewStreamableHTTPServer(srv)
	if err := httpSrv.Start(addr); err != nil {
		logger.Printf("HTTP 服务错误: %v", err)
		os.Exit(1)
	}
}