Proof of concept job status display in nushell

job-prompt.nu

# Jobs Prompt
# Proof-of-concept to display the status of ongoing (and completed) nushell jobs
# by showing their state in the $env.PROMPT_COMMAND closure. Abuses some tricks with $env, and hooks.
# Namely, it uses the `$env.config.hooks.pre_prompt` hook to call the `prompt` command, which
# saves it's result to `$env.JOBS_PROMPT`. Then, in the custom `$env.PROMPT_COMMAND` closure, that
# variable is read from, and used to display the text before your directory / real prompt.
#
# This was just a proof of concept that I worked on because I found it interesting, there are surely bugs,
# and the code is not good. But, it can serve as a base to work off of for anyone who likes the idea!
#
# Left as exercices to the reader:
# - Maybe investigate displaying the jobs prompt using the `$env.PROMPT_COMMAND_RIGHT` closure, so the statuses
#   are not in the way of your cwd
# - Maybe add a command that toggles the status display from being shown at all, when you don't want to see the updates
#   (just check for `$env.JOBS_PROMPT_ENABLED` in the `$env.PROMPT_COMMAND` closure in the `setup-env` command?)
# - Generally using this module to do anything of use: It currently does no real work, and is just designed to show that
#   with some effort it's possible to actually make this into something cool.
#
# To test it out, save this file somewhere, and run:
# ```nushell
# use job-prompt.nu *
# setup-env # Sets up the environment variables
# spawn <path> # Try spawning a task
# cleanup # Try removing finished tasks
# ```

use std/iter
# Gets cached jobs from env, safely.
export def --env 'get-jobs' []: [
  any -> table<video_path: path, uuid: string, id: int, last_message: any, get_status: closure, get_index: closure>
] {
  $env.VIDEO_PROCESSING_JOBS = (
    $env.VIDEO_PROCESSING_JOBS?
    | default []
  )
  $env.VIDEO_PROCESSING_JOBS
}

# Sets cached jobs in env, safely.
export def --env 'set-jobs' [
  index?: int # The index to change, if we are editing a specific index.
]: [
  record<get_index: closure, get_status: closure, video_path: path, uuid: string, id: int, last_message: any> -> table<video_path: path, uuid: string, id: int, last_message: any, get_index: closure, get_status: closure>,
  table<get_index: closure, get_status: closure, video_path: path, uuid: string, id: int, last_message: any> -> table<video_path: path, uuid: string, id: int, last_message: any, get_index: closure, get_status: closure>,
  nothing -> table<video_path: path, uuid: string, id: int, last_message: any, get_index: closure, get_status: closure>
] {
  let input = $in

  $env.VIDEO_PROCESSING_JOBS = (
    match ($input | describe -d | get type) {
      'record' => {
        # Upsert the table at the specified index.
        if ($index | is-not-empty) and ($index != -1) {
          get-jobs
          | upsert $index { $input }
        } else {
          error make {
            msg: $'Expected $index to be >= 0.'
            label: {
              span: (metadata $index).span
              text: $'This was not >= 0.'
            }
          }
        }
      }
      'list'|'table' => {
        # Replace the table.
        $input
      }
      'nothing' => {
        # Clear the table
        []
      }
      $other => {
        error make {
          msg: $'Expected pipeline input type to be one of: record, table, or nothing. Got: ($other).'
          label: {
            span: (metadata $input).span
            text: $'Unexpected type here.'
          }
        }
      }
    }
  )

  get-jobs
}

