26285265118

Committed 22 May 2026 11:31AM UTC coverage: 51.35% (+5.7%) from 45.617%

Build # 26285265118

Build Type

push

github

Committed by

k0d3r1s

Commit Message

Ignores local Claude memory files

The local memory server creates files that are not meant for version control. These rules ensure they stay out of the repository.

Coverage Stats

2909 of 5665 relevant lines covered (51.35%)

75.18 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

83.96

/internal/instance/client.go

package instance

import (
        "bufio"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "net"
        "time"

        "github.com/valksor/go-assern/internal/aggregator"
)

// ClientTimeout is the default timeout for client operations.
const ClientTimeout = 10 * time.Second

// ToolInfo represents tool information returned from a query.
type ToolInfo struct {
        Name        string          `json:"name"`
        Description string          `json:"description"`
        InputSchema json.RawMessage `json:"inputSchema,omitempty"`
}

// ListResult contains the result of querying a running instance.
type ListResult struct {
        Tools          []ToolInfo
        TokensByServer map[string]int
        TotalTokens    int
}

// Client connects to a running assern instance to query information.
type Client struct {
        socketPath string
        conn       net.Conn
        reader     *bufio.Reader
        requestID  int
}

// NewClient creates a new client for the given socket path.
func NewClient(socketPath string) *Client {
        return &Client{
                socketPath: socketPath,
                requestID:  0,
        }
}

// Connect establishes connection to the instance.
func (c *Client) Connect(ctx context.Context) error {
        var dialer net.Dialer
        conn, err := dialer.DialContext(ctx, "unix", c.socketPath)
        if err != nil {
                return fmt.Errorf("connect to socket: %w", err)
        }

        c.conn = conn
        c.reader = bufio.NewReader(conn)

        return nil
}

// Close closes the connection.
func (c *Client) Close() error {
        if c.conn != nil {
                return c.conn.Close()
        }

        return nil
}

// Initialize performs the MCP initialize handshake.
func (c *Client) Initialize(ctx context.Context) error {
        // Wait for handshake timeout to pass (server expects internal commands first)
        time.Sleep(handshakeTimeout + 10*time.Millisecond)

        c.requestID++
        initReq := map[string]any{
                keyJSONRPC: jsonrpcVersion,
                "id":       c.requestID,
                keyMethod:  "initialize",
                "params": map[string]any{
                        "protocolVersion": "2024-11-05",
                        "capabilities":    map[string]any{},
                        "clientInfo": map[string]any{
                                "name":    "assern-client",
                                "version": "1.0.0",
                        },
                },
        }

        if err := c.sendRequest(initReq); err != nil {
                return fmt.Errorf("send initialize: %w", err)
        }

        // Read initialize response
        var initResp struct {
                ID     int `json:"id"`
                Result any `json:"result"`
                Error  *struct {
                        Code    int    `json:"code"`
                        Message string `json:"message"`
                } `json:"error"`
        }

        if err := c.readResponse(&initResp); err != nil {
                return fmt.Errorf("read initialize response: %w", err)
        }

        if initResp.Error != nil {
                return fmt.Errorf("initialize error: %s", initResp.Error.Message)
        }

        // Send initialized notification
        initializedNotif := map[string]any{
                keyJSONRPC: jsonrpcVersion,
                keyMethod:  "notifications/initialized",
        }

        if err := c.sendRequest(initializedNotif); err != nil {
                return fmt.Errorf("send initialized notification: %w", err)
        }

        return nil
}

// ListTools queries the available tools from the running instance.
func (c *Client) ListTools(ctx context.Context) (*ListResult, error) {
        c.requestID++
        listReq := map[string]any{
                keyJSONRPC: jsonrpcVersion,
                "id":       c.requestID,
                keyMethod:  "tools/list",
                "params":   map[string]any{},
        }

        if err := c.sendRequest(listReq); err != nil {
                return nil, fmt.Errorf("send tools/list: %w", err)
        }

        var resp struct {
                ID     int `json:"id"`
                Result struct {
                        Tools []ToolInfo `json:"tools"`
                } `json:"result"`
                Error *struct {
                        Code    int    `json:"code"`
                        Message string `json:"message"`
                } `json:"error"`
        }

        if err := c.readResponse(&resp); err != nil {
                return nil, fmt.Errorf("read tools/list response: %w", err)
        }

        if resp.Error != nil {
                return nil, fmt.Errorf("tools/list error: %s", resp.Error.Message)
        }

        tokensByServer, totalTokens := estimateListTokens(resp.Result.Tools)

        return &ListResult{
                Tools:          resp.Result.Tools,
                TokensByServer: tokensByServer,
                TotalTokens:    totalTokens,
        }, nil
}

