summaryrefslogtreecommitdiff
path: root/internal/executor/agentmcp.go
blob: 0b36c143845d7b737a57b2c7a84bc09e1c754c2b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
package executor

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"sync"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/thepeterstone/claudomator/internal/role"
)

// Registry maps per-task MCP bearer tokens to a built agent MCP server. A token
// is minted when a runner spawns an agent subprocess and revoked when it exits,
// so the server resolves task context from the token alone and never trusts an
// agent-supplied task ID.
type Registry struct {
	mu      sync.RWMutex
	servers map[string]*mcp.Server
}

func NewRegistry() *Registry {
	return &Registry{servers: make(map[string]*mcp.Server)}
}

// Mint creates a token bound to a freshly built MCP server for ch.
func (r *Registry) Mint(ch AgentChannel) (string, error) {
	buf := make([]byte, 32)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	token := hex.EncodeToString(buf)
	r.mu.Lock()
	r.servers[token] = newAgentServer(ch)
	r.mu.Unlock()
	return token, nil
}

func (r *Registry) server(token string) (*mcp.Server, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	s, ok := r.servers[token]
	return s, ok
}

func (r *Registry) Revoke(token string) {
	r.mu.Lock()
	delete(r.servers, token)
	r.mu.Unlock()
}

type askUserInput struct {
	Question string   `json:"question" jsonschema:"the question to ask the user; phrase it as a real question ending with a question mark"`
	Options  []string `json:"options,omitempty" jsonschema:"optional list of suggested answer choices"`
}

type reportSummaryInput struct {
	Summary string `json:"summary" jsonschema:"a 2-5 sentence summary of what you did and the outcome"`
}

type spawnSubtaskInput struct {
	Name         string   `json:"name" jsonschema:"short descriptive name for the subtask"`
	Instructions string   `json:"instructions" jsonschema:"complete instructions for the subtask agent"`
	Model        string   `json:"model,omitempty" jsonschema:"optional model override, e.g. sonnet or opus"`
	MaxBudgetUSD float64  `json:"max_budget_usd,omitempty" jsonschema:"optional budget cap in USD"`
	Role         string   `json:"role,omitempty" jsonschema:"optional role name to dispatch the subtask through instead of a fixed model (e.g. an evaluator role); when set, model is ignored and the role's escalation ladder picks the provider/model"`
	DependsOn    []string `json:"depends_on,omitempty" jsonschema:"optional list of sibling subtask IDs (returned by prior spawn_subtask calls in this same decomposition) this subtask must wait for before it can run"`
}

type recordProgressInput struct {
	Message string `json:"message" jsonschema:"a short progress note describing what you are doing"`
}

type proposeEpicInput struct {
	Name        string   `json:"name" jsonschema:"short descriptive name for the epic; matched by exact name to reuse an existing epic instead of creating a duplicate"`
	Description string   `json:"description,omitempty" jsonschema:"optional longer description of the initiative"`
	StoryIDs    []string `json:"story_ids" jsonschema:"the story IDs to group under this epic"`
}

// proposeRoleConfigInput mirrors internal/role.RoleConfig's fields directly
// (same json tags) rather than defining a parallel shape, so the tool's
// input decodes straight into a role.RoleConfig with no field-by-field
// translation. See propose_role_config below.
type proposeRoleConfigInput struct {
	Role             string      `json:"role" jsonschema:"the role name this config applies to (e.g. an existing role like builder, or a new one)"`
	SystemPrompt     string      `json:"system_prompt,omitempty" jsonschema:"system prompt appended for tasks dispatched through this role"`
	Tools            []string    `json:"tools,omitempty" jsonschema:"optional tool allowlist for this role"`
	SandboxKind      string      `json:"sandbox_kind,omitempty" jsonschema:"optional sandbox kind for this role"`
	DefaultBudgetUSD float64     `json:"default_budget_usd,omitempty" jsonschema:"optional estimated budget in USD, used when the scheduler considers escalating this role's tasks"`
	EscalationLadder []role.Tier `json:"escalation_ladder,omitempty" jsonschema:"ordered list of escalation tiers; each has candidates (provider/model pairs), an optional selection_mode (round_robin|single), and max_retries"`
}

func (in proposeRoleConfigInput) toRoleConfig() role.RoleConfig {
	return role.RoleConfig{
		Role:             in.Role,
		SystemPrompt:     in.SystemPrompt,
		Tools:            in.Tools,
		SandboxKind:      in.SandboxKind,
		DefaultBudgetUSD: in.DefaultBudgetUSD,
		EscalationLadder: in.EscalationLadder,
	}
}

func textResult(text string) *mcp.CallToolResult {
	return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: text}}}
}

// newAgentServer builds an MCP server exposing the five agent tools bound to ch.
func newAgentServer(ch AgentChannel) *mcp.Server {
	s := mcp.NewServer(&mcp.Implementation{Name: "claudomator", Version: "1"}, nil)

	mcp.AddTool(s, &mcp.Tool{
		Name:        "ask_user",
		Description: "Ask the user a question when you genuinely need a decision to proceed. Your turn ends after calling this; the task is resumed with the user's answer. Prefer making a reasonable decision and noting it in report_summary over asking.",
	}, func(ctx context.Context, _ *mcp.CallToolRequest, in askUserInput) (*mcp.CallToolResult, any, error) {
		q := map[string]any{"text": in.Question}
		if len(in.Options) > 0 {
			q["options"] = in.Options
		}
		payload, _ := json.Marshal(q)
		ans, err := ch.AskUser(ctx, string(payload))
		if errors.Is(err, ErrAgentBlocked) {
			return textResult("Question recorded. End your turn now without calling any more tools; the task will be resumed once the user answers."), nil, nil
		}
		if err != nil {
			return nil, nil, err
		}
		return textResult(ans), nil, nil
	})

	mcp.AddTool(s, &mcp.Tool{
		Name:        "report_summary",
		Description: "Record a concise summary of what you accomplished. Call this before finishing.",
	}, func(ctx context.Context, _ *mcp.CallToolRequest, in reportSummaryInput) (*mcp.CallToolResult, any, error) {
		if err := ch.ReportSummary(ctx, in.Summary); err != nil {
			return nil, nil, err
		}
		return textResult("Summary recorded."), nil, nil
	})

	mcp.AddTool(s, &mcp.Tool{
		Name:        "spawn_subtask",
		Description: "Create a child task to be executed separately. Use this to break large work into focused pieces, then finish your turn.",
	}, func(ctx context.Context, _ *mcp.CallToolRequest, in spawnSubtaskInput) (*mcp.CallToolResult, any, error) {
		id, err := ch.SpawnSubtask(ctx, SubtaskSpec{
			Name:         in.Name,
			Instructions: in.Instructions,
			Model:        in.Model,
			MaxBudgetUSD: in.MaxBudgetUSD,
			Role:         in.Role,
			DependsOn:    in.DependsOn,
		})
		if err != nil {
			return nil, nil, err
		}
		return textResult("Created subtask " + id), nil, nil
	})

	mcp.AddTool(s, &mcp.Tool{
		Name:        "record_progress",
		Description: "Record a short progress note that appears in the task timeline.",
	}, func(ctx context.Context, _ *mcp.CallToolRequest, in recordProgressInput) (*mcp.CallToolResult, any, error) {
		if err := ch.RecordProgress(ctx, in.Message); err != nil {
			return nil, nil, err
		}
		return textResult("Noted."), nil, nil
	})

	mcp.AddTool(s, &mcp.Tool{
		Name:        "propose_epic",
		Description: "Group one or more stories under a new or existing epic (matched by exact name) when they form a cohesive initiative. Only call this when you've been given several story IDs and independently judge that they belong together.",
	}, func(ctx context.Context, _ *mcp.CallToolRequest, in proposeEpicInput) (*mcp.CallToolResult, any, error) {
		id, err := ch.ProposeEpic(ctx, EpicProposal{
			Name:        in.Name,
			Description: in.Description,
			StoryIDs:    in.StoryIDs,
		})
		if err != nil {
			return nil, nil, err
		}
		return textResult("Proposed epic " + id), nil, nil
	})

	mcp.AddTool(s, &mcp.Tool{
		Name:        "propose_role_config",
		Description: "Propose a new draft configuration version for a role, after reflecting on what happened (e.g. during a story retro). Creates a new draft role_configs row for a human to review and activate via POST /api/roles/{role}/activate -- it never changes what is currently active.",
	}, func(ctx context.Context, _ *mcp.CallToolRequest, in proposeRoleConfigInput) (*mcp.CallToolResult, any, error) {
		version, err := ch.ProposeRoleConfig(ctx, in.toRoleConfig())
		if err != nil {
			return nil, nil, err
		}
		return textResult(fmt.Sprintf("Proposed role config %s v%d (draft)", in.Role, version)), nil, nil
	})

	return s
}

func bearerToken(r *http.Request) string {
	h := r.Header.Get("Authorization")
	if h == "" {
		return ""
	}
	return strings.TrimSpace(strings.TrimPrefix(h, "Bearer "))
}

// NewAgentMCPHandler returns the HTTP handler for the per-task agent MCP server.
// It resolves the request's bearer token to the server for that task's run;
// unknown tokens yield a 400 from the underlying handler.
func NewAgentMCPHandler(reg *Registry) http.Handler {
	return mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server {
		s, ok := reg.server(bearerToken(r))
		if !ok {
			return nil
		}
		return s
	}, nil)
}