Skip to content

Commit fd1b20f

Browse files
code-ashercode-asher
committed
Add status watcher to MCP server
Since we can now get status updates from two places, they are placed in a queue so we can handle them one at a time.
1 parent d46ea56 commit fd1b20f

File tree

6 files changed

+456
-35
lines changed

6 files changed

+456
-35
lines changed

‎cli/exp_mcp.go

Lines changed: 187 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"context"
66
"encoding/json"
77
"errors"
8+
"net/url"
89
"os"
910
"path/filepath"
1011
"slices"
@@ -15,8 +16,10 @@ import (
1516
"github.com/spf13/afero"
1617
"golang.org/x/xerrors"
1718

19+
agentapi "github.com/coder/agentapi-sdk-go"
1820
"github.com/coder/coder/v2/buildinfo"
1921
"github.com/coder/coder/v2/cli/cliui"
22+
"github.com/coder/coder/v2/cli/cliutil"
2023
"github.com/coder/coder/v2/codersdk"
2124
"github.com/coder/coder/v2/codersdk/agentsdk"
2225
"github.com/coder/coder/v2/codersdk/toolsdk"
@@ -25,6 +28,7 @@ import (
2528

2629
const (
2730
envAppStatusSlug = "CODER_MCP_APP_STATUS_SLUG"
31+
envLLMAgentURL = "CODER_MCP_LLM_AGENT_URL"
2832
)
2933

3034
func (r *RootCmd) mcpCommand() *serpent.Command {
@@ -347,10 +351,20 @@ func (*RootCmd) mcpConfigureCursor() *serpent.Command {
347351
return cmd
348352
}
349353

354+
type reportTask struct {
355+
link string
356+
messageID int64
357+
selfReported bool
358+
state codersdk.WorkspaceAppStatusState
359+
summary string
360+
}
361+
350362
type mcpServer struct {
351363
agentClient *agentsdk.Client
352364
appStatusSlug string
353365
client *codersdk.Client
366+
llmClient *agentapi.Client
367+
queue *cliutil.Queue[reportTask]
354368
}
355369

356370
func (r *RootCmd) mcpServer() *serpent.Command {
@@ -359,12 +373,14 @@ func (r *RootCmd) mcpServer() *serpent.Command {
359373
instructions string
360374
allowedTools []string
361375
appStatusSlug string
376+
llmAgentURL url.URL
362377
)
363378
return &serpent.Command{
364379
Use: "server",
365380
Handler: func(inv *serpent.Invocation) error {
366381
srv := &mcpServer{
367382
appStatusSlug: appStatusSlug,
383+
queue: cliutil.NewQueue[reportTask](10),
368384
}
369385

370386
// Display client URL separately from authentication status.
@@ -408,8 +424,36 @@ func (r *RootCmd) mcpServer() *serpent.Command {
408424
cliui.Infof(inv.Stderr, "Task reporter : Enabled")
409425
}
410426

411-
// Start the server.
412-
return srv.start(inv, instructions, allowedTools)
427+
// Try to create a client for the LLM agent API, which is used to get the
428+
// screen status to make the status reporting more robust. No auth
429+
// needed, so no validation.
430+
if llmAgentURL.String() == "" {
431+
cliui.Infof(inv.Stderr, "LLM agent URL : Not configured")
432+
} else {
433+
cliui.Infof(inv.Stderr, "LLM agent URL : %s", llmAgentURL.String())
434+
llmClient, err := agentapi.NewClient(llmAgentURL.String())
435+
if err != nil {
436+
cliui.Infof(inv.Stderr, "Screen events : Disabled")
437+
cliui.Warnf(inv.Stderr, "%s must be set", envLLMAgentURL)
438+
} else {
439+
cliui.Infof(inv.Stderr, "Screen events : Enabled")
440+
srv.llmClient = llmClient
441+
}
442+
}
443+
444+
ctx, cancel := context.WithCancel(inv.Context())
445+
defer cancel()
446+
defer srv.queue.Close()
447+
448+
cliui.Infof(inv.Stderr, "Failed to watch screen events")
449+
// Start the reporter, watcher, and server.
450+
if srv.agentClient != nil && appStatusSlug != "" {
451+
srv.startReporter(ctx, inv)
452+
if srv.llmClient != nil {
453+
srv.startWatcher(ctx, inv)
454+
}
455+
}
456+
return srv.startServer(ctx, inv, instructions, allowedTools)
413457
},
414458
Short: "Start the Coder MCP server.",
415459
Middleware: serpent.Chain(
@@ -438,14 +482,142 @@ func (r *RootCmd) mcpServer() *serpent.Command {
438482
Value: serpent.StringOf(&appStatusSlug),
439483
Default: "",
440484
},
485+
{
486+
Flag: "llm-agent-url",
487+
Description: "The URL of the LLM agent API, used to listen for status updates.",
488+
Env: envLLMAgentURL,
489+
Value: serpent.URLOf(&llmAgentURL),
490+
},
441491
},
442492
}
443493
}
444494

445-
func (s *mcpServer) start(inv *serpent.Invocation, instructions string, allowedTools []string) error {
446-
ctx, cancel := context.WithCancel(inv.Context())
447-
defer cancel()
495+
func (s *mcpServer) startReporter(ctx context.Context, inv *serpent.Invocation) {
496+
var lastMessageID int64
497+
shouldUpdate := func(item reportTask) codersdk.WorkspaceAppStatusState {
498+
// Always send self-reported updates.
499+
if item.selfReported {
500+
return item.state
501+
}
502+
// Always send completed states.
503+
switch item.state {
504+
case codersdk.WorkspaceAppStatusStateComplete,
505+
codersdk.WorkspaceAppStatusStateFailure:
506+
return item.state
507+
}
508+
// Always send "working" when there is a new message, since this means the
509+
// user submitted a message through the API and we know the LLM will begin
510+
// work soon if it has not already.
511+
if item.messageID > lastMessageID {
512+
return codersdk.WorkspaceAppStatusStateWorking
513+
}
514+
// Otherwise, if the state is "working" and there have been no new messages,
515+
// it means either that the LLM is still working or it means the user has
516+
// interacted with the terminal directly. For now, we are ignoring these
517+
// updates. This risks missing cases where the user manually submits a new
518+
// prompt and the LLM becomes active and does not update itself, but it
519+
// avoids spamming useless status updates.
520+
return ""
521+
}
522+
var lastPayload agentsdk.PatchAppStatus
523+
go func() {
524+
for {
525+
// TODO: Even with the queue, there is still the potential that a message
526+
// from the screen watcher and a message from the LLM could arrive out of
527+
// order if the timing is just right. We might want to wait a bit, then
528+
// check if the status has changed before committing.
529+
item, ok := s.queue.Pop()
530+
if !ok {
531+
return
532+
}
533+
534+
state := shouldUpdate(item)
535+
if state == "" {
536+
continue
537+
}
538+
539+
if item.messageID != 0 {
540+
lastMessageID = item.messageID
541+
}
542+
543+
payload := agentsdk.PatchAppStatus{
544+
AppSlug: s.appStatusSlug,
545+
Message: item.summary,
546+
URI: item.link,
547+
State: state,
548+
}
549+
550+
// Preserve previous message and URI.
551+
if payload.Message == "" {
552+
payload.Message = lastPayload.Message
553+
}
554+
if payload.URI == "" {
555+
payload.URI = lastPayload.URI
556+
}
557+
558+
// Avoid sending duplicate updates.
559+
if lastPayload.State == payload.State &&
560+
lastPayload.URI == payload.URI &&
561+
lastPayload.Message == payload.Message {
562+
continue
563+
}
564+
565+
err := s.agentClient.PatchAppStatus(ctx, payload)
566+
if err != nil && !errors.Is(err, context.Canceled) {
567+
cliui.Warnf(inv.Stderr, "Failed to report task status: %s", err)
568+
}
569+
570+
lastPayload = payload
571+
}
572+
}()
573+
}
574+
575+
func (s *mcpServer) startWatcher(ctx context.Context, inv *serpent.Invocation) {
576+
eventsCh, errCh, err := s.llmClient.SubscribeEvents(ctx)
577+
if err != nil {
578+
cliui.Warnf(inv.Stderr, "Failed to watch screen events: %s", err)
579+
return
580+
}
581+
go func() {
582+
for {
583+
select {
584+
case <-ctx.Done():
585+
return
586+
case event := <-eventsCh:
587+
switch ev := event.(type) {
588+
case agentapi.EventStatusChange:
589+
// If the screen is stable, assume complete.
590+
state := codersdk.WorkspaceAppStatusStateWorking
591+
if ev.Status == agentapi.StatusStable {
592+
state = codersdk.WorkspaceAppStatusStateComplete
593+
}
594+
err := s.queue.Push(reportTask{
595+
state: state,
596+
})
597+
if err != nil {
598+
cliui.Warnf(inv.Stderr, "Failed to queue update: %s", err)
599+
return
600+
}
601+
case agentapi.EventMessageUpdate:
602+
err := s.queue.Push(reportTask{
603+
messageID: ev.Id,
604+
})
605+
if err != nil {
606+
cliui.Warnf(inv.Stderr, "Failed to queue update: %s", err)
607+
return
608+
}
609+
}
610+
case err := <-errCh:
611+
if !errors.Is(err, context.Canceled) {
612+
cliui.Warnf(inv.Stderr, "Received error from screen event watcher: %s", err)
613+
}
614+
return
615+
}
616+
}
617+
}()
618+
}
448619

620+
func (s *mcpServer) startServer(ctx context.Context, inv *serpent.Invocation, instructions string, allowedTools []string) error {
449621
cliui.Infof(inv.Stderr, "Starting MCP server")
450622

451623
cliui.Infof(inv.Stderr, "Instructions : %q", instructions)
@@ -476,8 +648,16 @@ func (s *mcpServer) start(inv *serpent.Invocation, instructions string, allowedT
476648

477649
// Add tool dependencies.
478650
toolOpts := []func(*toolsdk.Deps){
479-
toolsdk.WithAgentClient(s.agentClient),
480-
toolsdk.WithAppStatusSlug(s.appStatusSlug),
651+
toolsdk.WithTaskReporter(func(args toolsdk.ReportTaskArgs) error {
652+
// TODO: Is it OK to just push and return or should we wait for it to
653+
// actually get disatched to return any request errors?
654+
return s.queue.Push(reportTask{
655+
link: args.Link,
656+
selfReported: true,
657+
state: codersdk.WorkspaceAppStatusState(args.State),
658+
summary: args.Summary,
659+
})
660+
}),
481661
}
482662

483663
toolDeps, err := toolsdk.NewDeps(s.client, toolOpts...)

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy