·13 min read

Dhi Gets a Brain: In-editor Chat with Streaming RAG

How to build a VS Code chat panel with server-sent event streaming, prioritized context-slot assembly, hybrid RAG retrieval, and a three-process webview architecture — no API keys, all open-source models.

Cover Image for Dhi Gets a Brain: In-editor Chat with Streaming RAG

In Post 1 Dhi learned to finish your sentences: fill-in-the-middle autocomplete with Tree-sitter chunking and StarCoder2-3B. In Post 2 it learned the whole repo: six languages, hybrid BM25 + vector search, one-shot workspace indexing.

Neither of those is a conversation. Autocomplete reacts to a cursor position; it can't answer "why does this function retry three times?" A comment from Alex Shev on the architecture post pushed on exactly this gap:

"An AI IDE is not just chat plus autocomplete; it needs repo understanding, edit boundaries, tool execution, review surfaces, and a way to explain why a change was made."

Post 3 builds the piece that ties the first two posts together and starts answering that list: an in-editor chat panel with streaming responses and repo-aware context. Every answer is grounded in the same hybrid search from Post 2 — the chat doesn't just talk about your code, it reads it first.

Everything is at github.com/sochaty/dhi, tag post-3.


What Changed Between Post 2 and Post 3

ComponentPost 2Post 3
Interaction modelAutocomplete only (implicit, cursor-driven)+ Explicit chat panel (VS Code webview)
Response deliverySingle JSON responseServer-Sent Events, one token at a time
Context sourceRetrieved chunks onlySystem prompt + active file + retrieved chunks + conversation history
New endpointPOST /chat (StreamingResponse, text/event-stream)
New extension surfaceextension/src/chat/panel.ts, dhi.openChat command
CancellationN/AAbortController — a new message aborts the in-flight one

Part 1: Three Processes, One Conversation

A VS Code chat panel is not one program talking to a server — it's three processes relaying messages to each other, and the architecture rule from ARCHITECTURE.md is explicit about who's allowed to do what: all HTTP goes through DhiClient, and the webview never calls fetch() directly.

%%{init: {'theme': 'dark'}}%%
flowchart LR
    subgraph webview["Webview  (sandboxed, CSP-locked)"]
        UI["chat UI\ntextarea + bubbles"]
    end
    subgraph host["Extension Host  (Node.js)"]
        PANEL["ChatPanel\npanel.ts"]
        CLIENT["DhiClient\nclient/index.ts"]
    end
    subgraph server["Dhi Server  (FastAPI)"]
        EP["POST /chat\nmain.py"]
        CHAT["stream_chat()\nchat.py"]
    end
    OLLAMA["Ollama\nllama3.2:3b"]

    UI -- "postMessage\n{type:'chat.send'}" --> PANEL
    PANEL -- "await client.chat()" --> CLIENT
    CLIENT -- "fetch POST /chat" --> EP
    EP --> CHAT
    CHAT -- "prompt" --> OLLAMA
    OLLAMA -- "token stream" --> CHAT
    CHAT -- "SSE: data: {token}" --> EP
    EP -- "SSE" --> CLIENT
    CLIENT -- "yield token" --> PANEL
    PANEL -- "postMessage\n{type:'chat.token'}" --> UI

Three isolation boundaries, three reasons:

  1. Webview ↔ Extension Host — the webview runs in a sandboxed iframe with a strict Content-Security-Policy (default-src 'none'). It cannot make network requests at all. It can only postMessage to the extension host and receive messages back.
  2. Extension Host ↔ Server — only DhiClient is allowed to call fetch(). Providers and panels call methods on the client; nobody constructs a request URL by hand. This is what makes it possible to point the whole extension at a different server URL from one config setting.
  3. Server ↔ Ollama — the FastAPI process is the only thing that speaks to Ollama's /api/generate. The extension never sees an Ollama URL.

This is more ceremony than a single fetch() call, but it means the webview's CSP can stay maximally strict, and every outbound request funnels through one auditable chokepoint.


Part 2: Context Slot Assembly

A chat model has a fixed context window. Dhi's prompt has five things competing for space, and they don't have equal claim to it. chat.py documents the priority order at the top of the file:

