Deep dive kiến trúc ReAct runtime trong Go: Xây dựng AI Agent tương thích chuẩn MCP

J.A.R.V.I.S. là một AI Agent platform mình tự build từ đầu. Runtime ReAct (Reason + Act) viết hoàn toàn bằng Go 1.25+.

Mấy phần dưới đây là những chỗ mình đã thử sai, đổi hướng, rồi mới chốt được thiết kế hiện tại.

Multi-agent routing: mỗi loại việc một agent

Thay vì một agent “biết tất cả” xử lý mọi loại request, J.A.R.V.I.S. định tuyến request đến agent chuyên biệt (coding, research, scheduling, trò chuyện thông thường) dựa trên một bước intent classification nhẹ chạy trước khi gọi LLM chính.

Lý do thiết kế theo hướng này thay vì một mega-prompt duy nhất:

• Giảm chi phí token: agent chuyên biệt chỉ cần system prompt ngắn, tập trung đúng domain, thay vì nhồi tất cả instruction + tool definition của toàn hệ thống vào một context window.
• Tăng độ chính xác: một agent “coding” được tối ưu riêng cho việc đọc code, gọi tool search/file, trả lời ít bị trôi ý hơn so với một agent tổng quát phải cân bằng quá nhiều nhiệm vụ.
• Dễ mở rộng: thêm agent mới (ví dụ agent “finance” chuyên xử lý câu hỏi về ngân sách) chỉ cần đăng ký route mới, không phải sửa prompt của các agent khác.

Bước classification dùng một model nhỏ, latency thấp (dưới 100ms) để gắn nhãn intent, sau đó router mới chuyển tiếp sang agent phù hợp kèm theo context đã được rút gọn.

Một request đi qua runtime như sau:

flowchart TD
    U["Request từ user"] --> IC["Intent classification<br/>model nhỏ, dưới 100ms"]
    IC --> RT{"Router chọn agent"}
    RT --> A1["Agent coding"]
    RT --> A2["Agent research"]
    RT --> A3["Agent scheduling"]
    RT --> A4["Agent trò chuyện thông thường"]
    A1 --> CTX["Context đã được rút gọn"]
    A2 --> CTX
    A3 --> CTX
    A4 --> CTX
    CTX --> LLM["LLM chính của agent đó"]

Chạy tool song song bằng errgroup

Khi một task cần gọi nhiều tool độc lập (search web, query database, đọc file), runtime dùng errgroup để chạy song song thay vì tuần tự — đây là chỗ Go thể hiện rõ lợi thế so với các runtime viết bằng ngôn ngữ khác mà mình từng thử:

func ExecuteToolsParallel(ctx context.Context, calls []ToolCall) ([]ToolResult, error) {
	g, ctx := errgroup.WithContext(ctx)
	results := make([]ToolResult, len(calls))

	for i, call := range calls {
		i, call := i, call // capture biến vòng lặp đúng cách
		g.Go(func() error {
			res, err := executeTool(ctx, call)
			if err != nil {
				return fmt.Errorf("tool %s failed: %w", call.Name, err)
			}
			results[i] = res
			return nil
		})
	}

	if err := g.Wait(); err != nil {
		return nil, err
	}
	return results, nil
}

Điểm quan trọng của errgroup.WithContext: nếu một tool call lỗi, context sẽ bị cancel ngay lập tức, các tool call còn lại đang chạy sẽ nhận được tín hiệu huỷ thay vì tiếp tục chạy vô ích. Cách này giảm latency tổng của một “reasoning step” có 5 tool call từ ~4 giây (chạy tuần tự) xuống dưới 1 giây.

Nhánh chạy song song và nhánh lỗi trông như thế này:

flowchart TD
    S["Reasoning step<br/>5 tool call độc lập"] --> EG["errgroup: chạy song song<br/>mỗi goroutine một tool"]
    EG --> OK["Tất cả tool xong"]
    EG --> FAIL["Một tool lỗi"]
    OK --> RES["Gom kết quả: dưới 1 giây<br/>thay vì ~4 giây chạy tuần tự"]
    FAIL --> STOP["context bị cancel ngay<br/>các tool còn lại nhận tín hiệu huỷ, không chạy tiếp vô ích"]

