diff --git a/pico/cli.go b/pico/cli.go index 48e7909e..c1d1101f 100644 --- a/pico/cli.go +++ b/pico/cli.go @@ -2,6 +2,7 @@ package pico import ( "bufio" + "context" "encoding/json" "fmt" "log/slog" @@ -65,7 +66,7 @@ func (c *Cmd) notifications() error { return nil } -func (c *Cmd) logs() error { +func (c *Cmd) logs(ctx context.Context) error { sshClient, err := shared.CreateSSHClient( shared.GetEnv("PICO_SENDLOG_ENDPOINT", "send.pico.sh:22"), shared.GetEnv("PICO_SENDLOG_KEY", "ssh_data/term_info_ed25519"), @@ -76,16 +77,16 @@ func (c *Cmd) logs() error { if err != nil { return err } + defer sshClient.Close() session, err := sshClient.NewSession() - defer func() { - _ = session.Close() - }() if err != nil { return err } + defer session.Close() + stdoutPipe, err := session.StdoutPipe() if err != nil { return err @@ -96,6 +97,12 @@ func (c *Cmd) logs() error { return err } + go func() { + <-ctx.Done() + session.Close() + sshClient.Close() + }() + scanner := bufio.NewScanner(stdoutPipe) for scanner.Scan() { line := scanner.Text() @@ -183,7 +190,7 @@ func WishMiddleware(handler *CliHandler) wish.Middleware { opts.help() return } else if cmd == "logs" { - err = opts.logs() + err = opts.logs(sesh.Context()) if err != nil { wish.Fatalln(sesh, err) }