// estimateListTokens groups the estimated token cost of tool definitions by
// server, deriving the server from the tool's prefix (server_tool). Tools
// without a parseable prefix are bucketed under their own name.
func estimateListTokens(tools []ToolInfo) (map[string]int, int) {
        byServer := make(map[string]int)
        total := 0

        for _, tool := range tools {
                cost := aggregator.EstimateRawToolTokens(tool.Name, tool.Description, tool.InputSchema)

                server, _, err := aggregator.ParsePrefixedName(tool.Name)
                if err != nil {
                        // Names without a server prefix (e.g. the assern_* meta-tools have
                        // one, but a truly unprefixed name would not) go in one bucket
                        // rather than inventing a phantom server per tool.
                        server = "(unprefixed)"
                }

                byServer[server] += cost
                total += cost
        }

        return byServer, total
}

func (c *Client) sendRequest(req any) error {
        data, err := json.Marshal(req)
        if err != nil {
                return err
        }

        data = append(data, '\n')

        if _, err := c.conn.Write(data); err != nil {
                return err
        }

        return nil
}

func (c *Client) readResponse(resp any) error {
        if err := c.conn.SetReadDeadline(time.Now().Add(ClientTimeout)); err != nil {
                return err
        }
        defer func() { _ = c.conn.SetReadDeadline(time.Time{}) }()

        line, err := c.reader.ReadBytes('\n')
        if err != nil {
                return err
        }

        return json.Unmarshal(line, resp)
}

// QueryTools connects to a running instance and returns the available tools.
// This is a convenience function that handles the full connection lifecycle.
func QueryTools(ctx context.Context, socketPath string) (*ListResult, error) {
        client := NewClient(socketPath)

        if err := client.Connect(ctx); err != nil {
                return nil, err
        }
        defer func() { _ = client.Close() }()

        if err := client.Initialize(ctx); err != nil {
                return nil, err
        }

        return client.ListTools(ctx)
}

// ReloadResult contains the result of a reload operation.
type ReloadResult struct {
        Added   int      `json:"added"`
        Removed int      `json:"removed"`
        Errors  []string `json:"errors,omitempty"`
}

// Reload triggers a configuration reload on a running instance.
// This uses the internal command protocol (not MCP).
func Reload(ctx context.Context, socketPath string) (*ReloadResult, error) {
        var dialer net.Dialer
        conn, err := dialer.DialContext(ctx, "unix", socketPath)
        if err != nil {
                return nil, fmt.Errorf("connect to socket: %w", err)
        }
        defer func() { _ = conn.Close() }()

        // Send reload request
        reloadReq := map[string]any{
                keyJSONRPC: jsonrpcVersion,
                "id":       1,
                keyMethod:  "assern/reload",
        }
        if err := json.NewEncoder(conn).Encode(reloadReq); err != nil {
                return nil, fmt.Errorf("send reload request: %w", err)
        }

        // Set read deadline
        if err := conn.SetReadDeadline(time.Now().Add(ClientTimeout)); err != nil {
                return nil, fmt.Errorf("set read deadline: %w", err)
        }

        // Read response
        var resp struct {
                Result *ReloadResult `json:"result"`
                Error  *struct {
                        Code    int    `json:"code"`
                        Message string `json:"message"`
                } `json:"error"`
        }
        if err := json.NewDecoder(conn).Decode(&resp); err != nil {
                return nil, fmt.Errorf("read reload response: %w", err)
        }

        if resp.Error != nil {
                return nil, fmt.Errorf("reload error: %s", resp.Error.Message)
        }

        if resp.Result == nil {
                return nil, errors.New("empty reload response")
        }

        return resp.Result, nil
}