"""Streaming chat endpoint — context assembly + Ollama token stream.

Context slot order (highest → lowest trim priority):
  1. System prompt  (fixed, never trimmed)
  2. User message   (never trimmed)
  3. Active file     (trimmed to _MAX_FILE_CHARS)
  4. RAG chunks     (top _MAX_CHUNKS, each trimmed to _MAX_CHUNK_CHARS)
  5. History        (last _MAX_HISTORY_TURNS turns × 2 messages)
"""
%%{init: {'theme': 'dark'}}%%
flowchart TB
    SYS["1 · System prompt\nfixed — never trimmed"]
    MSG["2 · User message\nnever trimmed"]
    FILE["3 · Active file\ntrimmed to 2,000 chars"]
    RAG["4 · RAG chunks\ntop 3 · 400 chars each\nvia store.hybrid_query()"]
    HIST["5 · History\nlast 4 turns × 2 messages"]
    PROMPT["Flat prompt string\n→ Ollama /api/generate"]

    SYS --> PROMPT
    MSG --> PROMPT
    FILE --> PROMPT
    RAG --> PROMPT
    HIST --> PROMPT

    style SYS fill:#4a4a2e
    style MSG fill:#4a4a2e
    style FILE fill:#2e3a4a
    style RAG fill:#2e3a4a
    style HIST fill:#3a2e2e

The trim budget is deliberately asymmetric — the system prompt and the question you just typed are never touched, but everything that's context about the past (the open file, retrieved chunks, chat history) gets capped. assemble_prompt is a pure function, which is what makes it trivial to unit test without touching Ollama:

def assemble_prompt(request: ChatRequest, rag_chunks: list[str]) -> str:
    """Build a flat-string prompt from prioritised context slots."""
    parts: list[str] = [_SYSTEM_PROMPT]

    if request.file_content:
        content = request.file_content[:_MAX_FILE_CHARS]
        lang = request.language or "text"
        parts.append(f"# Active file: {request.file_path}\n```{lang}\n{content}\n```")

    if rag_chunks:
        trimmed = [c[:_MAX_CHUNK_CHARS] for c in rag_chunks[:_MAX_CHUNKS]]
        parts.append("# Relevant code from this repo\n" + "\n---\n".join(trimmed))

    for msg in request.history[-(_MAX_HISTORY_TURNS * 2):]:
        label = "User" if msg.role == "user" else "Assistant"
        parts.append(f"{label}: {msg.content}")

    parts.append(f"User: {request.message}\nAssistant:")
    return "\n\n".join(parts)

The RAG chunks are the interesting part — they come straight from Post 2's hybrid_query(), so a question like "why does the indexer skip .git?" pulls the actual _SKIP_DIRS chunk out of rag/indexer.py, not a hallucinated guess:

async def stream_chat(request: ChatRequest, store: Any) -> AsyncGenerator[str, None]:
    rag_chunks: list[str] = []
    if request.message.strip():
        rag_chunks = store.hybrid_query(request.message, n_results=_MAX_CHUNKS)

    prompt = assemble_prompt(request, rag_chunks)
    ...

This is the same ChunkStore instance the /complete and /search endpoints use — one index, three consumers.


Part 3: The SSE Contract

FastAPI's StreamingResponse turns an async generator into a text/event-stream response. The contract is two lines of protocol:

@app.post("/chat")
async def chat_endpoint(req: ChatRequestModel) -> StreamingResponse:
    """Stream a chat response as Server-Sent Events.

    Each event: ``data: {"token": "..."}\\n\\n``
    Final event: ``data: [DONE]\\n\\n``
    """
    chat_req = ChatReq(...)

    async def event_stream() -> AsyncGenerator[str, None]:
        try:
            async for token in stream_chat(chat_req, store):
                yield f"data: {json.dumps({'token': token})}\n\n"
        except Exception as exc:
            logging.exception("chat_endpoint error")
            yield f"data: {json.dumps({'error': str(exc)})}\n\n"
        finally:
            yield "data: [DONE]\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
    )

The try/except/finally shape matters: if Ollama errors mid-stream, the client gets a {"error": ...} event and still gets [DONE] — the finally block guarantees the stream terminates cleanly even on failure, so the client's read loop never hangs waiting for a close that isn't coming.

%%{init: {'theme': 'dark'}}%%
sequenceDiagram
    participant C as DhiClient (extension host)
    participant S as FastAPI /chat
    participant O as Ollama

    C->>S: POST /chat  {message, file_content, history}
    S->>O: POST /api/generate  {prompt, stream: true}
    loop each generated token
        O-->>S: {"response": "...", "done": false}
        S-->>C: data: {"token": "..."}\n\n
    end
    O-->>S: {"done": true}
    S-->>C: data: [DONE]\n\n
    Note over C,S: connection closes

Notice stream_chat's HTTP client sets read=None on the Ollama timeout — deliberately no read deadline, because token generation can legitimately take longer than any fixed timeout would allow. The only timeout that matters is the connect timeout; once the stream starts, it runs until Ollama says done: true.


Part 4: The Webview Message Protocol

The webview and the extension host agree on five message types. Nothing here is HTTP — it's postMessage in both directions:

%%{init: {'theme': 'dark'}}%%
sequenceDiagram
    participant UI as Webview UI
    participant Panel as ChatPanel (host)

    UI->>Panel: {type:'chat.send', data:{message}}
    activate Panel
    Note over Panel: abort() any in-flight request<br/>new AbortController
    loop streamed tokens
        Panel-->>UI: {type:'chat.token', data: token}
    end
    Panel-->>UI: {type:'chat.done'}
    deactivate Panel

    UI->>Panel: {type:'chat.clear'}
    Note over Panel: abort in-flight request<br/>reset history to []
    Panel-->>UI: {type:'chat.cleared'}

_handleMessage in panel.ts starts every new send by aborting whatever is still in flight:

private async _handleMessage(message: string): Promise<void> {
  this._abortController?.abort();
  this._abortController = new AbortController();

  const editor = this._lastEditor;
  const req = {
    message,
    file_path: editor?.document.uri.fsPath ?? '',
    language: editor?.document.languageId ?? '',
    file_content: editor?.document.getText() ?? '',
    history: this._history.slice(-MAX_HISTORY_MESSAGES),
  };

  let fullResponse = '';
  try {
    for await (const token of this._client.chat(req, this._abortController.signal)) {
      fullResponse += token;
      void this._panel.webview.postMessage({ type: 'chat.token', data: token });
    }
    this._history.push(
      { role: 'user', content: message },
      { role: 'assistant', content: fullResponse },
    );
    void this._panel.webview.postMessage({ type: 'chat.done' });
  } catch (err) {
    if (err instanceof Error && err.name !== 'AbortError') {
      void this._panel.webview.postMessage({ type: 'chat.error', data: err.message });
    }
  }
}

Two things worth calling out:

  • History only grows on success. fullResponse accumulates locally and is only pushed into _history after the loop completes without throwing. Cancel or error out mid-stream and nothing corrupt lands in the conversation history that gets sent on the next request.
  • AbortError is swallowed, not surfaced. Aborting is a normal outcome of sending a second message quickly — it should not paint an error bubble in the UI. Any other exception (network failure, server 500) does surface as chat.error.

Part 5: Parsing SSE by Hand on the Client

DhiClient.chat() is an async generator wrapped around the browser-ish fetch streaming API (Node's fetch in the extension host supports it). SSE framing means events can arrive split across chunk boundaries, so the client has to buffer:

async *chat(req: ChatRequest, signal?: AbortSignal): AsyncGenerator<string> {
  const resp = await fetch(`${this.baseUrl}/chat`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(req),
    signal,
  });
  if (!resp.ok || !resp.body) {
    const text = await resp.text();
    throw new Error(`Dhi server ${resp.status}: ${text}`);
  }
  const reader = resp.body.getReader();
  const decoder = new TextDecoder();
  let buffer = '';
  try {
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      buffer += decoder.decode(value, { stream: true });
      const lines = buffer.split('\n');
      buffer = lines.pop() ?? '';           // keep the trailing partial line
      for (const line of lines) {
        if (!line.startsWith('data: ')) continue;
        const payload = line.slice(6).trim();
        if (payload === '[DONE]') return;
        try {
          const parsed = JSON.parse(payload) as { token?: string; error?: string };
          if (parsed.error) throw new Error(parsed.error);
          if (parsed.token) yield parsed.token;
        } catch {
          // skip malformed SSE lines
        }
      }
    }
  } finally {
    reader.releaseLock();
  }
}

The lines.pop() trick is the whole trick: buffer.split('\n') always leaves the last element as whatever text came after the final \n in the chunk — which is either an empty string (chunk ended cleanly) or a partial line waiting for its terminator (chunk was cut mid-line). Popping it off and reassigning it to buffer means a line is only ever processed once it's known to be complete.


Part 6: A Gotcha — the Active Editor Disappears

The chat context is supposed to include "whatever file you're looking at." The obvious API is vscode.window.activeTextEditor. It's also wrong here, and the comment in panel.ts explains why:

// activeTextEditor becomes undefined when a webview gains focus;
// track the last real editor so file context is always available.
private _lastEditor: vscode.TextEditor | undefined;

The moment you click into the chat panel's textarea to type a question, VS Code's activeTextEditor goes undefined — a webview isn't a text editor, so there is no "active editor" while it has focus. If the chat sent activeTextEditor at request time, every question you type while looking at the chat panel would carry no file context at all, which is precisely the moment you're most likely to be asking about the file you just switched away from.

The fix is to track the last real editor separately, updated only on onDidChangeActiveTextEditor:

this._lastEditor = vscode.window.activeTextEditor;
vscode.window.onDidChangeActiveTextEditor(
  (editor) => { if (editor) this._lastEditor = editor; },
  undefined,
  context.subscriptions,
);

The callback only updates _lastEditor when editor is truthy — so focusing the webview (which fires the event with undefined) is a no-op, and _lastEditor keeps pointing at the last real file until you focus a different one.


Part 7: Streaming Markdown, Rendered Incrementally

The model's response arrives one token at a time, but it needs to render as formatted Markdown — code fences, inline code, bold — not a scrolling wall of raw text. The webview ships a deliberately small regex-based renderer rather than pulling in a Markdown library through the CSP:

function renderMd(text) {
  // fenced code blocks
  text = text.replace(/```(\w*)\n([\s\S]*?)```/g, (_, lang, code) =>
    '<pre><code>' + escape(code.trimEnd()) + '</code></pre>'
  );
  // inline code
  text = text.replace(/`([^`\n]+)`/g, (_, c) => '<code>' + escape(c) + '</code>');
  // bold
  text = text.replace(/\*\*(.+?)\*\*/g, '<strong>$1</strong>');
  // paragraphs
  return text.split(/\n\n+/).map(p => '<p>' + p.replace(/\n/g, '<br>') + '</p>').join('');
}

Every incoming token re-runs renderMd over the entire accumulated response and replaces the bubble's innerHTML:

function appendToken(token) {
  if (!currentEl) return;
  rawContent += token;
  currentEl.innerHTML = renderMd(rawContent) + '<span class="cursor"></span>';
  messagesEl.scrollTop = messagesEl.scrollHeight;
}

Re-rendering the whole string on every token sounds wasteful, but at the token rate a 3B model streams at, and for the length of a typical chat answer, this is invisible in practice — and it sidesteps a much harder problem: incrementally patching a partially-open code fence is fiddly (what do you render while you're two backticks into a closing ``` ?). Full re-render means the regex always sees the true current state, so a fence that just closed renders correctly on the very next token.


Part 8: Directory Structure

dhi/
├── docker-compose.yml
├── extension/
│   └── src/
│       ├── chat/
│       │   └── panel.ts          ← ChatPanel webview, SSE relay, AbortController
│       ├── client/
│       │   └── index.ts          ← DhiClient.chat() — the only fetch() to /chat
│       └── extension.ts          ← registers `dhi.openChat` command
└── server/
    ├── main.py                    ← POST /chat (StreamingResponse)
    ├── chat.py                    ← assemble_prompt(), stream_chat()
    ├── rag/
    │   └── store.py                ← hybrid_query() reused from post-2
    └── tests/
        └── test_chat.py            ← 12 tests: prompt assembly + streaming

Part 9: Tests

test_chat.py covers the context-slot logic in isolation — assemble_prompt is pure, so no Ollama mock is needed for most of it:

AreaTestsWhat they cover
TestAssemblePrompt9System prompt always present, file/RAG/history trimming, empty-slot omission
TestStreamChat3RAG query triggered on non-empty message, Ollama call shape, token yielding

One example that pins down a real trimming edge case — old history turns get dropped, but never the current message:

def test_trims_old_history_turns(self) -> None:
    history = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
    request = ChatRequest(message="current question", history=history)
    prompt = assemble_prompt(request, [])
    assert "msg0" not in prompt          # oldest turns trimmed
    assert "current question" in prompt  # current message never trimmed

The full suite runs with no live Ollama or Chroma — same rule as post-2: tests mock at layer boundaries.


Addressing Alex Shev's Comment

Post 3 answers two of the five points directly and sets up a third. The rest are the roadmap:

PointStatus
Repo understanding✅ — hybrid_query() from Post 2 feeds every chat answer, not just autocomplete
Edit boundaries✅ — the model only ever talks; nothing in Post 3 writes to a file. Chat is read-and-explain, not read-and-edit
Explain why◐ partial — the RAG chunks that informed an answer aren't yet surfaced in the UI, only used silently in the prompt
Tool execution❌ roadmap — no terminal or command execution wired into chat yet
Review surfaces❌ roadmap — no diff preview from a chat-suggested change yet

What's Next: Post 4

Post 3 gives Dhi a voice. The next post is about giving it hands: surfacing which retrieved chunks shaped an answer directly in the chat UI (closing out "explain why" fully), and a first pass at tool execution — feeding terminal output back into the conversation so Dhi can answer "why did that test fail" from the actual failure, not a guess.

The full series code is at github.com/sochaty/dhi. Check out post-3 to reproduce this post exactly. Star the repo if you want to build a production-grade open-source AI IDE without paying $20/month for the privilege.

Sources:

Comments