internal/mcp/mcp.go · Agent Time Ledger

302 lines · go

package mcp
 
import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"strings"
 
	"github.com/zionboggan/agent-time-ledger/internal/clock"
	"github.com/zionboggan/agent-time-ledger/internal/ledger"
)
 
type request struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      any             `json:"id,omitempty"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}
 
type response struct {
	JSONRPC string    `json:"jsonrpc"`
	ID      any       `json:"id,omitempty"`
	Result  any       `json:"result,omitempty"`
	Error   *rpcError `json:"error,omitempty"`
}
 
type message struct {
	Body   []byte
	Framed bool
}
 
type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}
 
type toolCallParams struct {
	Name      string          `json:"name"`
	Arguments json.RawMessage `json:"arguments,omitempty"`
}
 
type initializeParams struct {
	ProtocolVersion string `json:"protocolVersion"`
}
 
type callResult struct {
	Content []content `json:"content"`
	IsError bool      `json:"isError,omitempty"`
}
 
type content struct {
	Type string `json:"type"`
	Text string `json:"text"`
}
 
func Serve(stdin io.Reader, stdout, stderr io.Writer, service *ledger.Service) error {
	reader := bufio.NewReader(stdin)
	for {
		msg, err := readMessage(reader)
		if err == io.EOF {
			return nil
		}
		if err != nil {
			fmt.Fprintln(stderr, err)
			return err
		}
		if len(bytes.TrimSpace(msg.Body)) == 0 {
			continue
		}
 
		var req request
		if err := json.Unmarshal(msg.Body, &req); err != nil {
			_ = writeMessage(stdout, msg.Framed, response{
				JSONRPC: "2.0",
				Error:   &rpcError{Code: -32700, Message: "parse error"},
			})
			continue
		}
		if strings.HasPrefix(req.Method, "notifications/") {
			continue
		}
		resp := handle(req, service)
		if err := writeMessage(stdout, msg.Framed, resp); err != nil {
			return err
		}
	}
}
 
func handle(req request, service *ledger.Service) response {
	resp := response{JSONRPC: "2.0", ID: req.ID}
	switch req.Method {
	case "initialize":
		protocolVersion := "2024-11-05"
		var params initializeParams
		if len(req.Params) > 0 && json.Unmarshal(req.Params, &params) == nil && params.ProtocolVersion != "" {
			protocolVersion = params.ProtocolVersion
		}
		resp.Result = map[string]any{
			"protocolVersion": protocolVersion,
			"capabilities": map[string]any{
				"tools": map[string]any{"listChanged": false},
			},
			"serverInfo": map[string]any{
				"name":    "agent-time-ledger",
				"version": "0.1.0",
			},
		}
	case "ping":
		resp.Result = map[string]any{}
	case "tools/list":
		resp.Result = map[string]any{"tools": tools()}
	case "tools/call":
		result, err := handleToolCall(req.Params, service)
		if err != nil {
			resp.Result = callResult{
				Content: []content{{Type: "text", Text: err.Error()}},
				IsError: true,
			}
		} else {
			resp.Result = result
		}
	case "resources/list":
		resp.Result = map[string]any{"resources": []any{}}
	case "prompts/list":
		resp.Result = map[string]any{"prompts": []any{}}
	default:
		resp.Error = &rpcError{Code: -32601, Message: "method not found"}
	}
	return resp
}
 
func handleToolCall(raw json.RawMessage, service *ledger.Service) (callResult, error) {
	var params toolCallParams
	if err := json.Unmarshal(raw, &params); err != nil {
		return callResult{}, fmt.Errorf("invalid tool call params: %w", err)
	}
	args := map[string]any{}
	if len(params.Arguments) > 0 {
		if err := json.Unmarshal(params.Arguments, &args); err != nil {
			return callResult{}, fmt.Errorf("invalid tool arguments: %w", err)
		}
	}
 
	var result any
	var err error
	switch params.Name {
	case "time_now":
		timezone, _ := args["timezone"].(string)
		result, err = service.NowResponseIn(timezone)
	case "session_status":
		result, err = service.SessionStatus()
	case "mark_start":
		result, err = service.StartMark(requiredString(args, "name"))
	case "mark_elapsed":
		result, err = service.MarkElapsed(requiredString(args, "name"))
	case "stale_check":
		result, err = staleCheck(args, service)
	case "mark_list":
		result, err = service.ListMarks()
	case "mark_delete":
		name := requiredString(args, "name")
		err = service.DeleteMark(name)
		result = map[string]any{"ok": err == nil, "name": name, "confidence": clock.ConfidenceWallFallback}
	case "ledger_event":
		note, _ := args["note"].(string)
		err = service.LedgerEvent(note)
		result = map[string]any{"ok": err == nil, "confidence": clock.ConfidenceWallFallback}
	case "ledger_report":
		result, err = service.ReportToday()
	default:
		return callResult{}, fmt.Errorf("unknown tool %q", params.Name)
	}
	if err != nil {
		return callResult{}, err
	}
	text, err := json.MarshalIndent(result, "", "  ")
	if err != nil {
		return callResult{}, err
	}
	return callResult{Content: []content{{Type: "text", Text: string(text)}}}, nil
}
 
func staleCheck(args map[string]any, service *ledger.Service) (ledger.StaleResponse, error) {
	ttlText := requiredString(args, "ttl")
	ttl, err := clock.ParseDuration(ttlText)
	if err != nil {
		return ledger.StaleResponse{}, err
	}
	if markName, ok := args["mark"].(string); ok && markName != "" {
		return service.StaleFromMark(markName, ttl)
	}
	timestampText := requiredString(args, "timestamp")
	timestamp, err := clock.ParseRFC3339(timestampText)
	if err != nil {
		return ledger.StaleResponse{}, err
	}
	return service.StaleFromTimestamp(timestamp, ttl)
}
 
func requiredString(args map[string]any, name string) string {
	value, _ := args[name].(string)
	return value
}
 
func tools() []map[string]any {
	return []map[string]any{
		tool("time_now", "Return the current time with timezone and UTC offset.", map[string]any{
			"timezone": stringSchema("Optional IANA timezone, such as America/Chicago, America/New_York, or UTC"),
		}),
		tool("session_status", "Return active session status and elapsed time.", map[string]any{}),
		tool("mark_start", "Start or reset a named mark.", map[string]any{
			"name": stringSchema("Mark name"),
		}, "name"),
		tool("mark_elapsed", "Return elapsed time for a named mark.", map[string]any{
			"name": stringSchema("Mark name"),
		}, "name"),
		tool("mark_list", "Return all open marks with elapsed time.", map[string]any{}),
		tool("mark_delete", "Delete a named mark.", map[string]any{
			"name": stringSchema("Mark name"),
		}, "name"),
		tool("stale_check", "Check whether a timestamp or mark is stale for a TTL.", map[string]any{
			"timestamp": stringSchema("RFC3339 timestamp"),
			"mark":      stringSchema("Existing mark name"),
			"ttl":       stringSchema("Duration such as 15s, 30m, 2h, or 1d"),
		}, "ttl"),
		tool("ledger_event", "Append a local manual_note ledger event.", map[string]any{
			"note": stringSchema("Short local note"),
		}),
		tool("ledger_report", "Return today's local ledger report.", map[string]any{}),
	}
}
 
func tool(name, description string, properties map[string]any, required ...string) map[string]any {
	if required == nil {
		required = []string{}
	}
	return map[string]any{
		"name":        name,
		"description": description,
		"inputSchema": map[string]any{
			"type":                 "object",
			"properties":           properties,
			"required":             required,
			"additionalProperties": false,
		},
	}
}
 
func stringSchema(description string) map[string]any {
	return map[string]any{"type": "string", "description": description}
}
 
func readMessage(reader *bufio.Reader) (message, error) {
	line, err := reader.ReadString('\n')
	if err != nil {
		if err == io.EOF && strings.TrimSpace(line) != "" {
			return message{Body: []byte(line)}, nil
		}
		return message{}, err
	}
	line = strings.TrimRight(line, "\r\n")
	if strings.TrimSpace(line) == "" {
		return message{}, nil
	}
	if !strings.HasPrefix(strings.ToLower(line), "content-length:") {
		return message{Body: []byte(line)}, nil
	}
 
	lengthText := strings.TrimSpace(strings.TrimPrefix(strings.ToLower(line), "content-length:"))
	length, err := strconv.Atoi(lengthText)
	if err != nil {
		return message{}, fmt.Errorf("invalid Content-Length: %w", err)
	}
	for {
		header, err := reader.ReadString('\n')
		if err != nil {
			return message{}, err
		}
		if strings.TrimSpace(header) == "" {
			break
		}
	}
	body := make([]byte, length)
	_, err = io.ReadFull(reader, body)
	return message{Body: body, Framed: true}, err
}
 
func writeMessage(w io.Writer, framed bool, value any) error {
	data, err := json.Marshal(value)
	if err != nil {
		return err
	}
	if framed {
		_, err = fmt.Fprintf(w, "Content-Length: %d\r\n\r\n%s", len(data), data)
		return err
	}
	_, err = fmt.Fprintf(w, "%s\n", data)
	return err
}

1	package mcp
2
3	import (
4	"bufio"
5	"bytes"
6	"encoding/json"
7	"fmt"
8	"io"
9	"strconv"
10	"strings"
11
12	"github.com/zionboggan/agent-time-ledger/internal/clock"
13	"github.com/zionboggan/agent-time-ledger/internal/ledger"
14	)
15
16	type request struct {
17	JSONRPC string `json:"jsonrpc"`
18	ID any `json:"id,omitempty"`
19	Method string `json:"method"`
20	Params json.RawMessage `json:"params,omitempty"`
21	}
22
23	type response struct {
24	JSONRPC string `json:"jsonrpc"`
25	ID any `json:"id,omitempty"`
26	Result any `json:"result,omitempty"`
27	Error *rpcError `json:"error,omitempty"`
28	}
29
30	type message struct {
31	Body []byte
32	Framed bool
33	}
34
35	type rpcError struct {
36	Code int `json:"code"`
37	Message string `json:"message"`
38	}
39
40	type toolCallParams struct {
41	Name string `json:"name"`
42	Arguments json.RawMessage `json:"arguments,omitempty"`
43	}
44
45	type initializeParams struct {
46	ProtocolVersion string `json:"protocolVersion"`
47	}
48
49	type callResult struct {
50	Content []content `json:"content"`
51	IsError bool `json:"isError,omitempty"`
52	}
53
54	type content struct {
55	Type string `json:"type"`
56	Text string `json:"text"`
57	}
58
59	func Serve(stdin io.Reader, stdout, stderr io.Writer, service *ledger.Service) error {
60	reader := bufio.NewReader(stdin)
61	for {
62	msg, err := readMessage(reader)
63	if err == io.EOF {
64	return nil
65	}
66	if err != nil {
67	fmt.Fprintln(stderr, err)
68	return err
69	}
70	if len(bytes.TrimSpace(msg.Body)) == 0 {
71	continue
72	}
73
74	var req request
75	if err := json.Unmarshal(msg.Body, &req); err != nil {
76	_ = writeMessage(stdout, msg.Framed, response{
77	JSONRPC: "2.0",
78	Error: &rpcError{Code: -32700, Message: "parse error"},
79	})
80	continue
81	}
82	if strings.HasPrefix(req.Method, "notifications/") {
83	continue
84	}
85	resp := handle(req, service)
86	if err := writeMessage(stdout, msg.Framed, resp); err != nil {
87	return err
88	}
89	}
90	}
91
92	func handle(req request, service *ledger.Service) response {
93	resp := response{JSONRPC: "2.0", ID: req.ID}
94	switch req.Method {
95	case "initialize":
96	protocolVersion := "2024-11-05"
97	var params initializeParams
98	if len(req.Params) > 0 && json.Unmarshal(req.Params, &params) == nil && params.ProtocolVersion != "" {
99	protocolVersion = params.ProtocolVersion
100	}
101	resp.Result = map[string]any{
102	"protocolVersion": protocolVersion,
103	"capabilities": map[string]any{
104	"tools": map[string]any{"listChanged": false},
105	},
106	"serverInfo": map[string]any{
107	"name": "agent-time-ledger",
108	"version": "0.1.0",
109	},
110	}
111	case "ping":
112	resp.Result = map[string]any{}
113	case "tools/list":
114	resp.Result = map[string]any{"tools": tools()}
115	case "tools/call":
116	result, err := handleToolCall(req.Params, service)
117	if err != nil {
118	resp.Result = callResult{
119	Content: []content{{Type: "text", Text: err.Error()}},
120	IsError: true,
121	}
122	} else {
123	resp.Result = result
124	}
125	case "resources/list":
126	resp.Result = map[string]any{"resources": []any{}}
127	case "prompts/list":
128	resp.Result = map[string]any{"prompts": []any{}}
129	default:
130	resp.Error = &rpcError{Code: -32601, Message: "method not found"}
131	}
132	return resp
133	}
134
135	func handleToolCall(raw json.RawMessage, service *ledger.Service) (callResult, error) {
136	var params toolCallParams
137	if err := json.Unmarshal(raw, &params); err != nil {
138	return callResult{}, fmt.Errorf("invalid tool call params: %w", err)
139	}
140	args := map[string]any{}
141	if len(params.Arguments) > 0 {
142	if err := json.Unmarshal(params.Arguments, &args); err != nil {
143	return callResult{}, fmt.Errorf("invalid tool arguments: %w", err)
144	}
145	}
146
147	var result any
148	var err error
149	switch params.Name {
150	case "time_now":
151	timezone, _ := args["timezone"].(string)
152	result, err = service.NowResponseIn(timezone)
153	case "session_status":
154	result, err = service.SessionStatus()
155	case "mark_start":
156	result, err = service.StartMark(requiredString(args, "name"))
157	case "mark_elapsed":
158	result, err = service.MarkElapsed(requiredString(args, "name"))
159	case "stale_check":
160	result, err = staleCheck(args, service)
161	case "mark_list":
162	result, err = service.ListMarks()
163	case "mark_delete":
164	name := requiredString(args, "name")
165	err = service.DeleteMark(name)
166	result = map[string]any{"ok": err == nil, "name": name, "confidence": clock.ConfidenceWallFallback}
167	case "ledger_event":
168	note, _ := args["note"].(string)
169	err = service.LedgerEvent(note)
170	result = map[string]any{"ok": err == nil, "confidence": clock.ConfidenceWallFallback}
171	case "ledger_report":
172	result, err = service.ReportToday()
173	default:
174	return callResult{}, fmt.Errorf("unknown tool %q", params.Name)
175	}
176	if err != nil {
177	return callResult{}, err
178	}
179	text, err := json.MarshalIndent(result, "", " ")
180	if err != nil {
181	return callResult{}, err
182	}
183	return callResult{Content: []content{{Type: "text", Text: string(text)}}}, nil
184	}
185
186	func staleCheck(args map[string]any, service *ledger.Service) (ledger.StaleResponse, error) {
187	ttlText := requiredString(args, "ttl")
188	ttl, err := clock.ParseDuration(ttlText)
189	if err != nil {
190	return ledger.StaleResponse{}, err
191	}
192	if markName, ok := args["mark"].(string); ok && markName != "" {
193	return service.StaleFromMark(markName, ttl)
194	}
195	timestampText := requiredString(args, "timestamp")
196	timestamp, err := clock.ParseRFC3339(timestampText)
197	if err != nil {
198	return ledger.StaleResponse{}, err
199	}
200	return service.StaleFromTimestamp(timestamp, ttl)
201	}
202
203	func requiredString(args map[string]any, name string) string {
204	value, _ := args[name].(string)
205	return value
206	}
207
208	func tools() []map[string]any {
209	return []map[string]any{
210	tool("time_now", "Return the current time with timezone and UTC offset.", map[string]any{
211	"timezone": stringSchema("Optional IANA timezone, such as America/Chicago, America/New_York, or UTC"),
212	}),
213	tool("session_status", "Return active session status and elapsed time.", map[string]any{}),
214	tool("mark_start", "Start or reset a named mark.", map[string]any{
215	"name": stringSchema("Mark name"),
216	}, "name"),
217	tool("mark_elapsed", "Return elapsed time for a named mark.", map[string]any{
218	"name": stringSchema("Mark name"),
219	}, "name"),
220	tool("mark_list", "Return all open marks with elapsed time.", map[string]any{}),
221	tool("mark_delete", "Delete a named mark.", map[string]any{
222	"name": stringSchema("Mark name"),
223	}, "name"),
224	tool("stale_check", "Check whether a timestamp or mark is stale for a TTL.", map[string]any{
225	"timestamp": stringSchema("RFC3339 timestamp"),
226	"mark": stringSchema("Existing mark name"),
227	"ttl": stringSchema("Duration such as 15s, 30m, 2h, or 1d"),
228	}, "ttl"),
229	tool("ledger_event", "Append a local manual_note ledger event.", map[string]any{
230	"note": stringSchema("Short local note"),
231	}),
232	tool("ledger_report", "Return today's local ledger report.", map[string]any{}),
233	}
234	}
235
236	func tool(name, description string, properties map[string]any, required ...string) map[string]any {
237	if required == nil {
238	required = []string{}
239	}
240	return map[string]any{
241	"name": name,
242	"description": description,
243	"inputSchema": map[string]any{
244	"type": "object",
245	"properties": properties,
246	"required": required,
247	"additionalProperties": false,
248	},
249	}
250	}
251
252	func stringSchema(description string) map[string]any {
253	return map[string]any{"type": "string", "description": description}
254	}
255
256	func readMessage(reader *bufio.Reader) (message, error) {
257	line, err := reader.ReadString('\n')
258	if err != nil {
259	if err == io.EOF && strings.TrimSpace(line) != "" {
260	return message{Body: []byte(line)}, nil
261	}
262	return message{}, err
263	}
264	line = strings.TrimRight(line, "\r\n")
265	if strings.TrimSpace(line) == "" {
266	return message{}, nil
267	}
268	if !strings.HasPrefix(strings.ToLower(line), "content-length:") {
269	return message{Body: []byte(line)}, nil
270	}
271
272	lengthText := strings.TrimSpace(strings.TrimPrefix(strings.ToLower(line), "content-length:"))
273	length, err := strconv.Atoi(lengthText)
274	if err != nil {
275	return message{}, fmt.Errorf("invalid Content-Length: %w", err)
276	}
277	for {
278	header, err := reader.ReadString('\n')
279	if err != nil {
280	return message{}, err
281	}
282	if strings.TrimSpace(header) == "" {
283	break
284	}
285	}
286	body := make([]byte, length)
287	_, err = io.ReadFull(reader, body)
288	return message{Body: body, Framed: true}, err
289	}
290
291	func writeMessage(w io.Writer, framed bool, value any) error {
292	data, err := json.Marshal(value)
293	if err != nil {
294	return err
295	}
296	if framed {
297	_, err = fmt.Fprintf(w, "Content-Length: %d\r\n\r\n%s", len(data), data)
298	return err
299	}
300	_, err = fmt.Fprintf(w, "%s\n", data)
301	return err
302	}