1	package instance
2
3	import (
4	"bufio"
5	"context"
6	"encoding/json"
7	"errors"
8	"fmt"
9	"net"
10	"time"
11
12	"github.com/valksor/go-assern/internal/aggregator"
13	)
14
15	// ClientTimeout is the default timeout for client operations.
16	const ClientTimeout = 10 * time.Second
17
18	// ToolInfo represents tool information returned from a query.
19	type ToolInfo struct {
20	Name string `json:"name"`
21	Description string `json:"description"`
22	InputSchema json.RawMessage `json:"inputSchema,omitempty"`
23	}
24
25	// ListResult contains the result of querying a running instance.
26	type ListResult struct {
27	Tools []ToolInfo
28	TokensByServer map[string]int
29	TotalTokens int
30	}
31
32	// Client connects to a running assern instance to query information.
33	type Client struct {
34	socketPath string
35	conn net.Conn
36	reader *bufio.Reader
37	requestID int
38	}
39
40	// NewClient creates a new client for the given socket path.
41	func NewClient(socketPath string) *Client {	11✔
42	return &Client{	11✔
43	socketPath: socketPath,	11✔
44	requestID: 0,	11✔
45	}	11✔
46	}	11✔
47
48	// Connect establishes connection to the instance.
49	func (c *Client) Connect(ctx context.Context) error {	9✔
50	var dialer net.Dialer	9✔
51	conn, err := dialer.DialContext(ctx, "unix", c.socketPath)	9✔
52	if err != nil {	11✔
53	return fmt.Errorf("connect to socket: %w", err)	2✔
54	}	2✔
55
56	c.conn = conn	7✔
57	c.reader = bufio.NewReader(conn)	7✔
58		7✔
59	return nil	7✔
60	}
61
62	// Close closes the connection.
63	func (c *Client) Close() error {	8✔
64	if c.conn != nil {	15✔
65	return c.conn.Close()	7✔
66	}	7✔
67
68	return nil	1✔
69	}
70
71	// Initialize performs the MCP initialize handshake.
72	func (c *Client) Initialize(ctx context.Context) error {	6✔
73	// Wait for handshake timeout to pass (server expects internal commands first)	6✔
74	time.Sleep(handshakeTimeout + 10*time.Millisecond)	6✔
75		6✔
76	c.requestID++	6✔
77	initReq := map[string]any{	6✔
78	keyJSONRPC: jsonrpcVersion,	6✔
79	"id": c.requestID,	6✔
80	keyMethod: "initialize",	6✔
81	"params": map[string]any{	6✔
82	"protocolVersion": "2024-11-05",	6✔
83	"capabilities": map[string]any{},	6✔
84	"clientInfo": map[string]any{	6✔
85	"name": "assern-client",	6✔
86	"version": "1.0.0",	6✔
87	},	6✔
88	},	6✔
89	}	6✔
90		6✔
91	if err := c.sendRequest(initReq); err != nil {	6✔
92	return fmt.Errorf("send initialize: %w", err)	×
93	}	×
94
95	// Read initialize response
96	var initResp struct {	6✔
97	ID int `json:"id"`	6✔
98	Result any `json:"result"`	6✔
99	Error *struct {	6✔
100	Code int `json:"code"`	6✔
101	Message string `json:"message"`	6✔
102	} `json:"error"`	6✔
103	}	6✔
104		6✔
105	if err := c.readResponse(&initResp); err != nil {	6✔
106	return fmt.Errorf("read initialize response: %w", err)	×
107	}	×
108
109	if initResp.Error != nil {	6✔
110	return fmt.Errorf("initialize error: %s", initResp.Error.Message)	×
111	}	×
112
113	// Send initialized notification
114	initializedNotif := map[string]any{	6✔
115	keyJSONRPC: jsonrpcVersion,	6✔
116	keyMethod: "notifications/initialized",	6✔
117	}	6✔
118		6✔
119	if err := c.sendRequest(initializedNotif); err != nil {	6✔
120	return fmt.Errorf("send initialized notification: %w", err)	×
121	}	×
122
123	return nil	6✔
124	}
125
126	// ListTools queries the available tools from the running instance.
127	func (c Client) ListTools(ctx context.Context) (ListResult, error) {	5✔
128	c.requestID++	5✔
129	listReq := map[string]any{	5✔
130	keyJSONRPC: jsonrpcVersion,	5✔
131	"id": c.requestID,	5✔
132	keyMethod: "tools/list",	5✔
133	"params": map[string]any{},	5✔
134	}	5✔
135		5✔
136	if err := c.sendRequest(listReq); err != nil {	5✔
137	return nil, fmt.Errorf("send tools/list: %w", err)	×
138	}	×
139
140	var resp struct {	5✔
141	ID int `json:"id"`	5✔
142	Result struct {	5✔
143	Tools []ToolInfo `json:"tools"`	5✔
144	} `json:"result"`	5✔
145	Error *struct {	5✔
146	Code int `json:"code"`	5✔
147	Message string `json:"message"`	5✔
148	} `json:"error"`	5✔
149	}	5✔
150		5✔
151	if err := c.readResponse(&resp); err != nil {	5✔
152	return nil, fmt.Errorf("read tools/list response: %w", err)	×
153	}	×
154
155	if resp.Error != nil {	5✔
156	return nil, fmt.Errorf("tools/list error: %s", resp.Error.Message)	×
157	}	×
158
159	tokensByServer, totalTokens := estimateListTokens(resp.Result.Tools)	5✔
160		5✔
161	return &ListResult{	5✔
162	Tools: resp.Result.Tools,	5✔
163	TokensByServer: tokensByServer,	5✔
164	TotalTokens: totalTokens,	5✔
165	}, nil	5✔
166	}
167
168	// estimateListTokens groups the estimated token cost of tool definitions by
169	// server, deriving the server from the tool's prefix (server_tool). Tools
170	// without a parseable prefix are bucketed under their own name.
171	func estimateListTokens(tools []ToolInfo) (map[string]int, int) {	7✔
172	byServer := make(map[string]int)	7✔
173	total := 0	7✔
174		7✔
175	for _, tool := range tools {	66✔
176	cost := aggregator.EstimateRawToolTokens(tool.Name, tool.Description, tool.InputSchema)	59✔
177		59✔
178	server, _, err := aggregator.ParsePrefixedName(tool.Name)	59✔
179	if err != nil {	60✔
180	// Names without a server prefix (e.g. the assern_* meta-tools have	1✔
181	// one, but a truly unprefixed name would not) go in one bucket	1✔
182	// rather than inventing a phantom server per tool.	1✔
183	server = "(unprefixed)"	1✔
184	}	1✔
185
186	byServer[server] += cost	59✔
187	total += cost	59✔
188	}
189
190	return byServer, total	7✔
191	}
192
193	func (c *Client) sendRequest(req any) error {	17✔
194	data, err := json.Marshal(req)	17✔
195	if err != nil {	17✔
196	return err	×
197	}	×
198
199	data = append(data, '\n')	17✔
200		17✔
201	if _, err := c.conn.Write(data); err != nil {	17✔
202	return err	×
203	}	×
204
205	return nil	17✔
206	}
207
208	func (c *Client) readResponse(resp any) error {	11✔
209	if err := c.conn.SetReadDeadline(time.Now().Add(ClientTimeout)); err != nil {	11✔
210	return err	×
211	}	×
212	defer func() { _ = c.conn.SetReadDeadline(time.Time{}) }()	22✔
213
214	line, err := c.reader.ReadBytes('\n')	11✔
215	if err != nil {	11✔
216	return err	×
217	}	×
218
219	return json.Unmarshal(line, resp)	11✔
220	}
221
222	// QueryTools connects to a running instance and returns the available tools.
223	// This is a convenience function that handles the full connection lifecycle.
224	func QueryTools(ctx context.Context, socketPath string) (*ListResult, error) {	7✔
225	client := NewClient(socketPath)	7✔
226		7✔
227	if err := client.Connect(ctx); err != nil {	9✔
228	return nil, err	2✔
229	}	2✔
230	defer func() { _ = client.Close() }()	10✔
231
232	if err := client.Initialize(ctx); err != nil {	5✔
233	return nil, err	×
234	}	×
235
236	return client.ListTools(ctx)	5✔
237	}
238
239	// ReloadResult contains the result of a reload operation.
240	type ReloadResult struct {
241	Added int `json:"added"`
242	Removed int `json:"removed"`
243	Errors []string `json:"errors,omitempty"`
244	}
245
246	// Reload triggers a configuration reload on a running instance.
247	// This uses the internal command protocol (not MCP).
248	func Reload(ctx context.Context, socketPath string) (*ReloadResult, error) {	4✔
249	var dialer net.Dialer	4✔
250	conn, err := dialer.DialContext(ctx, "unix", socketPath)	4✔
251	if err != nil {	5✔
252	return nil, fmt.Errorf("connect to socket: %w", err)	1✔
253	}	1✔
254	defer func() { _ = conn.Close() }()	6✔
255
256	// Send reload request
257	reloadReq := map[string]any{	3✔
258	keyJSONRPC: jsonrpcVersion,	3✔
259	"id": 1,	3✔
260	keyMethod: "assern/reload",	3✔
261	}	3✔
262	if err := json.NewEncoder(conn).Encode(reloadReq); err != nil {	3✔
263	return nil, fmt.Errorf("send reload request: %w", err)	×
264	}	×
265
266	// Set read deadline
267	if err := conn.SetReadDeadline(time.Now().Add(ClientTimeout)); err != nil {	3✔
268	return nil, fmt.Errorf("set read deadline: %w", err)	×
269	}	×
270
271	// Read response
272	var resp struct {	3✔
273	Result *ReloadResult `json:"result"`	3✔
274	Error *struct {	3✔
275	Code int `json:"code"`	3✔
276	Message string `json:"message"`	3✔
277	} `json:"error"`	3✔
278	}	3✔
279	if err := json.NewDecoder(conn).Decode(&resp); err != nil {	4✔
280	return nil, fmt.Errorf("read reload response: %w", err)	1✔
281	}	1✔
282
283	if resp.Error != nil {	3✔
284	return nil, fmt.Errorf("reload error: %s", resp.Error.Message)	1✔
285	}	1✔
286
287	if resp.Result == nil {	1✔
288	return nil, errors.New("empty reload response")	×
289	}	×
290
291	return resp.Result, nil	1✔
292	}

valksor / go-assern / 26285265118

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous