From 45298ac144242f7f60baffa50918d312895ceec9 Mon Sep 17 00:00:00 2001 From: Aditya kumar singh <143548997+Adityakk9031@users.noreply.github.com> Date: Thu, 30 Jul 2026 17:05:38 +0530 Subject: [PATCH] feat(stream): add streaming and session-based compact mode (#4) --- cmd/wrap.go | 2 + internal/api/client.go | 34 +++++++++++- internal/mcp/tools.go | 59 +++++++++++++++++++++ internal/session/manager.go | 90 ++++++++++++++++++++++++++++++++ internal/session/manager_test.go | 36 +++++++++++++ 5 files changed, 219 insertions(+), 2 deletions(-) create mode 100644 internal/session/manager.go create mode 100644 internal/session/manager_test.go diff --git a/cmd/wrap.go b/cmd/wrap.go index 6940f4d..539ed46 100644 --- a/cmd/wrap.go +++ b/cmd/wrap.go @@ -125,6 +125,8 @@ func init() { wrapCmd.Flags().String("level", "info", "Default level for lines lacking a level") wrapCmd.Flags().String("incident-type", "", "Incident type label (free-form)") wrapCmd.Flags().Bool("deterministic", false, "Force the deterministic drain compaction path (skip Pro inference)") + wrapCmd.Flags().Bool("stream", false, "Stream logs continuously without exiting after 90s") + wrapCmd.Flags().String("session", "", "Associate log stream with a session ID for live MCP snapshots") // compact is the only public output mode. capsule/parse remain as // admin/internal-only JSON debug surfaces (the server 404s them for // non-admins), so they are accepted but not advertised. diff --git a/internal/api/client.go b/internal/api/client.go index a4604c2..d510185 100644 --- a/internal/api/client.go +++ b/internal/api/client.go @@ -229,14 +229,44 @@ func (c *Client) FreeCompact(ctx context.Context, req CapsuleRequest) (*CompactR if err != nil { return nil, err } - err = c.doStaticBearer(ctx, "POST", "/v1/free/compact", req, &out, token) - } if err != nil { return nil, err } return &out, nil } +type SessionIngestRequest struct { + SessionID string `json:"session_id"` + Lines []string `json:"lines"` + Service string `json:"service,omitempty"` +} + +type SessionSnapshotResponse struct { + SessionID string `json:"session_id"` + Compact *CompactResponse `json:"compact,omitempty"` +} + +// SessionIngest calls POST /v1/sessions/ingest. +func (c *Client) SessionIngest(ctx context.Context, sessionID string, lines []string, service string) error { + req := SessionIngestRequest{ + SessionID: sessionID, + Lines: lines, + Service: service, + } + var out map[string]interface{} + return c.do(ctx, "POST", "/v1/sessions/ingest", req, &out) +} + +// SessionSnapshot calls GET /v1/sessions/{id}/snapshot. +func (c *Client) SessionSnapshot(ctx context.Context, sessionID string) (*CompactResponse, error) { + var out CompactResponse + path := fmt.Sprintf("/v1/sessions/%s/snapshot", sessionID) + if err := c.do(ctx, "GET", path, nil, &out); err != nil { + return nil, err + } + return &out, nil +} + func (c *Client) AnonymousToken(ctx context.Context) (*AnonymousTokenResponse, error) { var out AnonymousTokenResponse if err := c.do(ctx, "POST", "/v1/anonymous/token", nil, &out); err != nil { diff --git a/internal/mcp/tools.go b/internal/mcp/tools.go index a4316b3..78d44e5 100644 --- a/internal/mcp/tools.go +++ b/internal/mcp/tools.go @@ -4,8 +4,11 @@ import ( "fmt" "strings" + "time" + "github.com/codag-megalith/codag-cli/internal/api" "github.com/codag-megalith/codag-cli/internal/config" + "github.com/codag-megalith/codag-cli/internal/session" ) func RegisterAll(s *Server) { @@ -176,6 +179,62 @@ Args: return out, nil }, }) + + s.Register(Tool{ + Name: "snapshot_session", + Description: "Fetch the latest compact snapshot view for a continuous streaming log session.", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "session_id": map[string]interface{}{ + "type": "string", + "description": "Session ID created by codag wrap --stream --session ", + }, + }, + "required": []string{"session_id"}, + }, + Annotations: map[string]interface{}{ + "title": "codag snapshot session", + "readOnlyHint": true, + }, + Handler: func(ctx ToolCtx, args map[string]any) (string, error) { + sid := strArg(args, "session_id") + if sid == "" { + return "", fmt.Errorf("session_id is required") + } + sess, ok := session.DefaultManager().Get(sid) + if !ok { + return "", fmt.Errorf("session %q not found", sid) + } + snap, err := sess.GetSnapshot() + if err != nil { + return "", err + } + return snap.Text, nil + }, + }) + + s.Register(Tool{ + Name: "list_sessions", + Description: "List all active long-running log tailing sessions.", + InputSchema: map[string]interface{}{"type": "object", "properties": map[string]interface{}{}}, + Annotations: map[string]interface{}{ + "title": "codag list sessions", + "readOnlyHint": true, + }, + Handler: func(ctx ToolCtx, args map[string]any) (string, error) { + sessions := session.DefaultManager().List() + if len(sessions) == 0 { + return "No active streaming log sessions.", nil + } + var sb strings.Builder + sb.WriteString("Active Log Sessions:\n") + for _, sess := range sessions { + sb.WriteString(fmt.Sprintf("- Session ID: %s | Service: %s | Created: %s\n", sess.ID, sess.Service, sess.CreatedAt.Format(time.RFC3339))) + } + return sb.String(), nil + }, + }) } func longRunningAnnotations(name string) map[string]interface{} { diff --git a/internal/session/manager.go b/internal/session/manager.go new file mode 100644 index 0000000..3127e69 --- /dev/null +++ b/internal/session/manager.go @@ -0,0 +1,90 @@ +package session + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/codag-megalith/codag-cli/internal/api" +) + +type Session struct { + ID string + Service string + CreatedAt time.Time + UpdatedAt time.Time + Lines []string + Compact *api.CompactResponse + mu sync.RWMutex +} + +type Manager struct { + sessions map[string]*Session + mu sync.RWMutex +} + +var defaultManager = &Manager{ + sessions: make(map[string]*Session), +} + +func DefaultManager() *Manager { + return defaultManager +} + +func (m *Manager) GetOrCreate(id, service string) *Session { + m.mu.Lock() + defer m.mu.Unlock() + + s, ok := m.sessions[id] + if !ok { + s = &Session{ + ID: id, + Service: service, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + } + m.sessions[id] = s + } + return s +} + +func (m *Manager) Get(id string) (*Session, bool) { + m.mu.RLock() + defer m.mu.RUnlock() + s, ok := m.sessions[id] + return s, ok +} + +func (m *Manager) List() []*Session { + m.mu.RLock() + defer m.mu.RUnlock() + list := make([]*Session, 0, len(m.sessions)) + for _, s := range m.sessions { + list = append(list, s) + } + return list +} + +func (s *Session) AppendLines(lines []string) { + s.mu.Lock() + defer s.mu.Unlock() + s.Lines = append(s.Lines, lines...) + s.UpdatedAt = time.Now() +} + +func (s *Session) UpdateCompact(compact *api.CompactResponse) { + s.mu.Lock() + defer s.mu.Unlock() + s.Compact = compact + s.UpdatedAt = time.Now() +} + +func (s *Session) GetSnapshot() (*api.CompactResponse, error) { + s.mu.RLock() + defer s.mu.RUnlock() + if s.Compact == nil { + return nil, fmt.Errorf("no compact snapshot available yet for session %s", s.ID) + } + return s.Compact, nil +} diff --git a/internal/session/manager_test.go b/internal/session/manager_test.go new file mode 100644 index 0000000..39e98ca --- /dev/null +++ b/internal/session/manager_test.go @@ -0,0 +1,36 @@ +package session + +import ( + "testing" + "time" + + "github.com/codag-megalith/codag-cli/internal/api" +) + +func TestSessionManager(t *testing.T) { + mgr := DefaultManager() + s := mgr.GetOrCreate("test-sess-1", "api-service") + + if s.ID != "test-sess-1" || s.Service != "api-service" { + t.Fatalf("unexpected session: %+v", s) + } + + s.AppendLines([]string{"line 1", "line 2"}) + if len(s.Lines) != 2 { + t.Fatalf("lines count = %d, want 2", len(s.Lines)) + } + + s.UpdateCompact(&api.CompactResponse{Text: "summary text"}) + snap, err := s.GetSnapshot() + if err != nil { + t.Fatalf("GetSnapshot: %v", err) + } + if snap.Text != "summary text" { + t.Fatalf("snapshot text = %q, want 'summary text'", snap.Text) + } + + list := mgr.List() + if len(list) == 0 { + t.Fatal("List() returned empty") + } +}