Skip to content

Commit d9c367d

Browse files
djthorpedjthorpe
authored
Adding a heartbeat MCP server (connector) (#50)
* Updates * Updated heartbeats * Updates * Updated * Added README * Updated * Updated * Updates after PR comments * Updates * Updates after PR comments * Updating tests * Updated tool coverage * Updated after PR comments * Updates
1 parent 294ef78 commit d9c367d

78 files changed

Lines changed: 5653 additions & 1279 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎cmd/llm/agent.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,16 @@ import (
1212
"strings"
1313

1414
// Packages
15-
server "github.com/mutablelogic/go-server"
1615
otel "github.com/mutablelogic/go-client/pkg/otel"
1716
agent "github.com/mutablelogic/go-llm/pkg/agent"
1817
httpclient "github.com/mutablelogic/go-llm/pkg/httpclient"
1918
opt "github.com/mutablelogic/go-llm/pkg/opt"
2019
schema "github.com/mutablelogic/go-llm/pkg/schema"
2120
uitable "github.com/mutablelogic/go-llm/pkg/ui/table"
21+
server "github.com/mutablelogic/go-server"
2222
httpresponse "github.com/mutablelogic/go-server/pkg/httpresponse"
2323
types "github.com/mutablelogic/go-server/pkg/types"
24-
"go.opentelemetry.io/otel/attribute"
24+
attribute "go.opentelemetry.io/otel/attribute"
2525
)
2626

2727
///////////////////////////////////////////////////////////////////////////////

‎cmd/llm/credentials.go‎

Lines changed: 54 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,12 @@ import (
99
"time"
1010

1111
// Packages
12-
server "github.com/mutablelogic/go-server"
1312
goclient "github.com/mutablelogic/go-client"
1413
oauth "github.com/mutablelogic/go-client/pkg/oauth"
1514
otel "github.com/mutablelogic/go-client/pkg/otel"
1615
mcpclient "github.com/mutablelogic/go-llm/pkg/mcp/client"
1716
schema "github.com/mutablelogic/go-llm/pkg/schema"
17+
server "github.com/mutablelogic/go-server"
1818
types "github.com/mutablelogic/go-server/pkg/types"
1919
attribute "go.opentelemetry.io/otel/attribute"
2020
)
@@ -65,62 +65,64 @@ func (cmd *LoginCommand) Run(g server.Cmd) (err error) {
6565
var captured *oauth.OAuthCredentials
6666

6767
var c *mcpclient.Client
68-
c, err = mcpclient.New(cmd.URL, cmd.ClientName, "0", func(ctx context.Context, discoveryURL string) error {
69-
flow := c.OAuth()
70-
71-
metadata, err := flow.Discover(ctx, discoveryURL)
72-
if err != nil {
73-
return fmt.Errorf("discover: %w", err)
74-
}
68+
c, err = mcpclient.New(cmd.URL, cmd.ClientName, "0",
69+
mcpclient.WithAuth(func(ctx context.Context, discoveryURL string) error {
70+
flow := c.OAuth()
7571

76-
switch {
77-
case cmd.ClientCredentials:
78-
// Machine-to-machine: client ID+secret required.
79-
if cmd.ClientID == "" {
80-
return fmt.Errorf("--client-id is required for --client-credentials flow")
81-
}
82-
if cmd.ClientSecret == "" {
83-
return fmt.Errorf("--client-secret is required for --client-credentials flow")
84-
}
85-
creds := &oauth.OAuthCredentials{
86-
ClientID: cmd.ClientID,
87-
ClientSecret: cmd.ClientSecret,
88-
TokenURL: metadata.TokenEndpoint,
89-
Metadata: metadata,
90-
}
91-
captured, err = flow.AuthorizeWithCredentials(ctx, creds, cmd.Scopes...)
92-
93-
case cmd.Device:
94-
// Device authorization grant (RFC 8628).
95-
creds, err := loginCreds(ctx, flow, metadata, cmd.ClientID, cmd.ClientSecret, cmd.ClientName)
96-
if err != nil {
97-
return err
98-
}
99-
captured, err = flow.AuthorizeWithDevice(ctx, creds, func(userCode, verificationURI string) error {
100-
fmt.Printf("Visit %s and enter code: %s\n", verificationURI, userCode)
101-
return nil
102-
}, cmd.Scopes...)
103-
104-
default:
105-
// Browser-based Authorization Code + PKCE via loopback redirect.
106-
listener, err := net.Listen("tcp", "127.0.0.1:0")
72+
metadata, err := flow.Discover(ctx, discoveryURL)
10773
if err != nil {
108-
return fmt.Errorf("start callback listener: %w", err)
74+
return fmt.Errorf("discover: %w", err)
10975
}
110-
defer listener.Close()
111-
redirectURI := fmt.Sprintf("http://127.0.0.1:%d/callback", listener.Addr().(*net.TCPAddr).Port)
11276

113-
creds, err := loginBrowserCreds(ctx, flow, metadata, cmd.ClientID, cmd.ClientSecret, cmd.ClientName, redirectURI)
114-
if err != nil {
115-
return err
77+
switch {
78+
case cmd.ClientCredentials:
79+
// Machine-to-machine: client ID+secret required.
80+
if cmd.ClientID == "" {
81+
return fmt.Errorf("--client-id is required for --client-credentials flow")
82+
}
83+
if cmd.ClientSecret == "" {
84+
return fmt.Errorf("--client-secret is required for --client-credentials flow")
85+
}
86+
creds := &oauth.OAuthCredentials{
87+
ClientID: cmd.ClientID,
88+
ClientSecret: cmd.ClientSecret,
89+
TokenURL: metadata.TokenEndpoint,
90+
Metadata: metadata,
91+
}
92+
captured, err = flow.AuthorizeWithCredentials(ctx, creds, cmd.Scopes...)
93+
94+
case cmd.Device:
95+
// Device authorization grant (RFC 8628).
96+
creds, err := loginCreds(ctx, flow, metadata, cmd.ClientID, cmd.ClientSecret, cmd.ClientName)
97+
if err != nil {
98+
return err
99+
}
100+
captured, err = flow.AuthorizeWithDevice(ctx, creds, func(userCode, verificationURI string) error {
101+
fmt.Printf("Visit %s and enter code: %s\n", verificationURI, userCode)
102+
return nil
103+
}, cmd.Scopes...)
104+
105+
default:
106+
// Browser-based Authorization Code + PKCE via loopback redirect.
107+
listener, err := net.Listen("tcp", "127.0.0.1:0")
108+
if err != nil {
109+
return fmt.Errorf("start callback listener: %w", err)
110+
}
111+
defer listener.Close()
112+
redirectURI := fmt.Sprintf("http://127.0.0.1:%d/callback", listener.Addr().(*net.TCPAddr).Port)
113+
114+
creds, err := loginBrowserCreds(ctx, flow, metadata, cmd.ClientID, cmd.ClientSecret, cmd.ClientName, redirectURI)
115+
if err != nil {
116+
return err
117+
}
118+
captured, err = flow.AuthorizeWithBrowser(ctx, creds, listener, func(authURL string) error {
119+
fmt.Printf("Opening browser for authorization...\n%s\n", authURL)
120+
return openBrowser(authURL)
121+
}, cmd.Scopes...)
116122
}
117-
captured, err = flow.AuthorizeWithBrowser(ctx, creds, listener, func(authURL string) error {
118-
fmt.Printf("Opening browser for authorization...\n%s\n", authURL)
119-
return openBrowser(authURL)
120-
}, cmd.Scopes...)
121-
}
122-
return err
123-
}, opts...)
123+
return err
124+
}),
125+
mcpclient.WithClientOpt(opts...))
124126
if err != nil {
125127
return fmt.Errorf("create client: %w", err)
126128
}

‎cmd/llm/main.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ type CLI struct {
1717
ConnectorCommands
1818
CredentialsCommands
1919
GenerateCommands
20+
MCPCommands
2021
ModelCommands
2122
SessionCommands
2223
ToolCommands

‎cmd/llm/mcp.go‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
//go:build !client
2+
3+
package main
4+
5+
import (
6+
"context"
7+
"encoding/json"
8+
"fmt"
9+
"net/url"
10+
"os"
11+
"path/filepath"
12+
13+
// Packages
14+
heartbeat "github.com/mutablelogic/go-llm/pkg/heartbeat"
15+
mcpserver "github.com/mutablelogic/go-llm/pkg/mcp/server"
16+
server "github.com/mutablelogic/go-server"
17+
gocmd "github.com/mutablelogic/go-server/pkg/cmd"
18+
httprouter "github.com/mutablelogic/go-server/pkg/httprouter"
19+
version "github.com/mutablelogic/go-server/pkg/version"
20+
)
21+
22+
///////////////////////////////////////////////////////////////////////////////
23+
// TYPES
24+
25+
type MCPCommands struct {
26+
HeartbeatMCP HeartbeatMCPCommand `cmd:"" name:"heartbeat" help:"Run the heartbeat MCP server." group:"SERVER"`
27+
}
28+
29+
type HeartbeatMCPCommand struct {
30+
gocmd.RunServer
31+
}
32+
33+
///////////////////////////////////////////////////////////////////////////////
34+
// COMMANDS
35+
36+
func (cmd *HeartbeatMCPCommand) Run(ctx server.Cmd) error {
37+
// Create the file-backed heartbeat store
38+
store, err := cmd.HeartbeatStore(ctx.Name())
39+
if err != nil {
40+
return err
41+
}
42+
43+
// Create the MCP server
44+
srv, err := mcpserver.New("heartbeat", version.Version())
45+
if err != nil {
46+
return fmt.Errorf("mcp server: %w", err)
47+
}
48+
49+
// Create the heartbeat manager; on each fire, upsert a resource so connected
50+
// clients receive notifications/resources/list_changed automatically.
51+
mgr, err := heartbeat.New(store,
52+
heartbeat.WithLogger(ctx.Logger()),
53+
heartbeat.WithOnFire(func(_ context.Context, h *heartbeat.Heartbeat) {
54+
u, _ := url.Parse("heartbeat:" + h.ID)
55+
raw, _ := json.Marshal(h)
56+
srv.AddResources(&heartbeatResource{
57+
uri: u.String(),
58+
name: h.Message,
59+
data: raw,
60+
}) //nolint:errcheck
61+
}))
62+
if err != nil {
63+
return fmt.Errorf("heartbeat manager: %w", err)
64+
}
65+
66+
// Register the heartbeat tools with the MCP server
67+
tools, err := mgr.ListTools(ctx.Context())
68+
if err != nil {
69+
return fmt.Errorf("listing heartbeat tools: %w", err)
70+
}
71+
if err := srv.AddTools(tools...); err != nil {
72+
return fmt.Errorf("registering heartbeat tools: %w", err)
73+
}
74+
75+
// Mount the MCP handler on the router at the HTTP prefix
76+
cmd.Register(func(router *httprouter.Router, _ server.Cmd) error {
77+
return router.RegisterFunc("", srv.Handler().ServeHTTP, false, nil)
78+
})
79+
80+
// Run the heartbeat manager alongside the HTTP server
81+
go func() {
82+
if err := mgr.Run(ctx.Context()); err != nil {
83+
ctx.Logger().ErrorContext(ctx.Context(), "heartbeat manager error", "error", err)
84+
}
85+
}()
86+
87+
return cmd.RunServer.Run(ctx)
88+
}
89+
90+
///////////////////////////////////////////////////////////////////////////////
91+
// PRIVATE TYPES
92+
93+
// heartbeatResource implements llm.Resource for a fired heartbeat, exposing
94+
// the message as the name and the marshalled JSON as readable content.
95+
type heartbeatResource struct {
96+
uri string
97+
name string
98+
data []byte
99+
}
100+
101+
func (r *heartbeatResource) URI() string { return r.uri }
102+
func (r *heartbeatResource) Name() string { return r.name }
103+
func (r *heartbeatResource) Description() string { return "" }
104+
func (r *heartbeatResource) MIMEType() string { return "application/json" }
105+
func (r *heartbeatResource) Read(_ context.Context) ([]byte, error) {
106+
return r.data, nil
107+
}
108+
109+
///////////////////////////////////////////////////////////////////////////////
110+
// PRIVATE METHODS
111+
112+
// HeartbeatStore returns the heartbeat store, creating it lazily.
113+
// Heartbeats are stored in the user's cache directory.
114+
func (cmd *HeartbeatMCPCommand) HeartbeatStore(execName string) (*heartbeat.Store, error) {
115+
cache, err := os.UserCacheDir()
116+
if err != nil {
117+
return nil, fmt.Errorf("failed to determine cache directory: %w", err)
118+
}
119+
store, err := heartbeat.NewStore(filepath.Join(cache, execName, "heartbeats"))
120+
if err != nil {
121+
return nil, fmt.Errorf("failed to create heartbeat store: %w", err)
122+
}
123+
return store, nil
124+
}

‎cmd/llm/mcp_client.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
//go:build client
2+
3+
package main
4+
5+
type MCPCommands struct{}

‎cmd/llm/server.go‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,18 @@
33
package main
44

55
import (
6+
"context"
67
"fmt"
78
"os"
89
"path/filepath"
910

1011
// Packages
1112
goclient "github.com/mutablelogic/go-client"
13+
llm "github.com/mutablelogic/go-llm"
1214
homeassistant "github.com/mutablelogic/go-llm/pkg/homeassistant"
1315
httphandler "github.com/mutablelogic/go-llm/pkg/httphandler"
1416
manager "github.com/mutablelogic/go-llm/pkg/manager"
17+
mcpclient "github.com/mutablelogic/go-llm/pkg/mcp/client"
1518
newsapi "github.com/mutablelogic/go-llm/pkg/newsapi"
1619
anthropic "github.com/mutablelogic/go-llm/pkg/provider/anthropic"
1720
eliza "github.com/mutablelogic/go-llm/pkg/provider/eliza"
@@ -155,7 +158,13 @@ func (cmd *RunServer) WithManager(ctx server.Cmd, fn func(*manager.Manager, stri
155158

156159
// Add the MCP connector factory so CreateConnector probes servers on registration.
157160
// clientOpts already includes trace, tracer and timeout flags.
158-
opts = append(opts, manager.WithConnectorFactory(manager.MCPConnectorFactory(ctx.Name(), ctx.Version(), clientOpts...)))
161+
opts = append(opts, manager.WithConnectorFactory(manager.MCPConnectorFactory(ctx.Name(), ctx.Version(),
162+
mcpclient.WithClientOpt(clientOpts...),
163+
mcpclient.OptOnResourceUpdated(func(c context.Context, r llm.Resource) {
164+
ctx.Logger().Info("resource updated", "uri", r.URI(), "name", r.Name())
165+
}),
166+
)))
167+
opts = append(opts, manager.WithLogger(ctx.Logger()))
159168

160169
// Create the manager
161170
mgr, err := manager.NewManager(ctx.Name(), ctx.Version(), opts...)

‎cmd/llm/tool.go‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,16 @@
11
package main
22

33
import (
4+
"encoding/json"
45
"fmt"
56

67
// Packages
7-
server "github.com/mutablelogic/go-server"
88
otel "github.com/mutablelogic/go-client/pkg/otel"
99
httpclient "github.com/mutablelogic/go-llm/pkg/httpclient"
1010
opt "github.com/mutablelogic/go-llm/pkg/opt"
1111
schema "github.com/mutablelogic/go-llm/pkg/schema"
1212
uitable "github.com/mutablelogic/go-llm/pkg/ui/table"
13+
server "github.com/mutablelogic/go-server"
1314
types "github.com/mutablelogic/go-server/pkg/types"
1415
attribute "go.opentelemetry.io/otel/attribute"
1516
)
@@ -20,6 +21,7 @@ import (
2021
type ToolCommands struct {
2122
ListTools ListToolsCommand `cmd:"" name:"tools" help:"List tools." group:"TOOL"`
2223
GetTool GetToolCommand `cmd:"" name:"tool" help:"Get tool." group:"TOOL"`
24+
CallTool CallToolCommand `cmd:"" name:"call" help:"Call a tool with JSON input." group:"TOOL"`
2325
}
2426

2527
type ListToolsCommand struct {
@@ -31,6 +33,11 @@ type GetToolCommand struct {
3133
Name string `arg:"" name:"name" help:"Tool name"`
3234
}
3335

36+
type CallToolCommand struct {
37+
Name string `arg:"" name:"name" help:"Tool name"`
38+
Input string `arg:"" name:"input" help:"JSON input for the tool" optional:""`
39+
}
40+
3441
///////////////////////////////////////////////////////////////////////////////
3542
// COMMANDS
3643

@@ -95,3 +102,32 @@ func (cmd *GetToolCommand) Run(ctx server.Cmd) (err error) {
95102
fmt.Println(tool)
96103
return nil
97104
}
105+
106+
func (cmd *CallToolCommand) Run(ctx server.Cmd) (err error) {
107+
client, err := clientFor(ctx)
108+
if err != nil {
109+
return err
110+
}
111+
112+
// OTEL
113+
parent, endSpan := otel.StartSpan(ctx.Tracer(), ctx.Context(), "CallToolCommand",
114+
attribute.String("tool", cmd.Name),
115+
)
116+
defer func() { endSpan(err) }()
117+
118+
// Parse optional JSON input
119+
var input json.RawMessage
120+
if cmd.Input != "" {
121+
if err := json.Unmarshal([]byte(cmd.Input), &input); err != nil {
122+
return fmt.Errorf("invalid JSON input: %w", err)
123+
}
124+
}
125+
126+
response, err := client.CallTool(parent, cmd.Name, input)
127+
if err != nil {
128+
return err
129+
}
130+
131+
fmt.Println(response)
132+
return nil
133+
}

0 commit comments

Comments
 (0)