# Spawn a video processing task, linked to the prompt.
export def --env 'spawn' [
  video_path: path
]: [
  any -> record<video_path: path, uuid: string, id: int, last_message: any, get_status: closure, get_index: closure>
] {
  let uuid = random uuid

  let id = job spawn --tag $'[($uuid) - Uploading video: ($video_path)]' {
    let statuses = [
      {complete: false status: 'Transcoding'}
      {complete: false status: 'Filtering'}
      {complete: false status: 'Uploading'}
      {complete: true status: 'Uploaded'}
    ]

    let job_id = job id
    for $index in 0..<($statuses | length) {
      let $status = $statuses | get $index
      let $message = {
        job: $job_id
        status: $status.status
        complete: $status.complete
        result: (if ($status.complete) { $'https://example.com/video.mp4' })
      }

      try {
        $message | job send --tag=($job_id) 0
      }

      if $status.complete {
        break
      } else {
        while true {
          sleep 1sec;
          match (try { job recv --tag=($job_id) --timeout=(1sec) }) {
            {received: true} => {
              break
            }
          }
        }
      }
    }

    # while true {
    #   let index = random int 0..<($statuses | length)
    #   let status = $statuses | get $index
    #   let message = {
    #     job: $job_id
    #     status: $status.status
    #     complete: $status.complete
    #     result: (if ($status.complete) { $'https://example.com/video.mp4' })
    #   }

    #   try {
    #     $message
    #     | job send --tag=($job_id) 0
    #   }

    #   if $status.complete {
    #     break
    #   }
    # }
  }

  let get_index = {||
    get-jobs
    | iter find-index {|job| $job.uuid == $uuid }
  }

  let get_env_object = {||
    let index = do $get_index
    let item = if ($index != -1) {
      get-jobs
      | get $index
    }
    {index: $index item: $item}
  }

  let match_message = {|message: any|
    match $message {
      {complete: true , result: $result} => { $result | ansi link --text $"[(ansi light_green)Uploaded ✅(ansi reset)]" }
      {complete: false , status: $status} => { $"[(ansi yellow)($status)(ansi reset)]" }
      _ => { $"[(ansi red)❌ Bad state(ansi reset)]" }
    }
  }

  # Could also do this via key-value storage (e.g. `use std-rfc/kv`).
  let update_env = {|message: any|
    let env_object = do $get_env_object

    if ($env_object.index == -1) {
      log error $'❌ Could not find index for job w/ ID: ($id) and Video Path: ($video_path)'
      return
    }

    # print $'Updating ($env_object.index) to have { last_message: ($message) }'
    $env_object.item
    | upsert last_message { $message }
    | set-jobs $env_object.index
  }

  # Gets the current status of a job
  let get_status = {|last_message: any|
    # Default `last_message` if it's null.
    let last_message = try {
      do $get_env_object
      | get -o item
      | get -o last_message
    }

    # If the job is no longer running, should we return the last message or an error?

    # Try to get a new message specifically for this job.
    let message = try {
      let out = job recv --tag $id --timeout 0sec
      if ($out | is-not-empty) {
        # Tell the job that we received the message.
        try {
          {received: true message: $out} | job send --tag $id $id
        }

        # print $'Got valid response for ID: ($id). Message: ($out)'
      }

      $out
    } catch {|error|
      # print $'Got invalid response for ID: ($id). Using $last_message: ($last_message)'
      $last_message
    }

    # print $'Got message: ($message)'

    {
      message: $message
      display: (do $match_message $message)
    }
  }

  let job_object = {
    uuid: $uuid
    video_path: $video_path
    id: $id
    last_message: null
    get_status: $get_status
    get_index: $get_index
  }

  get-jobs
  | append $job_object
  | set-jobs

  print $"Started new Job w/ ID: '($job_object.id)', Video Path: '($job_object.video_path)'."
  $job_object
}

# Get the prompt string to display.
export def --env 'prompt' [] {
  let results = get-jobs
  | each {|job|
    let message = try {
      do --ignore-errors --env $job.get_status $job.last_message
    }

    let new_job = ($job | upsert last_message { $message.message })
    # print $'Updating job from ($job) to: ($new_job).'
    {
      job: $new_job
      display: $message.display
    }
  }

  let any_complete = $results | any {|result|
    match $result.job.last_message {
      {complete: true} => { true }
      _ => { false }
    }
  }

  $env.VIDEO_PROCESSING_JOBS = $results.job

  let output = $results.display
  | str trim
  | compact -e
  | if ($in | is-not-empty) {
    if $any_complete {
      append $"\(cleanup: `(ansi yellow)cleanup(ansi reset)`\)"
    } else {
      $in
    }
    | str join $"(ansi reset) "
  } else {
    ""
  }

  $env.JOBS_PROMPT = $output

  $output
}

export def --env 'cleanup' [] {
  $env.VIDEO_PROCESSING_JOBS = (
    get-jobs
    | where {|job|
      match $job.last_message {
        {complete: true} => { false }
        _ => { true }
      }
    }
  )
  ignore
}

# Get the closure that generates the prompt displayed for $env.PROMPT_COMMAND.
export def --env 'get-prompt-closure' [] {
  let original = $env | get -o PROMPT_COMMAND
  let get_prompt = {||
    let original_text = do $original
    let prompt_text = prompt

    if ($prompt_text | is-not-empty) {
      $"($prompt_text) (ansi reset)($original_text)"
    } else {
      $original_text
    }
  }

  $get_prompt
}

export def --env 'get-prompt-hook' [] {
  let pre_prompt = "$env.JOBS_PROMPT = (prompt)"
  $env.config.hooks.pre_prompt = (
    $env.config.hooks.pre_prompt?
    | default []
    | append $pre_prompt
  )
  ignore
}

# Set up the jobs module
export def --env 'setup-env' [] {
  # Setup the VIDEO_PROCESSING_JOBS env var
  get-jobs

  # Setup the prompt hook
  get-prompt-hook

  # Setup the JOBS_PROMPT env var
  $env.JOBS_PROMPT = ($env.JOBS_PROMPT? | default '')

  let previous_command = $env.PROMPT_COMMAND?

  # Setup the prompt closure
  $env.PROMPT_COMMAND = {||
    let text = if ($previous_command | is-not-empty) {
      do $previous_command
    } else {
      (pwd)
    }

    let prompt_env = $env.JOBS_PROMPT?
    if ($prompt_env | is-not-empty) {
      $'($prompt_env) (ansi reset)($text)'
    } else {
      $text
    }
  }

  ignore
}

videos_nu.mp4