Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 10 additions & 8 deletions cmd/marvel/ctxforward.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,13 @@ type statuslinePayload struct {
}

// renderForward parses one statusline payload and returns the status text
// to print plus the context percentage to forward (send=false when the
// payload carries no forwardable figure). Pure so it is table-testable.
func renderForward(raw []byte) (line string, pct float64, send bool) {
// to print, the context percentage and model name to forward (send=false
// when the payload carries no forwardable figure). Pure so it is
// table-testable.
func renderForward(raw []byte) (line string, pct float64, model string, send bool) {
var p statuslinePayload
if err := json.Unmarshal(raw, &p); err != nil {
return "marvel ctx-forward: unreadable payload", 0, false
return "marvel ctx-forward: unreadable payload", 0, "", false
}

// Subagent shape: summarize the task rows. No RPC — the daemon has
Expand All @@ -70,17 +71,17 @@ func renderForward(raw []byte) (line string, pct float64, send bool) {
}
}
}
return fmt.Sprintf("agents %d/%d running · max CTX %.0f%%", running, len(p.Tasks), maxPct), 0, false
return fmt.Sprintf("agents %d/%d running · max CTX %.0f%%", running, len(p.Tasks), maxPct), 0, "", false
}

if p.ContextWindow == nil || p.ContextWindow.UsedPercentage == nil {
// Session too young to have a measurement. Show something
// stable rather than flickering an error.
return fmt.Sprintf("%s · CTX –", orUnknown(p.Model.DisplayName)), 0, false
return fmt.Sprintf("%s · CTX –", orUnknown(p.Model.DisplayName)), 0, "", false
}
pct = *p.ContextWindow.UsedPercentage
line = fmt.Sprintf("%s · CTX %.0f%% · $%.2f", orUnknown(p.Model.DisplayName), pct, p.Cost.TotalCostUSD)
return line, pct, true
return line, pct, p.Model.DisplayName, true
}

func orUnknown(s string) string {
Expand All @@ -101,7 +102,7 @@ func newCtxForwardCmd() *cobra.Command {
fmt.Println("marvel ctx-forward")
return nil
}
line, pct, send := renderForward(raw)
line, pct, model, send := renderForward(raw)
fmt.Println(line)

socket := os.Getenv("MARVEL_SOCKET")
Expand All @@ -113,6 +114,7 @@ func newCtxForwardCmd() *cobra.Command {
params, _ := json.Marshal(map[string]any{
"session_key": workspace + "/" + session,
"context_percent": pct,
"model": model,
})
// Best-effort by design; see the failure posture above.
_, _ = daemon.SendRequest(socket, daemon.Request{
Expand Down
23 changes: 14 additions & 9 deletions cmd/marvel/ctxforward_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,22 @@ import (
func TestRenderForward(t *testing.T) {
t.Parallel()
tests := []struct {
name string
payload string
wantSend bool
wantPct float64
wantIn string
name string
payload string
wantSend bool
wantPct float64
wantModel string
wantIn string
}{
{
name: "main payload with measurement",
payload: `{"model":{"display_name":"Haiku 4.5"},
"cost":{"total_cost_usd":0.0898},
"context_window":{"used_percentage":17}}`,
wantSend: true,
wantPct: 17,
wantIn: "CTX 17%",
wantSend: true,
wantPct: 17,
wantModel: "Haiku 4.5",
wantIn: "CTX 17%",
},
{
name: "young session, null percentage",
Expand Down Expand Up @@ -55,13 +57,16 @@ func TestRenderForward(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
line, pct, send := renderForward([]byte(tt.payload))
line, pct, model, send := renderForward([]byte(tt.payload))
if send != tt.wantSend {
t.Fatalf("send = %v, want %v", send, tt.wantSend)
}
if send && pct != tt.wantPct {
t.Errorf("pct = %v, want %v", pct, tt.wantPct)
}
if send && model != tt.wantModel {
t.Errorf("model = %q, want %q", model, tt.wantModel)
}
if !strings.Contains(line, tt.wantIn) {
t.Errorf("line = %q, want it to contain %q", line, tt.wantIn)
}
Expand Down
53 changes: 38 additions & 15 deletions cmd/marvel/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -1466,7 +1466,7 @@ func sortSessions(sessions []api.Session, ws *watchSort) {
less = sessions[i].Workspace < sessions[j].Workspace
case "state":
less = string(sessions[i].State) < string(sessions[j].State)
case "agent":
case "runtime":
ai, aj := sessions[i].Runtime.Name, sessions[j].Runtime.Name
if ai == "" {
ai = sessions[i].Runtime.Command
Expand All @@ -1475,6 +1475,17 @@ func sortSessions(sessions []api.Session, ws *watchSort) {
aj = sessions[j].Runtime.Command
}
less = ai < aj
case "llm":
less = sessions[i].ContextModel < sessions[j].ContextModel
case "health":
hi, hj := string(sessions[i].HealthState), string(sessions[j].HealthState)
if hi == "" {
hi = "unknown"
}
if hj == "" {
hj = "unknown"
}
less = hi < hj
case "desk":
less = sessions[i].PaneID < sessions[j].PaneID
default:
Expand Down Expand Up @@ -1529,11 +1540,18 @@ func formatBytes(n int64) string {
func renderSessionTable(sessions []api.Session) string {
var buf bytes.Buffer
w := tabwriter.NewWriter(&buf, 0, 4, 2, ' ', 0)
_, _ = fmt.Fprintf(w, "WORKSPACE\tTEAM\tROLE\tGEN\tNAME\tSTATE\tHEALTH\tCTX%%\tCPU%%\tRSS\tDESK\tAGENT\n")
_, _ = fmt.Fprintf(w, "WORKSPACE\tTEAM\tROLE\tGEN\tAGENT NAME\tSTATE\tHEALTH\tCTX%%\tCPU%%\tRSS\tDESK\tRUNTIME\tLLM\n")
for _, s := range sessions {
agent := s.Runtime.Name
if agent == "" {
agent = s.Runtime.Command
runtimeName := s.Runtime.Name
if runtimeName == "" {
runtimeName = s.Runtime.Command
}
// LLM is the model as the metering producer named it: the
// stream accountant's raw model for headless sessions, the
// statusline feed's display name for interactive ones.
llm := s.ContextModel
if llm == "" {
llm = "-"
}
// CTX% has two producers: the cooperative heartbeat RPC (the
// simulator) and the usage accountant fed by adapter streams.
Expand Down Expand Up @@ -1577,8 +1595,8 @@ func renderSessionTable(sessions []api.Session) string {
if health == "" {
health = "unknown"
}
_, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
s.Workspace, s.Team, s.Role, gen, s.Name, s.State, health, ctx, cpu, rss, desk, agent)
_, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
s.Workspace, s.Team, s.Role, gen, s.Name, s.State, health, ctx, cpu, rss, desk, runtimeName, llm)
}
_ = w.Flush()
return buf.String()
Expand All @@ -1593,12 +1611,13 @@ func renderWatch(ws *watchSort, interval time.Duration) string {
if ws.showHelp {
fmt.Fprintf(&buf, "\n")
fmt.Fprintf(&buf, " Sort keys (toggle asc/desc):\n")
fmt.Fprintf(&buf, " w workspace t team r role\n")
fmt.Fprintf(&buf, " g generation n name s state\n")
fmt.Fprintf(&buf, " c context d desk a agent\n")
fmt.Fprintf(&buf, " p cpu m memory\n")
fmt.Fprintf(&buf, " w workspace t team R role\n")
fmt.Fprintf(&buf, " g generation n agent name s state\n")
fmt.Fprintf(&buf, " c context d desk r runtime\n")
fmt.Fprintf(&buf, " l llm h health\n")
fmt.Fprintf(&buf, " p cpu m memory (rss)\n")
fmt.Fprintf(&buf, "\n")
fmt.Fprintf(&buf, " h toggle help q quit\n")
fmt.Fprintf(&buf, " ? toggle help q quit\n")
fmt.Fprintf(&buf, "\n")
return buf.String()
}
Expand All @@ -1609,7 +1628,7 @@ func renderWatch(ws *watchSort, interval time.Duration) string {
} else {
sortLabel += " asc"
}
fmt.Fprintf(&buf, "sort: %s h:help q:quit\n\n", sortLabel)
fmt.Fprintf(&buf, "sort: %s ?:help q:quit\n\n", sortLabel)

sessions, err := fetchSessions()
if err != nil {
Expand Down Expand Up @@ -1687,7 +1706,11 @@ func watchSessionsLoop(interval time.Duration) error {
case 'n':
toggleSort(ws, "name", false)
case 'r':
toggleSort(ws, "runtime", false)
case 'R':
toggleSort(ws, "role", false)
case 'l':
toggleSort(ws, "llm", false)
case 'g':
toggleSort(ws, "generation", false)
case 't':
Expand All @@ -1696,11 +1719,11 @@ func watchSessionsLoop(interval time.Duration) error {
toggleSort(ws, "workspace", false)
case 's':
toggleSort(ws, "state", false)
case 'a':
toggleSort(ws, "agent", false)
case 'd':
toggleSort(ws, "desk", false)
case 'h':
toggleSort(ws, "health", false)
case '?':
ws.showHelp = !ws.showHelp
default:
continue
Expand Down
19 changes: 16 additions & 3 deletions cmd/marvel/render_test.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package main

import (
"regexp"
"strings"
"testing"
"time"
Expand Down Expand Up @@ -30,14 +31,26 @@ func TestFormatBytes(t *testing.T) {
}
}

// splitColumns splits a tabwriter-rendered line on runs of two or more
// spaces, so multi-word headers ("AGENT NAME") stay one column.
func splitColumns(line string) []string {
var out []string
for _, f := range regexp.MustCompile(`\s{2,}`).Split(strings.TrimSpace(line), -1) {
if f != "" {
out = append(out, f)
}
}
return out
}

func column(t *testing.T, table, name string) string {
t.Helper()
lines := strings.Split(strings.TrimRight(table, "\n"), "\n")
if len(lines) < 2 {
t.Fatalf("table has no rows:\n%s", table)
}
header := strings.Fields(lines[0])
row := strings.Fields(lines[1])
header := splitColumns(lines[0])
row := splitColumns(lines[1])
if len(header) != len(row) {
t.Fatalf("header has %d columns, row has %d:\n%s", len(header), len(row), table)
}
Expand Down Expand Up @@ -131,7 +144,7 @@ func TestRenderSessionTableWithHeartbeat(t *testing.T) {
}); err != nil {
t.Fatalf("create session: %v", err)
}
if err := store.UpdateSessionHeartbeat("ws/agent-0", 55.4); err != nil {
if err := store.UpdateSessionHeartbeat("ws/agent-0", 55.4, ""); err != nil {
t.Fatalf("heartbeat: %v", err)
}

Expand Down
2 changes: 1 addition & 1 deletion internal/api/bolt_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -407,7 +407,7 @@ func TestBoltStore_HeartbeatContextReadingSurvivesRehydrate(t *testing.T) {
}); err != nil {
t.Fatalf("create session: %v", err)
}
if err := s1.UpdateSessionHeartbeat("ws/agent-0", 64.0); err != nil {
if err := s1.UpdateSessionHeartbeat("ws/agent-0", 64.0, ""); err != nil {
t.Fatalf("heartbeat: %v", err)
}
if err := s1.CloseBolt(); err != nil {
Expand Down
8 changes: 7 additions & 1 deletion internal/api/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -502,7 +502,7 @@ func (s *Store) UpdatePolicy(key string, fn func(*Policy) error) error {
// dominant write rate for marvel's bbolt usage; if it surfaces as a
// performance issue, batch by waiting N heartbeats before persisting
// (or move heartbeat state into a separate in-memory-only path).
func (s *Store) UpdateSessionHeartbeat(key string, contextPercent float64) error {
func (s *Store) UpdateSessionHeartbeat(key string, contextPercent float64, model string) error {
s.mu.Lock()
defer s.mu.Unlock()
sess, ok := s.sessions[key]
Expand All @@ -511,6 +511,12 @@ func (s *Store) UpdateSessionHeartbeat(key string, contextPercent float64) error
}
sess.ContextPercent = contextPercent
sess.LastHeartbeat = time.Now().UTC()
// A cooperative reporter that knows its model names it (the
// statusline feed does); one that doesn't sends "" and any
// prior reading stands.
if model != "" {
sess.ContextModel = model
}
// ContextAt is the single "measured" sentinel for the context column,
// shared with the usage accountant's path below. A cooperative
// heartbeat is a measurement too, so it stamps it.
Expand Down
6 changes: 3 additions & 3 deletions internal/api/store_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ func TestUpdateSessionHeartbeat(t *testing.T) {
t.Fatalf("create session: %v", err)
}

if err := s.UpdateSessionHeartbeat("test-ws/agent-0", 42.5); err != nil {
if err := s.UpdateSessionHeartbeat("test-ws/agent-0", 42.5, ""); err != nil {
t.Fatalf("update heartbeat: %v", err)
}

Expand All @@ -180,7 +180,7 @@ func TestUpdateSessionHeartbeat(t *testing.T) {
}

// Not found case
if err := s.UpdateSessionHeartbeat("test-ws/nonexistent", 10); !errors.Is(err, ErrNotFound) {
if err := s.UpdateSessionHeartbeat("test-ws/nonexistent", 10, ""); !errors.Is(err, ErrNotFound) {
t.Fatalf("expected ErrNotFound, got %v", err)
}
}
Expand Down Expand Up @@ -361,7 +361,7 @@ func TestUpdateSessionHeartbeatStampsContextAt(t *testing.T) {
if err := s.CreateSession(sess); err != nil {
t.Fatalf("create session: %v", err)
}
if err := s.UpdateSessionHeartbeat("test-ws/agent-0", 42.5); err != nil {
if err := s.UpdateSessionHeartbeat("test-ws/agent-0", 42.5, ""); err != nil {
t.Fatalf("update heartbeat: %v", err)
}
got, _ := s.GetSession("test-ws/agent-0")
Expand Down
6 changes: 5 additions & 1 deletion internal/daemon/daemon.go
Original file line number Diff line number Diff line change
Expand Up @@ -952,6 +952,10 @@ func (d *Daemon) handleScale(params json.RawMessage) Response {
type heartbeatParams struct {
SessionKey string `json:"session_key"`
ContextPercent float64 `json:"context_percent"`
// Model is the model as the reporter names it, "" when the
// reporter does not know (the simulator). The statusline feed
// sends the harness's display name.
Model string `json:"model,omitempty"`
}

func (d *Daemon) handleHeartbeat(params json.RawMessage) Response {
Expand All @@ -960,7 +964,7 @@ func (d *Daemon) handleHeartbeat(params json.RawMessage) Response {
return Response{Error: fmt.Sprintf("bad params: %v", err)}
}

if err := d.store.UpdateSessionHeartbeat(p.SessionKey, p.ContextPercent); err != nil {
if err := d.store.UpdateSessionHeartbeat(p.SessionKey, p.ContextPercent, p.Model); err != nil {
return Response{Error: err.Error()}
}

Expand Down