Mỗi loại tool một timeout riêng

Không phải tool nào cũng cần cùng một timeout. Tool gọi search engine có thể chờ 3 giây, nhưng tool query database nội bộ chỉ nên chờ tối đa 500ms trước khi coi là lỗi. Runtime của J.A.R.V.I.S. gắn timeout riêng cho từng loại tool thông qua context.WithTimeout lồng bên trong mỗi goroutine, thay vì dùng một timeout chung cho cả batch.

Bộ nhớ 3 tầng: Vector, Graph và Redis

Một AI Agent hữu ích cần nhớ được nhiều hơn những gì vừa nói trong 5 phút trước — nhưng nhớ “tất cả” theo kiểu nhồi toàn bộ lịch sử vào context window vừa tốn token vừa không hiệu quả khi hội thoại kéo dài nhiều tuần. J.A.R.V.I.S. tách bộ nhớ thành 3 tầng, mỗi tầng phục vụ một loại truy vấn khác nhau:

• Vector memory (MongoDB Atlas Vector Search): semantic recall cho hội thoại dài hạn. Khi user hỏi lại một chủ đề đã bàn vài tuần trước, hệ thống search theo similarity embedding thay vì tìm kiếm từ khoá chính xác.
• Graph memory: lưu quan hệ giữa entity (người, dự án, công cụ) để trả lời câu hỏi suy luận nhiều bước — ví dụ “dự án mà Trinh đang làm cùng với team Backend là gì?” cần nối nhiều cạnh quan hệ thay vì chỉ tìm một đoạn text liên quan.
• Redis: short-term working memory cho context cửa sổ hiện tại, TTL theo phiên làm việc. Đây là tầng nhanh nhất, phục vụ trực tiếp cho vòng ReAct hiện tại mà không cần round-trip tới vector DB.

Khi xử lý một câu hỏi, runtime truy vấn cả 3 tầng song song (cũng dùng errgroup như phần tool execution ở trên), rồi hợp nhất kết quả theo thứ tự ưu tiên: Redis (ngữ cảnh tức thời) → Graph (quan hệ) → Vector (ngữ nghĩa dài hạn).

Ba tầng đó chạy song song, chỉ khác nhau ở lúc hợp nhất:

flowchart TD
    Q["Câu hỏi mới"] --> RV["Truy vấn cả 3 tầng song song<br/>cũng dùng errgroup"]
    RV --> RM["Redis<br/>working memory, TTL theo phiên"]
    RV --> GM["Graph memory<br/>quan hệ giữa entity"]
    RV --> VM["Vector memory<br/>MongoDB Atlas Vector Search"]
    RM --> MRG["Hợp nhất theo thứ tự ưu tiên"]
    GM --> MRG
    VM --> MRG
    MRG --> OUT["Ngữ cảnh cho vòng ReAct hiện tại"]

MCP client/server và streaming SSE

J.A.R.V.I.S. implement cả MCP (Model Context Protocol) client lẫn server, cho phép plug tool ngoài mà không cần sửa core runtime:

• MCP server: expose các tool nội bộ của J.A.R.V.I.S. (truy vấn lịch, đọc file dự án) ra ngoài theo chuẩn, để các client MCP khác (ví dụ một IDE hỗ trợ MCP) có thể gọi trực tiếp.
• MCP client: cho phép J.A.R.V.I.S. kết nối tới MCP server của bên thứ ba (ví dụ một server quản lý tài liệu nội bộ công ty) mà không cần viết integration riêng cho từng hệ thống.

Kết quả của mỗi bước xử lý được stream về client qua Server-Sent Events (SSE) đi qua Fastify gateway, giữ latency hiển thị token đầu tiên dưới 300ms — quan trọng để trải nghiệm “đang suy nghĩ” cảm giác tự nhiên thay vì người dùng nhìn màn hình trắng chờ cả câu trả lời xong mới hiện.

Thiết kế modular này — multi-agent routing, tool execution song song, bộ nhớ phân tầng, và MCP chuẩn hoá giao tiếp — giúp thêm agent mới hoặc tool mới chỉ mất vài giờ, thay vì phải đụng vào runtime lõi. Tiêu chí mình chốt từ bản nháp đầu tiên: mọi phần mở rộng phải additive. Không invasive.